"""Model router — resolves agent model tiers to specific LLM models.

Architecture:
- Agents specify `model_tier` in their definition (reasoning, capable, fast)
- Users pick a `quality_profile` in Settings (quality, balanced, economy)
- This module resolves: profile + provider + tier → specific model name

Configuration loaded from skills/_meta/model-routing/config/models.yaml.
Falls back to hardcoded defaults if YAML not found.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

log = logging.getLogger(__name__)

# Hardcoded fallback — used when YAML config not found
_FALLBACK_PROFILES: dict[str, dict[str, dict[str, str]]] = {
    "quality": {
        "anthropic": {
            "reasoning": "claude-opus-4-20250514",
            "capable": "claude-sonnet-4-20250514",
            "fast": "claude-haiku-4-5-20251001",
        },
        "openai": {
            "reasoning": "gpt-4.1",
            "capable": "gpt-4.1",
            "fast": "gpt-4.1-mini",
        },
    },
    "balanced": {
        "anthropic": {
            "reasoning": "claude-sonnet-4-20250514",
            "capable": "claude-sonnet-4-20250514",
            "fast": "claude-haiku-4-5-20251001",
        },
        "openai": {
            "reasoning": "gpt-4.1",
            "capable": "gpt-4.1-mini",
            "fast": "gpt-4.1-nano",
        },
    },
    "economy": {
        "anthropic": {
            "reasoning": "claude-sonnet-4-20250514",
            "capable": "claude-haiku-4-5-20251001",
            "fast": "claude-haiku-4-5-20251001",
        },
        "openai": {
            "reasoning": "gpt-4.1-mini",
            "capable": "gpt-4.1-mini",
            "fast": "gpt-4.1-nano",
        },
    },
}

_DEFAULT_TIER = "capable"

# arch:deterministic — fallback capability matrix when capabilities.yaml is
# missing or malformed. Conservative defaults: cloud tiers get reliable
# everything, local models get patched tool calling with unreliable JSON
# and no vision. The real vocabulary lives in the YAML; this is the
# safety net so first boot / corrupted YAML never crashes capability
# queries. Identical shape to capabilities.yaml.
_FALLBACK_CAPABILITIES: dict[str, Any] = {
    "profiles": {
        p: {
            t: {
                "tool_calling": "reliable",
                "json_mode": "reliable",
                "vision": "reliable" if p == "quality" else ("reliable" if t != "fast" else "unavailable"),
                "streaming": "reliable",
                "ctx_window": 200000 if p == "quality" else 128000,
            }
            for t in ("reasoning", "capable", "fast")
        }
        for p in ("quality", "balanced", "economy")
    },
    "local_models": {
        "_default": {
            "tool_calling": "patched",
            "json_mode": "unreliable",
            "vision": "unavailable",
            "streaming": "reliable",
            "ctx_window": 8000,
        }
    },
}

# Capability levels that mean "the feature works and you can rely on it
# either directly or via the skill-driven patches." Everything else
# (unreliable, unavailable, unknown) means callers MUST have a fallback
# or degrade explicitly. Kept as a module-level frozenset so downstream
# callers can import it without re-typing the literal.
CAPABILITY_OK_LEVELS = frozenset({"reliable", "patched"})


def _deep_merge_capabilities(base: dict, override: dict) -> dict:
    """Recursively merge `override` on top of `base`.

    Leaves in `override` replace leaves in `base`; nested dicts merge
    per-key. Lists and scalars in `override` fully replace their
    counterparts in `base` — deep-merging into lists is rarely what a
    YAML author means and causes surprising behavior. Returns a new
    dict; neither input is mutated.

    Used specifically by the capability-matrix loader so a user YAML
    that declares only some fields for a tier still inherits the
    remaining fields from `_FALLBACK_CAPABILITIES`. Without this, a
    single-field YAML override wiped all other fields for that tier.
    """
    if not isinstance(override, dict):
        return override if override is not None else base
    result: dict = {}
    # Start from base keys
    for key, base_val in base.items():
        if key in override:
            ov_val = override[key]
            if isinstance(base_val, dict) and isinstance(ov_val, dict):
                result[key] = _deep_merge_capabilities(base_val, ov_val)
            else:
                result[key] = ov_val
        else:
            result[key] = base_val
    # Add keys that are in override but not in base
    for key, ov_val in override.items():
        if key not in result:
            result[key] = ov_val
    return result


class ModelRouter:
    """Resolves model tiers to specific model names."""

    # arch:deterministic — cache TTL for the local-stack reachability probe.
    # 30 s balances "freshly accurate when Ollama just came up" against
    # "don't probe localhost on every single resolve()". Calls within the
    # window reuse the prior verdict.
    _LOCAL_REACHABLE_TTL_SECONDS = 30.0

    def __init__(
        self,
        provider: str = "anthropic",
        profile: str = "balanced",
        default_model: str = "",
        skills_dir: Path | None = None,
        *,
        mode: str = "cloud",
    ) -> None:
        self._provider = provider
        self._profile = profile
        self._default_model = default_model
        self._mode = mode
        self._profiles: dict[str, Any] = {}
        self._capabilities: dict[str, Any] = {}
        # Hybrid-routing state — populated by _load() from models.yaml's
        # `routing:` block, defaults preserve pre-routing behavior.
        self._cloud_use_local_for_tiers: frozenset[str] = frozenset()
        self._honor_prefer_local: bool = False
        # Reachability cache — _local_reachable_at is the monotonic time of
        # last probe, _local_reachable is the result. Both reset on
        # `invalidate_local_cache()`.
        self._local_reachable: bool | None = None
        self._local_reachable_at: float = 0.0
        self._load(skills_dir)

    def _load(self, skills_dir: Path | None) -> None:
        """Load model routing config from YAML, fall back to hardcoded."""
        self._profiles = _FALLBACK_PROFILES
        self._capabilities = _FALLBACK_CAPABILITIES

        if not skills_dir:
            return

        try:
            import yaml
        except Exception:
            return

        models_path = skills_dir / "_meta" / "model-routing" / "config" / "models.yaml"
        if models_path.is_file():
            try:
                data = yaml.safe_load(models_path.read_text(encoding="utf-8")) or {}
                self._profiles = data.get("profiles", {}) or _FALLBACK_PROFILES
                routing = data.get("routing", {}) or {}
                tiers = routing.get("cloud_use_local_for_tiers", []) or []
                if isinstance(tiers, list):
                    # Coerce to strings + lower() so a YAML author who writes
                    # `[Fast, FAST]` doesn't silently mismatch our internal
                    # tier-name vocabulary (which is lowercase).
                    self._cloud_use_local_for_tiers = frozenset(
                        str(t).strip().lower() for t in tiers if t
                    )
                self._honor_prefer_local = bool(routing.get("honor_prefer_local", False))
                log.debug(
                    "Model routing loaded from YAML: %d profiles, hybrid_tiers=%s, honor_prefer_local=%s",
                    len(self._profiles),
                    sorted(self._cloud_use_local_for_tiers),
                    self._honor_prefer_local,
                )
            except Exception:
                log.warning("Failed to load model routing YAML, using fallback")

        caps_path = skills_dir / "_meta" / "model-routing" / "config" / "capabilities.yaml"
        if caps_path.is_file():
            try:
                data = yaml.safe_load(caps_path.read_text(encoding="utf-8")) or {}
                # Deep-merge so a partial YAML (e.g. a user who only sets
                # `profiles.balanced.fast.vision: reliable`) still inherits
                # the other fields for that tier from `_FALLBACK_CAPABILITIES`.
                # Shallow-merge was a bug — it silently wiped ctx_window,
                # streaming, tool_calling, json_mode when a user declared
                # any single cap for a tier.
                merged = _deep_merge_capabilities(_FALLBACK_CAPABILITIES, data)
                self._capabilities = merged
                log.debug(
                    "Capability matrix loaded: %d profiles, %d local model families",
                    len(merged.get("profiles", {}) or {}),
                    len(merged.get("local_models", {}) or {}),
                )
            except Exception:
                log.warning("Failed to load capabilities YAML, using fallback")

    def resolve(
        self,
        model_tier: str | None = None,
        *,
        prefer_local: bool = False,
        force_cloud: bool = False,
    ) -> str:
        """Resolve a model tier to a specific model name.

        Args:
            model_tier: "reasoning", "capable", or "fast". None → default tier.
            prefer_local: caller hint that THIS specific call should run on
                the local stack if reachable (e.g. gated-content writes,
                privacy-sensitive flows). In cloud mode, honored only when
                routing.honor_prefer_local is True. In local mode, ignored
                (everything is already local).
            force_cloud: caller hint that THIS specific call MUST stay on
                the cloud provider regardless of hybrid-routing rules
                (e.g. interactive chat, where latency + quality matter
                more than cost). Always respected.

        Returns:
            Specific model name (e.g. "openai/gpt-5.4" or "ollama/qwen2.5:7b").
        """
        tier = (model_tier or _DEFAULT_TIER).lower()

        # ---- hybrid routing decision (cloud mode only) ----------------
        # In local mode the active provider IS ollama, so resolve() already
        # returns the local model. Hybrid only applies when the active
        # provider is a cloud provider AND this specific call wants local.
        if self._mode == "cloud" and not force_cloud:
            should_route_local = (
                tier in self._cloud_use_local_for_tiers
                or (prefer_local and self._honor_prefer_local)
            )
            if should_route_local and self._local_reachable_cached():
                local_model = self._lookup_local_model(tier)
                if local_model:
                    return self._ensure_ollama_prefix(local_model)
                # Fall through — local stack reachable but no model assigned
                # to this tier. The cloud answer is still better than 500ing.

        # ---- standard cloud / native-local resolution -----------------
        profile_data = self._profiles.get(self._profile, {})
        provider_data = profile_data.get(self._provider, {})
        model = provider_data.get(tier)

        if model:
            return self._ensure_prefix(model)

        # Fallback chain: default model → capable tier → provider default
        if self._default_model:
            return self._ensure_prefix(self._default_model)

        capable = provider_data.get("capable")
        if capable:
            return self._ensure_prefix(capable)

        return self._ensure_prefix(self._default_model or "claude-sonnet-4-20250514")

    def _ensure_prefix(self, model: str) -> str:
        """Ensure model name has provider/ prefix for LiteLLM routing."""
        if "/" in model:
            return model
        return f"{self._provider}/{model}"

    def _ensure_ollama_prefix(self, model: str) -> str:
        """Ensure a routed-to-local model gets the ollama/ prefix LiteLLM needs."""
        if "/" in model:
            return model
        return f"ollama/{model}"

    def _lookup_local_model(self, tier: str) -> str:
        """Read profiles.local.ollama.<tier> from the loaded YAML.

        Returns "" when no model is mapped to that tier. Used by the
        hybrid-routing path to find the local model that should answer
        a routed cloud-tier call.
        """
        local_profile = self._profiles.get("local", {})
        ollama_block = local_profile.get("ollama", {}) if isinstance(local_profile, dict) else {}
        return str(ollama_block.get(tier, "")) if isinstance(ollama_block, dict) else ""

    def _local_reachable_cached(self) -> bool:
        """Probe the local Ollama stack with a TTL cache.

        Returns the cached verdict if probed within `_LOCAL_REACHABLE_TTL_SECONDS`,
        otherwise re-probes /api/tags. Failures (Ollama down / not installed)
        return False — a missing local stack must NEVER raise here, the
        caller will fall through to the cloud path.
        """
        import time
        now = time.monotonic()
        if (
            self._local_reachable is not None
            and (now - self._local_reachable_at) < self._LOCAL_REACHABLE_TTL_SECONDS
        ):
            return self._local_reachable

        try:
            import httpx
            with httpx.Client(timeout=2.0) as client:
                r = client.get("http://localhost:11434/api/tags")
                self._local_reachable = r.status_code == 200
        except Exception:
            self._local_reachable = False
        self._local_reachable_at = now
        return self._local_reachable

    def invalidate_local_cache(self) -> None:
        """Forget the cached reachability verdict.

        Call this when Ollama state changes (model installed, daemon
        started/stopped) so the next resolve() re-probes immediately.
        """
        self._local_reachable = None
        self._local_reachable_at = 0.0

    def explain(
        self,
        model_tier: str | None = None,
        *,
        prefer_local: bool = False,
        force_cloud: bool = False,
    ) -> dict[str, Any]:
        """Pure read — describe what resolve() would return and why.

        Used by the /api/llm/local-stack endpoint and by tests to verify
        the routing decision without firing the LLM call. Never mutates
        cache state — calls _local_reachable_cached only if needed.
        """
        tier = (model_tier or _DEFAULT_TIER).lower()
        out: dict[str, Any] = {
            "tier": tier,
            "mode": self._mode,
            "prefer_local": bool(prefer_local),
            "force_cloud": bool(force_cloud),
            "hybrid_tiers": sorted(self._cloud_use_local_for_tiers),
            "honor_prefer_local": self._honor_prefer_local,
        }
        if self._mode == "cloud" and not force_cloud:
            should_route_local = (
                tier in self._cloud_use_local_for_tiers
                or (prefer_local and self._honor_prefer_local)
            )
            if should_route_local:
                local_model = self._lookup_local_model(tier)
                reachable = self._local_reachable_cached()
                if local_model and reachable:
                    out["routed_to"] = "local"
                    out["model"] = self._ensure_ollama_prefix(local_model)
                    out["reason"] = (
                        "tier in cloud_use_local_for_tiers"
                        if tier in self._cloud_use_local_for_tiers
                        else "prefer_local + honor_prefer_local"
                    )
                    return out
                out["routed_to"] = "cloud"
                out["model"] = self.resolve(model_tier, force_cloud=True)
                out["reason"] = (
                    "no local model mapped to tier"
                    if not local_model
                    else "local stack unreachable"
                )
                return out
        out["routed_to"] = "cloud" if self._mode == "cloud" else "local"
        out["model"] = self.resolve(model_tier, force_cloud=self._mode == "cloud")
        out["reason"] = "force_cloud" if force_cloud else "tier not in hybrid list"
        return out

    @property
    def profile(self) -> str:
        return self._profile

    @profile.setter
    def profile(self, value: str) -> None:
        self._profile = value

    def get_profile_summary(self) -> dict[str, str]:
        """Get the model assignments for the current profile."""
        profile_data = self._profiles.get(self._profile, {})
        provider_data = profile_data.get(self._provider, {})
        return dict(provider_data)

    # ----- capability matrix ---------------------------------------------
    #
    # Entry points for "does this tier/model support X?" queries. Callers
    # should check BEFORE assuming a feature works. The goal is to make
    # "cloud-only" features explicit and to prevent silent degradation on
    # local mode (journal_today regression of 2026-04-18).

    def capability(
        self,
        name: str,
        *,
        tier: str | None = None,
        model: str | None = None,
    ) -> Any:
        """Return the declared capability level for a tier or local model.

        Args:
            name: capability name — e.g. ``tool_calling``, ``json_mode``,
                ``vision``, ``streaming``, ``ctx_window``.
            tier: cloud-mode tier (``reasoning``/``capable``/``fast``).
                Defaults to the default tier.
            model: local-mode Ollama model name (e.g. ``qwen2.5:14b``).
                If provided, ``tier`` is ignored and capability is looked
                up in ``local_models`` by longest-prefix family match.

        Returns:
            For binary caps: ``"reliable"`` / ``"patched"`` / ``"unreliable"``
            / ``"unavailable"`` / ``"unknown"``. For ``ctx_window``: an int.
            Never raises — an unknown name returns ``"unknown"`` (or 0 for
            numeric caps) so callers can treat it as "do not rely on this."
        """
        table = self._lookup_capability_table(tier=tier, model=model)
        if name not in table:
            # ctx_window has a numeric default; everything else is a string
            if name == "ctx_window":
                return 0
            return "unknown"
        return table[name]

    def capability_ok(
        self,
        name: str,
        *,
        tier: str | None = None,
        model: str | None = None,
    ) -> bool:
        """True if the capability level indicates the feature works.

        "Works" means either ``reliable`` (direct support) or ``patched``
        (works via skill/prompt support such as local-intents). Anything
        else — ``unreliable`` / ``unavailable`` / ``unknown`` — returns
        False and callers must have a fallback or degrade explicitly.
        """
        level = self.capability(name, tier=tier, model=model)
        return isinstance(level, str) and level in CAPABILITY_OK_LEVELS

    def context_window(
        self,
        *,
        tier: str | None = None,
        model: str | None = None,
    ) -> int:
        """Return the declared context window in tokens (0 if unknown)."""
        raw = self.capability("ctx_window", tier=tier, model=model)
        try:
            return int(raw)
        except (TypeError, ValueError):
            return 0

    def _lookup_capability_table(
        self,
        *,
        tier: str | None,
        model: str | None,
    ) -> dict[str, Any]:
        """Resolve the correct capability block for (tier) or (model)."""
        if model:
            families = self._capabilities.get("local_models") or {}
            if not isinstance(families, dict):
                families = {}
            # Longest-prefix match — "qwen2.5:14b" matches "qwen2.5" before "_default".
            # Sorted descending by key length so sub-family overrides parent when
            # declared (e.g. "llama3.1" wins over "llama3" for "llama3.1:70b").
            for family in sorted(families.keys(), key=lambda k: -len(str(k))):
                if family == "_default":
                    continue
                if str(model).startswith(str(family)):
                    block = families.get(family)
                    if isinstance(block, dict):
                        return block
            default_block = families.get("_default")
            if isinstance(default_block, dict):
                return default_block
            return {}

        profiles = self._capabilities.get("profiles") or {}
        if not isinstance(profiles, dict):
            return {}
        profile_block = profiles.get(self._profile) or {}
        if not isinstance(profile_block, dict):
            return {}
        resolved_tier = tier or _DEFAULT_TIER
        tier_block = profile_block.get(resolved_tier) or {}
        if isinstance(tier_block, dict):
            return tier_block
        return {}
