"""Model routing with health tracking and fallback for the LangGraph agent.

When the primary model fails (quota, capacity, timeout), automatically
falls back to an alternative. Supports task-based routing (use haiku for
summarization, sonnet for complex reasoning).

Ported from Gemini CLI's modelRouterService.ts and Codex CLI's models_manager.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional

logger = logging.getLogger(__name__)


class ModelHealth(str, Enum):
    """Health status of a model."""
    HEALTHY = "healthy"
    RETRY_ONCE = "retry_once"   # Try once per turn, then skip
    TERMINAL = "terminal"        # Dead until explicit reset (quota exhausted)


@dataclass
class ModelInfo:
    """Information about an available model."""
    model_id: str                  # e.g., "claude-sonnet-4-20250514"
    family: str                    # haiku / sonnet / opus
    health: ModelHealth = ModelHealth.HEALTHY
    last_error: Optional[str] = None
    last_error_time: float = 0.0
    cost_per_1k_input: float = 0.003
    cost_per_1k_output: float = 0.015


# Default model families with cost info
_MODEL_FAMILIES = {
    "haiku": {"input": 0.001, "output": 0.005},
    "sonnet": {"input": 0.003, "output": 0.015},
    "opus": {"input": 0.015, "output": 0.075},
}


def _detect_family(model_id: str) -> str:
    """Detect model family from model ID."""
    lower = model_id.lower()
    for family in ("haiku", "sonnet", "opus"):
        if family in lower:
            return family
    return "sonnet"  # Default


class ModelRouter:
    """Route model requests with health tracking and fallback.

    Supports:
    - Error-driven fallback: when primary fails, try next in chain
    - Task-based routing: use cheap models for summarization
    - Per-turn retry resets: RETRY_ONCE models get one chance each turn
    """

    def __init__(self, primary_model_id: str) -> None:
        self._primary_id = primary_model_id
        primary_family = _detect_family(primary_model_id)

        # Build fallback chain: primary → sonnet → haiku
        self._models: list[ModelInfo] = []
        families_added = set()

        # Primary
        costs = _MODEL_FAMILIES.get(primary_family, _MODEL_FAMILIES["sonnet"])
        self._models.append(ModelInfo(
            model_id=primary_model_id,
            family=primary_family,
            cost_per_1k_input=costs["input"],
            cost_per_1k_output=costs["output"],
        ))
        families_added.add(primary_family)

        # Fallbacks (skip if already in chain)
        fallback_order = ["sonnet", "haiku"]
        for family in fallback_order:
            if family in families_added:
                continue
            # Use a default model ID for the family
            model_id = _default_model_id(family)
            costs = _MODEL_FAMILIES[family]
            self._models.append(ModelInfo(
                model_id=model_id,
                family=family,
                cost_per_1k_input=costs["input"],
                cost_per_1k_output=costs["output"],
            ))
            families_added.add(family)

    def select(self, task_hint: str = "") -> str:
        """Select the best available model.

        Args:
            task_hint: Optional hint for task-based routing.
                "summarize" → prefer haiku (cheap)
                "complex_reasoning" → prefer sonnet/opus
                "" → use primary

        Returns:
            Model ID string.
        """
        # Task-based routing
        if task_hint == "summarize":
            # Find cheapest healthy model
            for model in reversed(self._models):
                if model.health != ModelHealth.TERMINAL:
                    return model.model_id

        # Normal selection: first healthy model
        for model in self._models:
            if model.health == ModelHealth.HEALTHY:
                return model.model_id
            if model.health == ModelHealth.RETRY_ONCE:
                # Give it one more try this turn
                return model.model_id

        # All models unhealthy — try primary anyway
        logger.warning("All models unhealthy, falling back to primary: %s", self._primary_id)
        return self._primary_id

    def mark_error(self, model_id: str, error: Exception) -> None:
        """Classify an error and update model health.

        Error classification:
        - 429/quota/rate_limit → TERMINAL (don't retry this turn)
        - 529/overloaded/503 → RETRY_ONCE (try once more, then skip)
        - timeout → RETRY_ONCE
        - other → HEALTHY (transient, retry normally)
        """
        err_str = str(error).lower()
        model = self._find_model(model_id)
        if not model:
            return

        model.last_error = str(error)[:200]
        model.last_error_time = time.time()

        if any(kw in err_str for kw in ("429", "quota", "rate_limit", "rate limit")):
            model.health = ModelHealth.TERMINAL
            logger.warning("Model %s marked TERMINAL (quota/rate limit): %s",
                           model_id, str(error)[:100])
        elif any(kw in err_str for kw in ("529", "overloaded", "503", "502", "timeout")):
            model.health = ModelHealth.RETRY_ONCE
            logger.warning("Model %s marked RETRY_ONCE (transient): %s",
                           model_id, str(error)[:100])
        else:
            # Transient error, keep healthy
            logger.info("Model %s transient error (keeping HEALTHY): %s",
                        model_id, str(error)[:100])

    def mark_success(self, model_id: str) -> None:
        """Mark a model as healthy after a successful call."""
        model = self._find_model(model_id)
        if model and model.health == ModelHealth.RETRY_ONCE:
            model.health = ModelHealth.HEALTHY

    def reset_turn(self) -> None:
        """Reset RETRY_ONCE models for a new turn.

        Called at the start of each agent turn. RETRY_ONCE models get
        another chance. TERMINAL models stay dead.
        """
        for model in self._models:
            if model.health == ModelHealth.RETRY_ONCE:
                model.health = ModelHealth.HEALTHY

    def get_status(self) -> list[dict]:
        """Get health status of all models for display."""
        return [
            {
                "model_id": m.model_id,
                "family": m.family,
                "health": m.health.value,
                "last_error": m.last_error,
            }
            for m in self._models
        ]

    def _find_model(self, model_id: str) -> Optional[ModelInfo]:
        """Find a model by ID."""
        for model in self._models:
            if model.model_id == model_id:
                return model
        return None


def _default_model_id(family: str) -> str:
    """Get a default model ID for a family."""
    defaults = {
        "haiku": "claude-haiku-4-5-20251001",
        "sonnet": "claude-sonnet-4-5-20250929",
        "opus": "claude-opus-4-20250514",
    }
    return defaults.get(family, "claude-sonnet-4-5-20250929")
