"""Utility helpers for parsing AWS MSK event batches."""

import json
from typing import Any
from typing import Generator
from typing import Mapping

from kafka_utils.consumer.deserializer.simple_json import JSONDeserializer
from kafka_utils.consumer.deserializer.string import StringDeserializer
from kafka_utils.consumer.source.mapping import EventSourceMessage
from pydantic import ValidationError
from src.models.preference_change import InvalidPreferenceMessage
from src.models.preference_change import PreferenceChangeEnvelope
from src.models.preference_change import UpdatedSubscription


def parse(
        event: Mapping[str, Any] | None,
) -> tuple[list[PreferenceChangeEnvelope], list[InvalidPreferenceMessage]]:
    """Parse a raw AWS MSK → Lambda event payload into batch items.

    Extracts and validates preference change messages from MSK event records.

    Args:
        event: Raw AWS MSK event payload containing Kafka records

    Returns:
        A tuple of (valid_messages, invalid_messages) where valid_messages is a list of
        PreferenceChangeEnvelope objects and invalid_messages captures validation failures.
    """
    valid_messages: list[PreferenceChangeEnvelope] = []
    invalid_messages: list[InvalidPreferenceMessage] = []

    if not event:
        return valid_messages, invalid_messages

    for message_key, message_value_dict in _get_records(event):
        try:
            envelope = PreferenceChangeEnvelope.model_validate(message_value_dict)
            valid_messages.append(envelope)
        except ValidationError as exc:
            invalid_messages.append(
                InvalidPreferenceMessage(
                    message_key=message_key,
                    message_value=json.dumps(message_value_dict),
                    error_type=exc.__class__.__name__,
                    error_message=str(exc),
                )
            )

    return valid_messages, invalid_messages


def _get_records(
        event: Mapping[str, Any] | None
) -> Generator[tuple[str, dict], None, None]:
    """Flatten records from the source event.

    Deserializes Kafka message values from base64-encoded strings to JSON dictionaries.

    Args:
        event: AWS MSK event containing records

    Yields:
        Tuples of (message_key, message_value_dict).
    """
    if not event or 'records' not in event:
        return

    string_deserializer = StringDeserializer()
    json_deserializer = JSONDeserializer()
    for _, msk_message in EventSourceMessage(event):
        raw_key = msk_message.key
        message_key = string_deserializer.deserialize(raw_key) if raw_key is not None else ''
        message_value = string_deserializer.deserialize(msk_message.value)
        message_value = json_deserializer.deserialize(message_value)
        yield message_key, message_value


def flatten_events(envelopes: list[PreferenceChangeEnvelope]) -> dict[str, UpdatedSubscription]:
    """Flatten PreferenceChangeEnvelope list into SubscriptionUpdateBatchItem list.

    Deduplicate by subscription_id, keeping the latest update based on updated_at timestamp.

    Args:
        envelopes: List of preference change envelopes to flatten

    Returns:
        Dictionary mapping subscription IDs to their latest UpdatedSubscription
    """
    updated_subscriptions: dict[str, UpdatedSubscription] = {}

    for envelope in envelopes:
        for subscription in envelope.updated_subscriptions:
            subscription.updated_at = envelope.updated_at
            # Deduplicate: keep latest based on updated_at
            if subscription.id not in updated_subscriptions or \
                    subscription.updated_at > updated_subscriptions[subscription.id].updated_at:
                updated_subscriptions[subscription.id] = subscription

    return updated_subscriptions