"""Kafka executor helper functions."""
import json
from typing import Dict
from typing import List
from typing import Optional
from typing import Union

from kafka.consumer.fetcher import ConsumerRecord
from kafka.producer.future import RecordMetadata
from kafka.structs import OffsetAndMetadata
from kafka.structs import TopicPartition

from dbdeploy.base.config import CONTINUE_ON_DLQ
from dbdeploy.util.exceptions import KafkaSendError
from dbdeploy.util.kafka.structs import PartitionOffset


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


def on_send_error_callback(err: str) -> None:
    """Execute callback on message sending error."""
    raise KafkaSendError(f'KafkaProducer error: {err}')


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


def compare_offsets(
        latest_offsets: Dict[int, int],
        current_offsets: List[PartitionOffset]) -> None:
    """Compare consumer offsets and update finished/unfinished sets.

    Args:
        latest_offsets: partition <-> offset mapping
        current_offsets: current consumer offsets
    """
    for partition_offset in current_offsets:
        # This will prevent from pop()'ing finished partitions.
        partition = int(partition_offset.partition)
        if partition not in latest_offsets:
            continue
        latest_offset = latest_offsets.get(partition)
        if latest_offset is None:
            continue
        if partition_offset.offset >= latest_offset:
            latest_offsets.pop(partition, None)


def check_dlq_topic_records(
            polled_records: Dict[TopicPartition, List[ConsumerRecord]],
            partitions: List[TopicPartition],
        ) -> Optional[Dict[str, Union[bool, ConsumerRecord]]]:
    """Look for exceptions in DLQ."""
    for partition in partitions:
        records = polled_records.get(partition)
        if records is None:
            continue
        for record in records:
            return {
                'stop': not CONTINUE_ON_DLQ,
                'record': record
            }
    return None


def serialize_to_json(obj: None | Dict[str, str]) -> bytes | None:
    return None if obj is None else \
        json.dumps(obj, default=str).encode('utf-8')


def serialize_to_string(value: None | str) -> bytes | None:
    return None if value is None else value.encode('utf-8')