# fastkafka\registry.py
import logging
from inspect import signature, iscoroutinefunction
from typing import Any, Callable
from pydantic import BaseModel
from fastkafka.message import KafkaMessage
from fastkafka.di.di_container import resolve

__all__ = ["kafka_handler"]

logger = logging.getLogger(__name__)
handlers_registry: dict[str, list["CompiledHandler"]] = {}


class CompiledHandler:
    __slots__ = ("topic", "func", "sig", "data_model", "headers_model", "dependencies")

    def __init__(
        self,
        topic: str,
        func: Callable[..., Any],
        data_model: type[BaseModel] | None,
        headers_model: type[BaseModel] | None,
    ):
        self.topic = topic
        self.func = func
        self.sig = signature(func)
        self.data_model = data_model
        self.headers_model = headers_model

        # подготавливаем зависимости (DI) для параметров хендлера, которые не KafkaMessage и не data_model
        self.dependencies: dict[str, Any] = {}
        for name, param in self.sig.parameters.items():
            if param.annotation not in (KafkaMessage, data_model):
                self.dependencies[name] = resolve(param.annotation)

    async def handle(
        self, raw_data: Any, raw_headers: dict[str, str] | None, key: str | None
    ):
        # Валидация data
        msg_data = self.data_model(**raw_data) if self.data_model else raw_data
        # Валидация headers
        headers_src = raw_headers or {}
        msg_headers = (
            self.headers_model(**headers_src)
            if self.headers_model
            else headers_src
        )

        message = KafkaMessage(
            topic=self.topic, data=msg_data, headers=msg_headers, key=key
        )

        # Собираем kwargs: в хендлер прокидываем KafkaMessage во все параметры с type KafkaMessage или равные data_model
        kwargs = {
            name: (
                self.dependencies.get(name)
                if param.annotation not in (KafkaMessage, self.data_model)
                else message
            )
            for name, param in self.sig.parameters.items()
        }

        return await self.func(**kwargs)


def kafka_handler(
    topic: str,
    data_model: type[BaseModel] | None = None,
    headers_model: type[BaseModel] | None = None,
):
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        if not iscoroutinefunction(func):
            raise TypeError("Handler must be async")
        handlers_registry.setdefault(topic, []).append(
            CompiledHandler(topic, func, data_model, headers_model)
        )
        logger.debug(
            "Registered handler %s for topic %s (data_model=%s, headers_model=%s)",
            func.__name__,
            topic,
            getattr(data_model, "__name__", None),
            getattr(headers_model, "__name__", None),
        )
        return func

    return decorator
