"""
Defines the ``Station`` — the central handler registry and message dispatcher.

Handlers are registered with ``@Station.add_handler(source)`` where *source*
is either a ``BusType`` (category-level) or a line name string (instance-level).
When a message arrives, ``Station.dispatch`` resolves matching handlers by
checking the line's name first, then its ``bus_type``.
"""

import asyncio
import json
from collections.abc import Callable
from dataclasses import is_dataclass
from typing import Any, ClassVar, TypeVar

from pydantic import BaseModel

from ..hermes import Hermes
from ..utilities.serializers import JsonSerializer, Serializer
from .basekurier import BaseKurier
from .bus_types import BusType
from .exceptions import KurierDeserializationError
from .flow_root_meta_data import FlowRootMetaData

F = TypeVar("F", bound=Callable[..., Any])


class Station:
    """
    Central handler registry and message dispatcher.

    Handlers are keyed by ``(source, object_type_name)`` where *source* is
    a ``BusType`` enum member or a line name string, and *object_type_name*
    is the ``__name__`` of the handler's single typed parameter.
    """

    # Two-level dict: source → message_type_name → list of handler metadata
    station_registry: ClassVar[dict[str | BusType, dict[str, list[FlowRootMetaData]]]] = {}
    event_registry: ClassVar[dict[str | BusType, dict[str, list[FlowRootMetaData]]]] = {}

    @classmethod
    def add_handler(cls, source: str | BusType | None = None) -> Callable[[F], F]:
        """
        Decorator that registers a function as a message handler.

        The decorated function must accept exactly **one** typed parameter
        (the message type used for routing) plus an optional ``*context``
        variadic for the raw envelope dict.

        Args:
            source: A line name (str), a ``BusType``, or ``None``.
                    Defaults to ``BusType.IN_MEMORY`` when ``None``.

        Raises:
            ValueError: If the function does not have exactly one typed argument.
        """

        def wrapper(func: F) -> F:
            target_source = source if source is not None else BusType.IN_MEMORY

            is_async_func: bool = asyncio.iscoroutinefunction(func)
            annotations = cls._get_annotations(func)

            if len(annotations) != 1:
                error_msg = f"Passenger handler '{func.__name__}' can only have one argument (excluding return type). Found: {annotations}"
                Hermes.fatal(error_msg)
                raise ValueError(error_msg)

            message_type: type = annotations[0]
            message_type_str: str = annotations[0].__name__
            qualified_message_type_str = f"{message_type.__module__}.{message_type.__qualname__}"

            meta_data = FlowRootMetaData(
                source=target_source,
                is_async=is_async_func,
                parameter_type=message_type,
                function=func,
            )

            cls.station_registry.setdefault(target_source, {})
            # Register by BOTH to allow backward compatibility and explicit qualified routing
            cls.station_registry[target_source].setdefault(message_type_str, []).append(meta_data)
            cls.station_registry[target_source].setdefault(qualified_message_type_str, []).append(meta_data)

            source_log = target_source.name if isinstance(target_source, BusType) else f"Line:'{target_source}'"
            Hermes.info(f"Registered handler '{func.__name__}' for passenger '{message_type_str}' on {source_log}.")

            return func

        return wrapper

    @classmethod
    def add_event_handler(cls, event_type: str, source: str | BusType | None = None) -> Callable[[F], F]:
        """Decorator that registers a function as a handler for an explicit event type string."""

        def wrapper(func: F) -> F:
            target_source = source if source is not None else BusType.IN_MEMORY
            is_async_func: bool = asyncio.iscoroutinefunction(func)
            annotations = cls._get_annotations(func)

            if len(annotations) != 1:
                error_msg = f"Passenger handler '{func.__name__}' can only have one argument (excluding return type). Found: {annotations}"
                Hermes.fatal(error_msg)
                raise ValueError(error_msg)

            meta_data = FlowRootMetaData(
                source=target_source,
                is_async=is_async_func,
                parameter_type=annotations[0],
                function=func,
                event_type=event_type,
            )

            cls.event_registry.setdefault(target_source, {})
            cls.event_registry[target_source].setdefault(event_type, []).append(meta_data)

            source_log = target_source.name if isinstance(target_source, BusType) else f"Line:'{target_source}'"
            Hermes.info(f"Registered event handler '{func.__name__}' for event '{event_type}' on {source_log}.")
            return func

        return wrapper

    @classmethod
    async def dispatch(cls, line: str | BaseKurier, raw_envelope: Any) -> bool:
        """
        Deserialize an incoming envelope and invoke all matching handlers.

        Handler resolution order:
        1. Handlers registered to the line's **name** (instance-level).
        2. Handlers registered to the line's **bus_type** (category-level).

        Args:
            line: The originating line instance, or its name as a string
                  (auto-resolved via ``BaseKurier.from_registry``).
            raw_envelope: The raw data received from the broker (str, bytes, or dict).

        Returns:
            ``True`` if at least one handler was invoked, ``False`` otherwise.
        """
        line = BaseKurier.from_registry(line)

        serializer: Serializer = getattr(line, "serializer", JsonSerializer())
        try:
            envelope_data = serializer.deserialize_envelope(raw_envelope)
        except (json.JSONDecodeError, TypeError, ValueError) as e:
            Hermes.error(f"Failed to deserialize envelope on line '{line.name}': {e}")
            return False

        object_type_name = envelope_data.get("object_type_name")
        qualified_object_type_name = envelope_data.get("qualified_object_type_name")
        event_type = envelope_data.get("event_type")
        raw_payload = envelope_data.get("message_payload")
        correlation_id = envelope_data.get("correlation_id", "N/A")

        try:
            message_payload = serializer.deserialize(raw_payload) if isinstance(raw_payload, str | bytes) else raw_payload
        except Exception as e:
            Hermes.error(f"Failed to deserialize message payload on line '{line.name}': {e}")
            return False

        bus_type = getattr(line, "bus_type", None)

        Hermes.trace(f"Dispatching {object_type_name} via {line.name} (CID: {correlation_id})")

        if not object_type_name or message_payload is None:
            Hermes.warn(f"Received malformed envelope from {line.name}")
            return False

        applicable_handlers: list[FlowRootMetaData] = []

        def get_handlers_for_type(registry_for_source, type_name, qualified_name):
            handlers = []
            if qualified_name:
                handlers.extend(registry_for_source.get(qualified_name, []))
            if not handlers:
                handlers.extend(registry_for_source.get(type_name, []))
            return handlers

        if event_type:
            if line.name in cls.event_registry:
                applicable_handlers.extend(cls.event_registry[line.name].get(event_type, []))
            if bus_type and bus_type in cls.event_registry:
                applicable_handlers.extend(cls.event_registry[bus_type].get(event_type, []))

        # 1. Instance-level handlers (by line name)
        if line.name in cls.station_registry:
            specific_processors = get_handlers_for_type(cls.station_registry[line.name], object_type_name, qualified_object_type_name)
            applicable_handlers.extend(specific_processors)

        # 2. Category-level handlers (by BusType)
        if bus_type and bus_type in cls.station_registry:
            generic_processors = get_handlers_for_type(cls.station_registry[bus_type], object_type_name, qualified_object_type_name)
            applicable_handlers.extend(generic_processors)

        if not applicable_handlers:
            Hermes.info(f"No handlers found for '{object_type_name}' on {line.name}")
            return False

        num_dispatched = 0
        converted_values: dict[Any, Any] = {}

        for meta_data in applicable_handlers:
            func_name = meta_data.function.__name__
            parameter_type = meta_data.parameter_type
            if parameter_type not in converted_values:
                converted_values[parameter_type] = cls._convert_message_body_to_object(meta_data, message_payload)
            converted_value = converted_values[parameter_type]

            Hermes.debug(f"Invoking '{func_name}' (CID: {correlation_id})")

            if meta_data.is_async:
                await meta_data.function(converted_value, envelope_data)
            else:
                loop = asyncio.get_running_loop()
                await loop.run_in_executor(None, meta_data.function, converted_value, envelope_data)

            num_dispatched += 1

        return num_dispatched > 0

    @staticmethod
    def _convert_message_body_to_object(meta_data: FlowRootMetaData, message_payload: Any) -> Any:
        """
        Convert a deserialized message payload into the handler's target Python type.

        The payload has already been deserialized by the line's serializer,
        so it is a dict or primitive at this point. Supports Pydantic
        ``BaseModel`` (via ``model_validate``), dataclasses (via ``**dict``
        unpacking), and primitive types.

        Raises:
            KurierDeserializationError: If the payload cannot be converted to the target type.
        """
        target_type: type = meta_data.parameter_type

        if isinstance(target_type, type) and issubclass(target_type, BaseModel):
            try:
                if isinstance(message_payload, dict):
                    return target_type.model_validate(message_payload)
                return target_type.model_validate_json(str(message_payload))
            except Exception as e:
                raise KurierDeserializationError(f"Pydantic validation failed for {target_type.__name__}: {e}") from e

        elif is_dataclass(target_type):
            try:
                data = message_payload
                if isinstance(data, str):
                    data = json.loads(data)
                return target_type(**data)
            except Exception as e:
                raise KurierDeserializationError(f"Dataclass instantiation failed for {target_type.__name__}: {e}") from e

        else:
            try:
                if isinstance(message_payload, target_type):
                    return message_payload
                return target_type(message_payload)
            except Exception as e:
                raise KurierDeserializationError(f"Primitive conversion failed for {target_type.__name__}: {e}") from e

    @staticmethod
    def _get_annotations(func: F) -> list[type]:
        """Extract the type annotations of a function's parameters, excluding ``return``."""
        return [v for k, v in getattr(func, "__annotations__", {}).items() if k != "return"]
