from abc import ABC, abstractmethod
from typing import Any, overload

from instructor import AsyncInstructor
from pydantic import BaseModel

from yalc.clients.schemas import ClientCall, ClientMessage
from yalc.clients.strategy import ClientMetadataStrategy
from yalc.common.pricing import PricingService
from yalc.common.schemas import (
    LLMModel,
    LLMProvider,
    LLMRole,
    ResponseStats,
)
from yalc.common.utils import to_context_messages


class Client(ABC):
    """Abstract base class for provider-specific LLM clients.

    Use :func:`~yalc.clients.client_factory.create_client` to obtain a concrete instance.
    Subclasses register themselves via the ``@Client.provider(LLMProvider.X)`` decorator.
    """

    _registry: dict[LLMProvider, type["Client"]] = {}

    @classmethod
    def provider(cls, llm_provider: LLMProvider):
        """Class decorator that registers a Client subclass for a given provider."""

        def decorator(subclass: type["Client"]) -> type["Client"]:
            cls._registry[llm_provider] = subclass
            return subclass

        return decorator

    def __init__(
        self,
        model: LLMModel,
        instructor_client: AsyncInstructor,
        metadata_strategies: list[ClientMetadataStrategy] = [],
    ):
        self.metadata_strategies = metadata_strategies
        self.pricing_service = PricingService()
        self.instructor_client = instructor_client
        self.model = model

    @overload
    async def structured_response[T: BaseModel](
        self,
        response_type: type[T],
        messages: list[dict[str, str]],
        context: BaseModel,
    ) -> T: ...

    @overload
    async def structured_response[T: BaseModel](
        self,
        response_type: type[T],
        messages: list[dict[str, str]],
    ) -> tuple[T, ClientCall]: ...

    async def structured_response[T: BaseModel](
        self,
        response_type: type[T],
        messages: list[dict[str, str]],
        context: BaseModel | None = None,
    ) -> T | tuple[T, ClientCall]:
        """
        Sends messages to the LLM and returns a parsed Pydantic response.

        If ``context`` is provided, all registered metadata strategies are invoked and only
        the parsed response is returned. Otherwise, a ``(parsed, ClientCall)`` tuple is returned
        so the caller can inspect token usage and costs.
        """
        response, llm_call = await self._structured_response(
            response_type, messages
        )

        if context is not None:
            for strategy in self.metadata_strategies:
                strategy.handle(llm_call, context)
            return response

        return response, llm_call

    async def _structured_response[T: BaseModel](
        self,
        response_type: type[T],
        messages: list[dict[str, str]],
    ) -> tuple[T, ClientCall]:
        parsed, response = await self._response(
            response_type, messages
        )

        llm_call = self._create_llm_call(messages, parsed, response)
        return parsed, llm_call

    def _create_llm_call(
        self,
        messages: list[dict[str, str]],
        parsed: BaseModel,
        response: Any,
    ) -> ClientCall:
        context_messages = to_context_messages(messages)
        client_message = self._parse_client_message(parsed)

        return ClientCall(
            context_messages=context_messages,
            client_message=client_message,
            model_name=self.model,
            **self._get_response_stats(response).model_dump(),
        )

    @staticmethod
    def _parse_client_message(parsed: BaseModel) -> ClientMessage:
        return ClientMessage(response=parsed, role=LLMRole.ASSISTANT)

    @abstractmethod
    async def _response[T: BaseModel](
        self,
        response_type: type[T],
        messages: list[dict[str, str]],
    ) -> tuple[T, Any]:
        pass

    @abstractmethod
    def _get_response_stats(self, response: Any) -> ResponseStats:
        pass
