from __future__ import annotations

from dataclasses import dataclass, replace
from typing import Any

from app.config import config
from app.core.logging import get_audit_logger
from app.llm.contracts import LLMOutputKind, LLMProviderPort, LLMTask, ModelSelectionPolicy
from app.llm.models import (
    LLMErrorKind,
    LLMProviderError,
    LLMResolvedRoute,
    LLMRouteFallbackRule,
    LLMStructuredRequest,
    LLMStructuredResult,
    LLMTextRequest,
    LLMTextResult,
)


@dataclass(frozen=True, slots=True)
class TaskRouteProfile:
    output_kind: LLMOutputKind
    provider: str
    model: str
    fallback_provider: str | None = None
    fallback_model: str | None = None
    fallback_reason: str = "provider_unavailable_or_transient"
    metadata: dict[str, Any] | None = None

    def build_route(self, *, task: LLMTask) -> LLMResolvedRoute:
        fallback_rule = None
        if self.fallback_provider and self.fallback_model:
            fallback_rule = LLMRouteFallbackRule(
                provider=self.fallback_provider,
                model=self.fallback_model,
                reason=self.fallback_reason,
            )
        return LLMResolvedRoute(
            task=task,
            output_kind=self.output_kind,
            provider=self.provider,
            model=self.model,
            fallback_rule=fallback_rule,
            metadata=dict(self.metadata or {}),
        )


class DefaultModelSelectionPolicy(ModelSelectionPolicy):
    _ROUTES: dict[LLMTask, TaskRouteProfile] = {
        LLMTask.CLINICAL_DOCUMENT_EXTRACT: TaskRouteProfile(
            output_kind=LLMOutputKind.STRUCTURED_OBJECT,
            provider="gemini",
            model=config.GEMINI_MODEL_EXTRACT,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_CLINICAL,
            metadata={"required_capability": "structured_output", "volume": "high"},
        ),
        LLMTask.HISTORIA_CHUNK_SUMMARY: TaskRouteProfile(
            output_kind=LLMOutputKind.CONTROLLED_PLAIN_TEXT,
            provider="gemini",
            model=config.GEMINI_MODEL_EXTRACT,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_HISTORIA_LIGHT,
            metadata={"required_capability": "text_generation", "long_context": True, "volume": "high"},
        ),
        LLMTask.HISTORIA_FINAL_SUMMARY: TaskRouteProfile(
            output_kind=LLMOutputKind.CONTROLLED_PLAIN_TEXT,
            provider="gemini",
            model=config.GEMINI_MODEL_REASONING,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_HISTORIA_LIGHT,
            metadata={"required_capability": "text_generation", "long_context": True},
        ),
        LLMTask.SOAT_REASONING: TaskRouteProfile(
            output_kind=LLMOutputKind.STRUCTURED_OBJECT,
            provider="gemini",
            model=config.GEMINI_MODEL_REASONING,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_SOAT,
            metadata={"required_capability": "structured_output", "domain": "soat"},
        ),
        LLMTask.SOAT_CODE_GENERATION: TaskRouteProfile(
            output_kind=LLMOutputKind.STRUCTURED_COLLECTION,
            provider="gemini",
            model=config.GEMINI_MODEL_REASONING,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_SOAT,
            metadata={"required_capability": "structured_output", "domain": "soat"},
        ),
        LLMTask.SOAT_CHAT: TaskRouteProfile(
            output_kind=LLMOutputKind.CONTROLLED_PLAIN_TEXT,
            provider="gemini",
            model=config.GEMINI_MODEL_REASONING,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_SOAT,
            metadata={"required_capability": "text_generation", "long_context": True, "domain": "soat"},
        ),
        LLMTask.SOAT_MANUAL_CODE_LOOKUP: TaskRouteProfile(
            output_kind=LLMOutputKind.CONTROLLED_PLAIN_TEXT,
            provider="gemini",
            model=config.GEMINI_MODEL_REASONING,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_SOAT,
            metadata={"required_capability": "text_generation", "domain": "soat"},
        ),
        LLMTask.GLOSA_NOTE_GENERATION: TaskRouteProfile(
            output_kind=LLMOutputKind.CONTROLLED_PLAIN_TEXT,
            provider="gemini",
            model=config.GEMINI_MODEL_REASONING,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_SOAT,
            metadata={"required_capability": "text_generation", "domain": "soat"},
        ),
        LLMTask.CIE10_RESOLUTION: TaskRouteProfile(
            output_kind=LLMOutputKind.CONTROLLED_PLAIN_TEXT,
            provider="gemini",
            model=config.GEMINI_MODEL_EXTRACT,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_CIE10,
            metadata={"required_capability": "text_generation", "volume": "high"},
        ),
        LLMTask.CUPS_RESOLUTION: TaskRouteProfile(
            output_kind=LLMOutputKind.CONTROLLED_PLAIN_TEXT,
            provider="gemini",
            model=config.GEMINI_MODEL_EXTRACT,
            fallback_provider="groq",
            fallback_model=config.GROQ_MODEL_CUPS,
            metadata={"required_capability": "text_generation", "volume": "high"},
        ),
    }

    def resolve(self, task: LLMTask) -> LLMResolvedRoute:
        return self._ROUTES[task].build_route(task=task)


class DefaultModelRouter:
    def __init__(
        self,
        *,
        providers: dict[str, LLMProviderPort],
        policy: ModelSelectionPolicy,
        service: str = "llm_router",
    ) -> None:
        self._providers = dict(providers)
        self._policy = policy
        self._service = service
        self._audit_logger = get_audit_logger()

    def generate_text(self, request: LLMTextRequest) -> LLMTextResult:
        route = self._resolve_route(request.task)
        primary_request = replace(request, model=route.model)
        self._log_route("route.select", route.provider, route.model, route=route, fallback_used=False)
        try:
            provider = self._provider_or_raise(route.provider, route.model)
            return provider.generate_text(primary_request)
        except LLMProviderError as exc:
            return self._handle_text_fallback(route=route, request=primary_request, exc=exc)

    def generate_structured(self, request: LLMStructuredRequest) -> LLMStructuredResult:
        route = self._resolve_route(request.task)
        primary_request = replace(request, model=route.model)
        self._log_route("route.select", route.provider, route.model, route=route, fallback_used=False)
        try:
            provider = self._provider_or_raise(route.provider, route.model)
            return provider.generate_structured(primary_request)
        except LLMProviderError as exc:
            return self._handle_structured_fallback(route=route, request=primary_request, exc=exc)

    def _resolve_route(self, task: LLMTask | None) -> LLMResolvedRoute:
        if task is None:
            raise LLMProviderError(
                "La tarea LLM es requerida para resolver la ruta",
                kind=LLMErrorKind.INVALID_CONFIGURATION,
                provider="router",
            )
        return self._policy.resolve(task)

    def _provider_or_raise(self, provider_name: str, model: str) -> LLMProviderPort:
        provider = self._providers.get(provider_name)
        if provider is None:
            raise LLMProviderError(
                f"Proveedor {provider_name} no registrado",
                kind=LLMErrorKind.INVALID_CONFIGURATION,
                provider=provider_name,
                model=model,
            )
        availability = provider.availability()
        if not availability.configured:
            raise LLMProviderError(
                availability.reason or f"Proveedor {provider_name} no configurado",
                kind=LLMErrorKind.INVALID_CONFIGURATION,
                provider=provider_name,
                model=model,
            )
        return provider

    def _handle_text_fallback(
        self,
        *,
        route: LLMResolvedRoute,
        request: LLMTextRequest,
        exc: LLMProviderError,
    ) -> LLMTextResult:
        if not self._should_fallback(exc, route):
            raise
        fallback = route.fallback_rule
        assert fallback is not None
        self._log_route(
            "route.fallback",
            fallback.provider,
            fallback.model,
            route=route,
            fallback_used=True,
            reason=exc.kind.value,
        )
        provider = self._provider_or_raise(fallback.provider, fallback.model)
        result = provider.generate_text(replace(request, model=fallback.model))
        return replace(result, fallback=fallback)

    def _handle_structured_fallback(
        self,
        *,
        route: LLMResolvedRoute,
        request: LLMStructuredRequest,
        exc: LLMProviderError,
    ) -> LLMStructuredResult:
        if not self._should_fallback(exc, route):
            raise
        fallback = route.fallback_rule
        assert fallback is not None
        self._log_route(
            "route.fallback",
            fallback.provider,
            fallback.model,
            route=route,
            fallback_used=True,
            reason=exc.kind.value,
        )
        provider = self._provider_or_raise(fallback.provider, fallback.model)
        result = provider.generate_structured(replace(request, model=fallback.model))
        return replace(result, fallback=fallback)

    def _should_fallback(self, exc: LLMProviderError, route: LLMResolvedRoute) -> bool:
        if route.fallback_rule is None:
            return False
        return exc.kind in {
            LLMErrorKind.TRANSIENT,
            LLMErrorKind.RATE_LIMITED,
            LLMErrorKind.PROVIDER_UNAVAILABLE,
            LLMErrorKind.UNSUPPORTED_CAPABILITY,
            LLMErrorKind.INVALID_CONFIGURATION,
        }

    def _log_route(
        self,
        action: str,
        provider_name: str,
        model: str,
        *,
        route: LLMResolvedRoute,
        fallback_used: bool,
        reason: str | None = None,
    ) -> None:
        resource: dict[str, Any] = {
            "task": route.task.value,
            "output_kind": route.output_kind.value,
            "fallback_used": fallback_used,
            "selected_provider": provider_name,
            "selected_model": model,
        }
        if reason:
            resource["fallback_reason"] = reason
        self._audit_logger.llm_event(
            action=action,
            outcome="success",
            service=self._service,
            provider=provider_name,
            model=model,
            resource=resource,
        )
