"""
cortex/llm/router.py — Model Routing Policy (Issue #375).

Three-tier model routing system that directs LLM calls to the appropriate
model based on the operation being performed:

    FAST      — Small, cheap, low-latency (extraction, classification, routing)
    MAIN      — Higher quality (consolidation, handoff, deep summarization)
    REASONING — Best quality, expensive (contradiction analysis, reconciliation)

The ModelRouter wraps an LLMProvider and overrides the model parameter based
on the operation-to-tier mapping. When ``model_routing.enabled`` is False
(the default), the router is a transparent passthrough — zero behavior change.

Fallback chain: REASONING → MAIN → FAST. If a tier's model is unavailable
(empty string or provider error), the router falls back down the chain.

Usage::

    router = ModelRouter(provider=my_llm, config=my_config)
    # Explicit tier routing:
    resp = await router.route_complete(
        operation="extraction", system="...", user="..."
    )
    # Or get the model for a tier:
    model = router.model_for_operation("contradiction_analysis")

Dependencies: cortex.llm.provider (LLMProvider ABC).
Used by: All LLM callsites when model_routing.enabled=True.
"""

from __future__ import annotations

import enum
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

from cortex.llm.provider import (
    CompletionResponse,
    EmbedResponse,
    ExtractionResponse,
    LLMProvider,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Tier enum
# ---------------------------------------------------------------------------

class ModelTier(enum.Enum):
    """Three-tier model classification."""
    FAST = "fast"
    MAIN = "main"
    REASONING = "reasoning"


# ---------------------------------------------------------------------------
# Default operation → tier mapping
# ---------------------------------------------------------------------------

DEFAULT_OPERATION_MAP: Dict[str, ModelTier] = {
    # FAST tier — cheap, low-latency
    "extraction": ModelTier.FAST,
    "classification": ModelTier.FAST,
    "routing": ModelTier.FAST,
    "summarization_light": ModelTier.FAST,
    "nap_summary": ModelTier.FAST,
    "key_phrase_extraction": ModelTier.FAST,
    "brain_graph_nodes": ModelTier.FAST,
    "brain_graph_edges": ModelTier.FAST,
    "brain_graph_compact": ModelTier.FAST,
    "curiosity_seeds": ModelTier.FAST,
    "curiosity_themes": ModelTier.FAST,
    "curiosity_journal": ModelTier.FAST,
    # MAIN tier — higher quality
    "consolidation": ModelTier.MAIN,
    "handoff_build": ModelTier.MAIN,
    "summarization_deep": ModelTier.MAIN,
    "session_summary": ModelTier.MAIN,
    "cornerstone_proposal": ModelTier.MAIN,
    "deprecation_review": ModelTier.MAIN,
    "dialectic_synthesis": ModelTier.MAIN,
    "query_planning": ModelTier.MAIN,
    "agentic_refinement": ModelTier.MAIN,
    # REASONING tier — best quality, expensive
    "contradiction_analysis": ModelTier.REASONING,
    "reconciliation": ModelTier.REASONING,
    "entity_resolution": ModelTier.REASONING,
    "promotion_scoring": ModelTier.REASONING,
    "cornerstone_evolution": ModelTier.REASONING,
}


# ---------------------------------------------------------------------------
# Router config dataclass
# ---------------------------------------------------------------------------

@dataclass
class ModelRoutingConfig:
    """Configuration for the model routing system.

    Attributes:
        enabled: Master toggle. When False, router is a transparent passthrough.
        tier_models: Maps each tier to a model name/ID string.
        operation_map: Maps operation names to tier names (overrides defaults).
    """
    enabled: bool = False
    tier_models: Dict[str, str] = field(default_factory=lambda: {
        "fast": "",
        "main": "",
        "reasoning": "",
    })
    operation_map: Dict[str, str] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# ModelRouter
# ---------------------------------------------------------------------------

class ModelRouter(LLMProvider):
    """Routes LLM calls to the appropriate model tier based on operation.

    Extends LLMProvider so isinstance(router, LLMProvider) returns True,
    enabling proper type-checking at callsites (audit fix D-HIGH-1).

    When routing is disabled (default), all calls pass through to the
    underlying provider with no model override — backward compatible.

    When enabled, the router:
      1. Looks up the operation in the operation map → tier
      2. Resolves the tier to a model name from config
      3. Falls back through the tier chain if model is empty
      4. Logs which tier was requested and which model was actually used

    Args:
        provider: The underlying LLMProvider (or FallbackLLMProvider).
        config: CortexConfig instance (reads model_routing fields).
    """

    def __init__(
        self,
        provider: LLMProvider,
        config: Any = None,
    ) -> None:
        # Initialize LLMProvider base (sets up _usage_log)
        super().__init__()
        self._provider = provider
        self._routing_config = self._build_routing_config(config)
        self._operation_map = self._build_operation_map()
        # Share the underlying provider's usage log for unified tracking
        self._usage_log = provider._usage_log
        # LLM usage tracker for /api/v1/llm/usage endpoint (#419)
        self._usage_tracker: Any = None
        owner_id = getattr(config, "owner_id", "default") if config is not None else "default"
        try:
            from cortex.llm.usage_tracker import LLMUsageTracker
            self._usage_tracker = LLMUsageTracker.get(owner_id)
        except Exception as _tracker_exc:  # pragma: no cover
            logger.warning("LLMUsageTracker init failed: %s", _tracker_exc)

    @property
    def provider(self) -> LLMProvider:
        """Access the underlying LLM provider."""
        return self._provider

    @property
    def enabled(self) -> bool:
        """Whether model routing is active."""
        return self._routing_config.enabled

    # ------------------------------------------------------------------ #
    # Config parsing
    # ------------------------------------------------------------------ #

    def _build_routing_config(self, config: Any) -> ModelRoutingConfig:
        """Extract routing config from CortexConfig."""
        if config is None:
            return ModelRoutingConfig()

        enabled = getattr(config, "model_routing_enabled", False)
        tier_models = {
            "fast": getattr(config, "model_routing_fast", ""),
            "main": getattr(config, "model_routing_main", ""),
            "reasoning": getattr(config, "model_routing_reasoning", ""),
        }
        # Config-level operation overrides (dict of operation_name → tier_name)
        op_map_raw = getattr(config, "model_routing_operation_map", {})
        return ModelRoutingConfig(
            enabled=enabled,
            tier_models=tier_models,
            operation_map=op_map_raw if isinstance(op_map_raw, dict) else {},
        )

    def _build_operation_map(self) -> Dict[str, ModelTier]:
        """Merge default operation map with config overrides."""
        merged = dict(DEFAULT_OPERATION_MAP)
        for op_name, tier_name in self._routing_config.operation_map.items():
            try:
                tier = ModelTier(tier_name.lower())
                merged[op_name] = tier
            except ValueError:
                logger.warning(
                    "model_routing: unknown tier %r for operation %r — ignoring",
                    tier_name, op_name,
                )
        return merged

    # ------------------------------------------------------------------ #
    # Tier resolution
    # ------------------------------------------------------------------ #

    _FALLBACK_CHAIN = [ModelTier.REASONING, ModelTier.MAIN, ModelTier.FAST]

    def tier_for_operation(self, operation: str) -> ModelTier:
        """Look up which tier an operation maps to.

        Returns ModelTier.FAST as fallback if the operation is unknown.
        """
        return self._operation_map.get(operation, ModelTier.FAST)

    def model_for_tier(self, tier: ModelTier) -> Optional[str]:
        """Resolve a tier to a concrete model name.

        Walks the fallback chain (REASONING → MAIN → FAST) if the
        requested tier's model is empty.

        Returns None if no tier has a model configured (passthrough mode).
        """
        # Find position in fallback chain
        try:
            start_idx = self._FALLBACK_CHAIN.index(tier)
        except ValueError:
            start_idx = len(self._FALLBACK_CHAIN) - 1  # default to FAST

        for i in range(start_idx, len(self._FALLBACK_CHAIN)):
            candidate = self._FALLBACK_CHAIN[i]
            model = self._routing_config.tier_models.get(candidate.value, "")
            if model:
                if candidate != tier:
                    logger.info(
                        "model_routing: tier %s has no model, falling back to %s (%s)",
                        tier.value, candidate.value, model,
                    )
                return model
        return None

    def model_for_operation(self, operation: str) -> Optional[str]:
        """Resolve an operation name to a concrete model name.

        Combines tier_for_operation + model_for_tier with fallback chain.
        Returns None if routing is disabled or no model is configured.
        """
        if not self._routing_config.enabled:
            return None
        tier = self.tier_for_operation(operation)
        return self.model_for_tier(tier)

    # ------------------------------------------------------------------ #
    # LLMProvider ABC implementation — delegates to underlying provider
    # ------------------------------------------------------------------ #

    async def complete(
        self,
        system: str,
        user: str,
        model: Optional[str] = None,
        response_format: str = "text",
        max_tokens: Optional[int] = None,
        temperature: float = 0.0,
    ) -> CompletionResponse:
        """Delegate complete() to the underlying provider (LLMProvider ABC)."""
        return await self._provider.complete(
            system=system,
            user=user,
            model=model,
            response_format=response_format,
            max_tokens=max_tokens,
            temperature=temperature,
        )

    async def extract(
        self,
        system: str,
        user: str,
        schema: dict[str, Any],
        model: Optional[str] = None,
        max_tokens: Optional[int] = None,
        temperature: float = 0.0,
    ) -> ExtractionResponse:
        """Delegate extract() to the underlying provider (LLMProvider ABC)."""
        return await self._provider.extract(
            system=system,
            user=user,
            schema=schema,
            model=model,
            max_tokens=max_tokens,
            temperature=temperature,
        )

    async def embed(
        self,
        texts: list[str],
        model: Optional[str] = None,
        input_type: str = "query",
    ) -> EmbedResponse:
        """Delegate embed() to the underlying provider (LLMProvider ABC)."""
        return await self._provider.embed(
            texts=texts,
            model=model,
            input_type=input_type,
        )

    def usage_summary(self) -> dict[str, Any]:
        """Proxy usage_summary() to the underlying provider."""
        return self._provider.usage_summary()

    def provider_names(self) -> list[str]:
        """Proxy provider_names() if underlying provider supports it."""
        if hasattr(self._provider, "provider_names"):
            return self._provider.provider_names()
        return [type(self._provider).__name__]

    # ------------------------------------------------------------------ #
    # Routed LLM calls
    # ------------------------------------------------------------------ #

    async def route_complete(
        self,
        operation: str,
        system: str,
        user: str,
        *,
        model: Optional[str] = None,
        response_format: str = "text",
        max_tokens: Optional[int] = None,
        temperature: float = 0.0,
    ) -> CompletionResponse:
        """Route a completion call through the tier system.

        When routing is enabled and no explicit model is provided, the
        router selects the model based on the operation. When routing is
        disabled, delegates directly to the underlying provider.

        Args:
            operation: Operation name (e.g. 'extraction', 'reconciliation').
            system: System prompt.
            user: User message.
            model: Explicit model override (takes precedence over routing).
            response_format: 'text' or 'json'.
            max_tokens: Max output tokens.
            temperature: Sampling temperature.
        """
        resolved_model = self._resolve_model(operation, model)
        tier = self.tier_for_operation(operation).value
        logger.debug(
            "model_routing: operation=%s tier=%s model=%s",
            operation,
            tier,
            resolved_model or "(passthrough)",
        )
        kwargs: Dict[str, Any] = {
            "system": system,
            "user": user,
            "response_format": response_format,
            "temperature": temperature,
        }
        if resolved_model:
            kwargs["model"] = resolved_model
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens
        # Removed dead-code elif that could overwrite resolved model with empty
        # string when model="" was passed (audit fix D-HIGH-2)
        _t0 = time.monotonic(); _error = False; _result = None
        try:
            _result = await self._provider.complete(**kwargs)
        except Exception:
            _error = True; raise
        finally:
            _lms = (time.monotonic() - _t0) * 1000.0
            if self._usage_tracker is not None:
                try:
                    self._usage_tracker.record_call(tier=tier, operation=operation,
                        input_tokens=getattr(_result, "tokens_in", 0) if not _error else 0,
                        output_tokens=getattr(_result, "tokens_out", 0) if not _error else 0,
                        latency_ms=_lms, error=_error)
                except Exception: pass
        return _result

    async def route_extract(
        self,
        operation: str,
        system: str,
        user: str,
        schema: dict[str, Any],
        *,
        model: Optional[str] = None,
        max_tokens: Optional[int] = None,
        temperature: float = 0.0,
    ) -> ExtractionResponse:
        """Route an extraction call through the tier system."""
        resolved_model = self._resolve_model(operation, model)
        tier = self.tier_for_operation(operation).value
        logger.debug(
            "model_routing: operation=%s tier=%s model=%s (extract)",
            operation,
            tier,
            resolved_model or "(passthrough)",
        )
        kwargs: Dict[str, Any] = {
            "system": system,
            "user": user,
            "schema": schema,
            "temperature": temperature,
        }
        if resolved_model:
            kwargs["model"] = resolved_model
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens
        # Removed dead-code elif that could overwrite resolved model with empty
        # string when model="" was passed (audit fix D-HIGH-2)
        _t0 = time.monotonic(); _error = False; _result = None
        try:
            _result = await self._provider.extract(**kwargs)
        except Exception:
            _error = True; raise
        finally:
            _lms = (time.monotonic() - _t0) * 1000.0
            if self._usage_tracker is not None:
                try:
                    self._usage_tracker.record_call(tier=tier, operation=operation,
                        input_tokens=getattr(_result, "tokens_in", 0) if not _error else 0,
                        output_tokens=getattr(_result, "tokens_out", 0) if not _error else 0,
                        latency_ms=_lms, error=_error)
                except Exception: pass
        return _result

    async def route_embed(
        self,
        operation: str,
        texts: list[str],
        *,
        model: Optional[str] = None,
        input_type: str = "query",
    ) -> EmbedResponse:
        """Route an embed call (embeddings don't have tier routing — passthrough).

        Embeddings use a dedicated model (embedding_model in config), not the
        tier system. This method exists for API consistency.
        """
        return await self._provider.embed(
            texts=texts,
            model=model,
            input_type=input_type,
        )

    # ------------------------------------------------------------------ #
    # Internal
    # ------------------------------------------------------------------ #

    def _resolve_model(self, operation: str, explicit_model: Optional[str]) -> Optional[str]:
        """Pick the model: explicit override > routing > None (passthrough)."""
        if explicit_model:
            return explicit_model
        if not self._routing_config.enabled:
            return None
        return self.model_for_operation(operation)


# ---------------------------------------------------------------------------
# Helper for callsite annotation
# ---------------------------------------------------------------------------

def resolve_routed_model(
    llm: Any,
    operation: str,
    fallback_model: Optional[str] = None,
) -> Optional[str]:
    """Resolve the model for an operation, using the router if available.

    This is the primary integration point for callsites. Usage::

        model = resolve_routed_model(self.llm, "extraction", self.model)
        await self.llm.complete(system=..., user=..., model=model)

    When the llm object is a ModelRouter with routing enabled:
        - Returns the tier-appropriate model for the operation
    When routing is disabled or llm is a plain LLMProvider:
        - Returns the fallback_model (existing behavior)

    Args:
        llm: LLMProvider or ModelRouter instance.
        operation: Operation name (e.g. 'extraction', 'reconciliation').
        fallback_model: Model to use when routing is not available/enabled.

    Returns:
        Model name string, or None if no model override needed.
    """
    router = get_router(llm)
    if router is not None and router.enabled:
        routed = router.model_for_operation(operation)
        if routed:
            logger.debug(
                "resolve_routed_model: operation=%s → model=%s (tier=%s)",
                operation, routed, router.tier_for_operation(operation).value,
            )
            return routed
    return fallback_model


def get_router(llm: Any) -> Optional[ModelRouter]:
    """Extract the ModelRouter from an llm object, if present.

    Checks if llm IS a ModelRouter, or if it has a _router attribute
    (for wrapper providers that carry a router reference).

    Returns None for MagicMock objects and any non-ModelRouter types
    to ensure test mocks don't accidentally trigger routing.
    """
    if isinstance(llm, ModelRouter):
        return llm
    attr = getattr(llm, "_router", None)
    if isinstance(attr, ModelRouter):
        return attr
    return None
