"""Cost-aware model router.

Sends each task to the cheapest *capable* model, with graceful fallback down the
tiers (frontier → local → mock) so a missing key or a dead Ollama never breaks a
run. Consults the semantic cache before spending tokens, and records every call to
the metrics store. This is the component that operationalizes the cost story.
"""
from __future__ import annotations

from typing import Optional

from .caching import SemanticCache
from .config import Settings
from .metrics import MetricsStore
from .providers import LLMProvider, LLMRequest, LLMResponse, ProviderRegistry
from .providers.base import Usage

# Tasks that are safe + worthwhile to semantic-cache (deterministic-ish, repeatable).
_CACHEABLE_TASKS = {"classify", "normalize", "extract"}


class ModelRouter:
    def __init__(
        self,
        registry: ProviderRegistry,
        settings: Settings,
        metrics: MetricsStore,
        semantic_cache: Optional[SemanticCache] = None,
    ) -> None:
        self.registry = registry
        self.settings = settings
        self.metrics = metrics
        self.cache = semantic_cache or SemanticCache(
            threshold=settings.semantic_cache_threshold,
            enabled=settings.enable_semantic_cache,
        )

    # --- public API -----------------------------------------------------------
    def run(self, req: LLMRequest, run_id: str = "adhoc") -> LLMResponse:
        # 1) Semantic cache short-circuit.
        if req.task in _CACHEABLE_TASKS and self.settings.enable_semantic_cache:
            cache_key = f"{req.task}:{req.cacheable_prefix()}:{req.user_content}"
            cached = self.cache.get(cache_key)
            if cached is not None:
                resp = LLMResponse.build(
                    text=cached, model="semantic-cache", provider="cache",
                    usage=Usage(), latency_ms=0.2, cache_hit=True,
                    routing_reason="semantic cache hit (LLM call skipped)",
                )
                self.metrics.record_call(run_id, resp, req.task)
                return resp

        # 2) Route to candidates and execute with fallback.
        candidates = self._candidates(req)
        resp: Optional[LLMResponse] = None
        for provider, model, reason in candidates:
            resp = provider.complete(req, model)
            resp.routing_reason = reason + (
                f" → {resp.routing_reason}" if resp.routing_reason else ""
            )
            if not resp.error and resp.text:
                break  # success
        if resp is None:  # should never happen (mock is always last)
            resp = self.registry.mock.complete(req, "mock")

        # 3) Record + populate cache.
        self.metrics.record_call(run_id, resp, req.task)
        if (
            req.task in _CACHEABLE_TASKS
            and self.settings.enable_semantic_cache
            and resp.text
            and not resp.error
        ):
            self.cache.put(f"{req.task}:{req.cacheable_prefix()}:{req.user_content}", resp.text)
        return resp

    # --- routing policy -------------------------------------------------------
    def _candidates(self, req: LLMRequest) -> list[tuple[LLMProvider, str, str]]:
        """Ordered (provider, model, reason) list. First success wins; the list
        always ends with mock so a run can never hard-fail."""
        policy = self.settings.routing_policy
        if policy == "offline":
            return [(self.registry.mock, "mock", "policy=offline")]

        cheap = self._cheap_chain()
        smart = self._smart_chain()

        if policy == "cheap":
            chain = cheap
        elif policy == "smart":
            chain = smart
        else:  # auto
            chain = self._auto_chain(req, cheap, smart)

        # Always guarantee a terminal mock fallback.
        if not chain or chain[-1][0] is not self.registry.mock:
            chain = chain + [(self.registry.mock, "mock", "fallback=mock")]
        return chain

    def _auto_chain(self, req, cheap, smart):
        task = req.task
        if task in {"classify", "normalize", "validate", "format"}:
            return cheap  # easy → cheapest
        if task == "agent":
            return smart + cheap  # reasoning-heavy → frontier first
        if task == "extract":
            complexity = req.context.get("complexity", "simple")
            if complexity == "hard":
                return smart + cheap
            return cheap + smart  # simple extract → cheap first, frontier as backup
        return cheap

    def _cheap_chain(self) -> list[tuple[LLMProvider, str, str]]:
        chain: list[tuple[LLMProvider, str, str]] = []
        if self.registry.local and self.registry.local.available():
            chain.append((self.registry.local, self.settings.local_model,
                          "cheap: local Gemma (≈$0)"))
        if self.registry.anthropic and self.registry.anthropic.available():
            chain.append((self.registry.anthropic, self.settings.anthropic_model_cheap,
                          "cheap: Claude Haiku"))
        if self.registry.gemini and self.registry.gemini.available():
            chain.append((self.registry.gemini, self.settings.gemini_model,
                          "cheap: Gemini Flash"))
        if getattr(self.registry, "minicpm", None) and self.registry.minicpm.available():
            chain.append((self.registry.minicpm, self.settings.minicpm_model,
                          "cheap: MiniCPM-V (small model)"))
        return chain

    def _smart_chain(self) -> list[tuple[LLMProvider, str, str]]:
        chain: list[tuple[LLMProvider, str, str]] = []
        if self.registry.anthropic and self.registry.anthropic.available():
            chain.append((self.registry.anthropic, self.settings.anthropic_model_smart,
                          "smart: Claude Sonnet"))
        if self.registry.gemini and self.registry.gemini.available():
            chain.append((self.registry.gemini, self.settings.gemini_model,
                          "smart: Gemini Flash"))
        if getattr(self.registry, "minicpm", None) and self.registry.minicpm.available():
            chain.append((self.registry.minicpm, self.settings.minicpm_model,
                          "smart: MiniCPM-V (small model)"))
        if self.registry.local and self.registry.local.available():
            chain.append((self.registry.local, self.settings.local_model,
                          "smart: local Gemma (no frontier available)"))
        return chain