"""Kafka executor helper functions."""
from typing import Dict
from typing import List
from typing import Set

from garcon.activity import Activity

from feed_ingestion.util.kafka.structs import PartitonOffset
from feed_ingestion.util.sentry_util import send_error_or_warning


def on_send_success_callback(latest_offsets: Dict, record) -> None:
    """Execute callback on successfully sent messages."""
    latest_offsets[record.partition] = record.offset


def on_send_error_callback(activity: Activity, err) -> Dict:
    """Execute callback on message sendign error."""
    send_error_or_warning(err)
    activity.logger.error('KafkaProducer error: %', err)
    return {'stop': True}


def get_offsets_per_partiton(topic: str, offsets: Dict) -> List:
    """Return a list of consumer's topic-offest tuples."""
    return [PartitonOffset(str(tp.partition), offsets[tp].offset)
            for tp in offsets.keys() if tp.topic == topic]


def compare_current_vs_latest_offsets(
        latest_offsets: Dict,
        current_offsets: List,
        finished_partitions: Set,
        unfinished_partitions: Set) -> None:
    """Compare consumer offsets and update finished/unfinished sets."""
    for offset in current_offsets:
        # This will prevent from pop()'ing finished partitions.
        if offset.partition in finished_partitions:
            continue
        partition_latest_offset = latest_offsets.get(int(offset.partition))
        if not partition_latest_offset and partition_latest_offset != 0:
            continue
        if offset.offset >= partition_latest_offset:
            unfinished_partitions.pop()
            finished_partitions.add(offset.partition)


def check_dlq_topic_records(
        sf_query_id: str,
        polled_records: Dict,
        unfinished_partitions: List,
        end_offsets: Dict) -> Dict:
    """Look for exeptions in DLQ."""
    for topic_partition in unfinished_partitions:
        records = polled_records.get(topic_partition)
        end_offset = end_offsets.get(topic_partition)
        if not records or not end_offset:
            unfinished_partitions.pop()
            continue
        for record in records:
            if record.offset > end_offset:
                unfinished_partitions.pop()
                break
            if sf_query_id.encode() in record.key:
                return {'stop': True, 'record': record}
    return {}