"""Model registry: 10 providers × 19 models with cost/capability metadata.

Cost figures are per-million-tokens (March 2026).  The router uses these to
pick the cheapest model that satisfies the task requirements.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum


class TaskTier(str, Enum):
    """Complexity tier that determines which model class is needed."""

    TRIVIAL = "trivial"          # autocomplete, short answers
    OPERATIONAL = "operational"  # small edits, simple generation
    STANDARD = "standard"        # implementation, moderate complexity
    COMPLEX = "complex"          # multi-file refactor, architecture
    DEEP_REASONING = "deep"      # hard bugs, thinking-heavy tasks

    @property
    def rank(self) -> int:
        return _TIER_RANKS[self]


_TIER_RANKS: dict["TaskTier", int] = {
    TaskTier.TRIVIAL: 0,
    TaskTier.OPERATIONAL: 1,
    TaskTier.STANDARD: 2,
    TaskTier.COMPLEX: 3,
    TaskTier.DEEP_REASONING: 4,
}


class Capability(str, Enum):
    CODE = "code"
    REASONING = "reasoning"
    THINKING = "thinking"        # extended thinking / chain-of-thought
    FAST = "fast"
    VISION = "vision"
    LONG_CONTEXT = "long_context"  # >200K tokens
    TOOL_USE = "tool_use"
    TOOLS = "tool_use"


@dataclass(frozen=True)
class ModelSpec:
    """Immutable specification for a single model deployment."""

    id: str                          # unique key, e.g. "anthropic/opus-4.6"
    provider: str                    # provider name matching config key
    litellm_model: str               # string passed to litellm.completion()
    display_name: str
    input_cost: float                # $/1M input tokens
    output_cost: float               # $/1M output tokens
    context_window: int              # max tokens
    params_b: float = 0.0            # parameter count in billions (0 = unknown)
    is_local: bool = False           # runs locally (Ollama, llama.cpp, etc.)
    is_free: bool = False            # zero cost (free tier / local)
    capabilities: frozenset[Capability] = field(default_factory=frozenset)
    min_tier: TaskTier = TaskTier.TRIVIAL
    max_tier: TaskTier = TaskTier.DEEP_REASONING
    cache_read_cost: float = 0.0     # $/1M tokens for prompt cache read
    batch_discount: float = 1.0      # multiplier (0.5 = 50% off)
    speed_tps: int = 80              # approx output tokens/sec
    notes: str = ""

    @property
    def effective_cost(self) -> float:
        """Blended cost per 1M tokens (assuming 1:1 input:output ratio)."""
        return (self.input_cost + self.output_cost) / 2

    def summary_dict(self) -> dict:
        """Return a dict suitable for API/CLI display."""
        return {
            "id": self.id,
            "provider": self.provider,
            "litellm_model": self.litellm_model,
            "display_name": self.display_name,
            "params_b": self.params_b,
            "is_local": self.is_local,
            "is_free": self.is_free,
            "input_cost": self.input_cost,
            "output_cost": self.output_cost,
            "effective_cost": self.effective_cost,
            "context_window": self.context_window,
            "capabilities": sorted(c.value for c in self.capabilities),
            "min_tier": self.min_tier.value,
            "max_tier": self.max_tier.value,
            "speed_tps": self.speed_tps,
            "notes": self.notes,
        }


# ============================================================
# MODEL REGISTRY — 19 models across 10 providers
# Sorted by effective cost ascending within each tier
# ============================================================

MODELS: dict[str, ModelSpec] = {}


def _r(spec: ModelSpec) -> ModelSpec:
    MODELS[spec.id] = spec
    return spec


# --- Tier: TRIVIAL / OPERATIONAL (< $3 effective) ---

_r(ModelSpec(
    id="cerebras/llama-3.3-70b",
    provider="cerebras",
    litellm_model="cerebras/llama-3.3-70b",
    display_name="Llama 3.3 70B (Cerebras)",
    input_cost=0.60, output_cost=0.60,
    context_window=128_000,
    params_b=70.0,
    capabilities=frozenset({Capability.CODE, Capability.FAST}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.OPERATIONAL,
    speed_tps=2200,
    notes="Fastest inference available; good for simple tasks",
))

_r(ModelSpec(
    id="deepseek/v3",
    provider="deepseek",
    litellm_model="deepseek/deepseek-chat",
    display_name="DeepSeek V3",
    input_cost=0.27, output_cost=1.10,
    context_window=128_000,
    params_b=671.0,
    capabilities=frozenset({Capability.CODE, Capability.REASONING}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.STANDARD,
    cache_read_cost=0.07,
    speed_tps=60,
    notes="Best cost/quality ratio for routine coding; MoE 671B (37B active)",
))

_r(ModelSpec(
    id="groq/llama-3.3-70b",
    provider="groq",
    litellm_model="groq/llama-3.3-70b-versatile",
    display_name="Llama 3.3 70B (Groq)",
    input_cost=0.59, output_cost=0.79,
    context_window=128_000,
    params_b=70.0,
    capabilities=frozenset({Capability.CODE, Capability.FAST}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.OPERATIONAL,
    speed_tps=1200,
    notes="Free tier available; ultra-fast",
))

_r(ModelSpec(
    id="anthropic/haiku-4.5",
    provider="anthropic",
    litellm_model="anthropic/claude-haiku-4-5-20251001",
    display_name="Claude Haiku 4.5",
    input_cost=1.00, output_cost=5.00,
    context_window=200_000,
    params_b=8.0,
    capabilities=frozenset({Capability.CODE, Capability.TOOL_USE}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.STANDARD,
    cache_read_cost=0.10,
    batch_discount=0.5,
    speed_tps=150,
    notes="Fast, cheap, good for debug and validation",
))

_r(ModelSpec(
    id="google/gemini-2.5-flash",
    provider="google",
    litellm_model="gemini/gemini-2.5-flash-preview-05-20",
    display_name="Gemini 2.5 Flash",
    input_cost=0.15, output_cost=0.60,
    context_window=1_048_576,
    params_b=0.0,
    is_free=True,
    capabilities=frozenset({Capability.CODE, Capability.FAST, Capability.LONG_CONTEXT, Capability.THINKING}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.STANDARD,
    speed_tps=350,
    notes="Free tier: 500 req/day; great for planning",
))

# --- Tier: STANDARD (< $10 effective) ---

_r(ModelSpec(
    id="google/gemini-2.5-pro",
    provider="google",
    litellm_model="gemini/gemini-2.5-pro-preview-05-06",
    display_name="Gemini 2.5 Pro",
    input_cost=1.25, output_cost=10.00,
    context_window=1_048_576,
    params_b=0.0,
    is_free=True,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING,
                           Capability.LONG_CONTEXT, Capability.TOOL_USE}),
    min_tier=TaskTier.OPERATIONAL, max_tier=TaskTier.COMPLEX,
    speed_tps=200,
    notes="Free tier: 1000 req/day via AI Studio; excellent for architecture",
))

_r(ModelSpec(
    id="google/gemini-3-flash-preview",
    provider="google",
    litellm_model="gemini/gemini-3-flash-preview",
    display_name="Gemini 3 Flash Preview",
    input_cost=0.50, output_cost=3.00,
    context_window=1_048_576,
    params_b=0.0,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING, Capability.TOOL_USE, Capability.LONG_CONTEXT}),
    min_tier=TaskTier.OPERATIONAL, max_tier=TaskTier.COMPLEX,
    cache_read_cost=0.05,
    batch_discount=0.5,
    speed_tps=250,
    notes="Balanced preview model for agentic coding and long-context tasks",
))

_r(ModelSpec(
    id="openai/gpt-4.1",
    provider="openai",
    litellm_model="openai/gpt-4.1",
    display_name="GPT-4.1",
    input_cost=2.00, output_cost=8.00,
    context_window=1_047_576,
    params_b=0.0,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.TOOL_USE,
                           Capability.LONG_CONTEXT}),
    min_tier=TaskTier.OPERATIONAL, max_tier=TaskTier.COMPLEX,
    speed_tps=100,
    notes="1M context; solid coding; direct GPT-4.5 successor",
))

_r(ModelSpec(
    id="mistral/codestral-2501",
    provider="mistral",
    litellm_model="mistral/codestral-2501",
    display_name="Codestral 2501",
    input_cost=0.30, output_cost=0.90,
    context_window=256_000,
    params_b=22.0,
    capabilities=frozenset({Capability.CODE, Capability.FAST}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.STANDARD,
    speed_tps=180,
    notes="Specialized code model; very cheap",
))

# --- Tier: COMPLEX ($10–20 effective) ---

_r(ModelSpec(
    id="openrouter/kimi-k2.5",
    provider="openrouter",
    litellm_model="openrouter/moonshotai/kimi-k2.5",
    display_name="Kimi K2.5 (Moonshot)",
    input_cost=1.00, output_cost=4.00,
    context_window=131_072,
    params_b=0.0,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING}),
    min_tier=TaskTier.OPERATIONAL, max_tier=TaskTier.COMPLEX,
    speed_tps=80,
    notes="Agent Swarm mode; good thinking; via OpenRouter",
))

_r(ModelSpec(
    id="openrouter/swe-1.5",
    provider="openrouter",
    litellm_model="openrouter/all-hands/openhands-lm-32b-v0.1",
    display_name="SWE-1.5 (OpenHands)",
    input_cost=0.80, output_cost=0.80,
    context_window=131_072,
    params_b=32.0,
    capabilities=frozenset({Capability.CODE, Capability.TOOL_USE}),
    min_tier=TaskTier.OPERATIONAL, max_tier=TaskTier.STANDARD,
    speed_tps=90,
    notes="Specialized for SWE tasks; cheap via OpenRouter",
))

_r(ModelSpec(
    id="openrouter/step-3.5-flash:free",
    provider="openrouter",
    litellm_model="openrouter/stepfun/step-3.5-flash:free",
    display_name="Step 3.5 Flash (free)",
    input_cost=0.0, output_cost=0.0,
    context_window=256_000,
    params_b=196.0,
    is_free=True,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING, Capability.FAST, Capability.LONG_CONTEXT}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.COMPLEX,
    speed_tps=240,
    notes="Free via OpenRouter; flash MoE with strong reasoning and long context",
))

_r(ModelSpec(
    id="openrouter/nemotron-3-nano-30b-a3b:free",
    provider="openrouter",
    litellm_model="openrouter/nvidia/nemotron-3-nano-30b-a3b:free",
    display_name="Nemotron 3 Nano 30B A3B (free)",
    input_cost=0.0, output_cost=0.0,
    context_window=256_000,
    params_b=30.0,
    is_free=True,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING, Capability.TOOL_USE, Capability.FAST, Capability.LONG_CONTEXT}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.COMPLEX,
    speed_tps=200,
    notes="Free via OpenRouter; efficient agentic MoE model",
))

_r(ModelSpec(
    id="openrouter/gpt-oss-120b",
    provider="openrouter",
    litellm_model="openrouter/openai/gpt-oss-120b",
    display_name="GPT-OSS 120B",
    input_cost=0.039, output_cost=0.19,
    context_window=131_072,
    params_b=120.0,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING, Capability.TOOL_USE, Capability.LONG_CONTEXT}),
    min_tier=TaskTier.OPERATIONAL, max_tier=TaskTier.DEEP_REASONING,
    speed_tps=100,
    notes="Open-weight MoE from OpenAI; high-reasoning and agentic",
))

_r(ModelSpec(
    id="openrouter/qwen3-235b-a22b-2507",
    provider="openrouter",
    litellm_model="openrouter/qwen/qwen3-235b-a22b-2507",
    display_name="Qwen3 235B A22B Instruct 2507",
    input_cost=0.071, output_cost=0.10,
    context_window=262_144,
    params_b=235.0,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING, Capability.LONG_CONTEXT}),
    min_tier=TaskTier.OPERATIONAL, max_tier=TaskTier.DEEP_REASONING,
    speed_tps=90,
    notes="Large Qwen MoE; multilingual reasoning and coding",
))

_r(ModelSpec(
    id="openai/gpt-5.4",
    provider="openai",
    litellm_model="openai/gpt-5.4",
    display_name="GPT-5.4",
    input_cost=2.50, output_cost=15.00,
    context_window=1_047_576,
    params_b=0.0,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING,
                           Capability.LONG_CONTEXT, Capability.TOOL_USE, Capability.VISION}),
    min_tier=TaskTier.STANDARD, max_tier=TaskTier.DEEP_REASONING,
    speed_tps=90,
    notes="Deep reasoning; good for hard debugging",
))

# --- Tier: DEEP REASONING ($15+ effective) ---

_r(ModelSpec(
    id="anthropic/opus-4.6",
    provider="anthropic",
    litellm_model="anthropic/claude-opus-4-6",
    display_name="Claude Opus 4.6",
    input_cost=5.00, output_cost=25.00,
    context_window=1_000_000,
    params_b=0.0,
    capabilities=frozenset({Capability.CODE, Capability.REASONING, Capability.THINKING,
                           Capability.LONG_CONTEXT, Capability.TOOL_USE, Capability.VISION}),
    min_tier=TaskTier.COMPLEX, max_tier=TaskTier.DEEP_REASONING,
    cache_read_cost=0.50,
    batch_discount=0.5,
    speed_tps=60,
    notes="Highest quality coding; use only for complex tasks",
))

_r(ModelSpec(
    id="deepseek/r1",
    provider="deepseek",
    litellm_model="deepseek/deepseek-reasoner",
    display_name="DeepSeek R1",
    input_cost=0.55, output_cost=2.19,
    context_window=128_000,
    params_b=671.0,
    capabilities=frozenset({Capability.REASONING, Capability.THINKING}),
    min_tier=TaskTier.STANDARD, max_tier=TaskTier.DEEP_REASONING,
    speed_tps=40,
    notes="Best cost/quality for deep reasoning; MoE 671B (37B active)",
))

# --- Local models (Ollama — $0) ---

_r(ModelSpec(
    id="ollama/qwen2.5-coder-3b",
    provider="ollama",
    litellm_model="ollama/qwen2.5-coder:3b",
    display_name="Qwen2.5-Coder 3B (local)",
    input_cost=0.0, output_cost=0.0,
    context_window=32_768,
    params_b=3.0,
    is_local=True,
    is_free=True,
    capabilities=frozenset({Capability.CODE, Capability.FAST}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.OPERATIONAL,
    speed_tps=20,
    notes="Jetson Orin 8GB / RPi 5 compatible",
))

_r(ModelSpec(
    id="ollama/qwen2.5-coder-7b",
    provider="ollama",
    litellm_model="ollama/qwen2.5-coder:7b",
    display_name="Qwen2.5-Coder 7B (local)",
    input_cost=0.0, output_cost=0.0,
    context_window=32_768,
    params_b=7.0,
    is_local=True,
    is_free=True,
    capabilities=frozenset({Capability.CODE, Capability.FAST}),
    min_tier=TaskTier.TRIVIAL, max_tier=TaskTier.STANDARD,
    speed_tps=15,
    notes="Good local coding model; 16GB RAM minimum",
))


# ============================================================
# Tier-based model selection helpers
# ============================================================

def models_for_tier(tier: TaskTier, required_caps: frozenset[Capability] | None = None,
                    available_providers: set[str] | None = None) -> list[ModelSpec]:
    """Return models suitable for a tier, sorted by effective cost ascending."""
    results = []
    for m in MODELS.values():
        if tier.rank < m.min_tier.rank:
            continue
        if tier.rank > m.max_tier.rank:
            continue
        if required_caps and not required_caps.issubset(m.capabilities):
            continue
        if available_providers and m.provider not in available_providers:
            continue
        results.append(m)
    return sorted(results, key=lambda m: m.effective_cost)


def cheapest_model(tier: TaskTier, available_providers: set[str] | None = None,
                   required_caps: frozenset[Capability] | None = None) -> ModelSpec | None:
    """Return the single cheapest model that satisfies constraints."""
    candidates = models_for_tier(tier, required_caps, available_providers)
    return candidates[0] if candidates else None


_MODEL_COMPAT_ALIASES: dict[str, str] = {
    # ── Tier aliases (canonical, used everywhere) ──
    "balanced": "google/gemini-3-flash-preview",
    "smart": "google/gemini-3-flash-preview",
    "cheap": "google/gemini-2.5-flash",
    "free": "google/gemini-2.5-flash",
    "local": "ollama/qwen2.5-coder-7b",
    "premium": "anthropic/opus-4.6",
    # ── Legacy Anthropic redirects (avoid expensive models) ──
    "sonnet": "google/gemini-3-flash-preview",
    "sonnet-4": "google/gemini-3-flash-preview",
    "sonnet-4.6": "google/gemini-3-flash-preview",
    "claude-sonnet": "google/gemini-3-flash-preview",
    "openai/claude-sonnet": "google/gemini-3-flash-preview",
    "anthropic/sonnet-4.6": "google/gemini-3-flash-preview",
    "anthropic/claude-sonnet-4-6-20250514": "google/gemini-3-flash-preview",
    "anthropic/claude-sonnet-4-20250514": "google/gemini-3-flash-preview",
    # ── Legacy Anthropic haiku redirects ──
    "haiku": "google/gemini-2.5-flash",
    "haiku-4.5": "anthropic/haiku-4.5",
}


def get_model(model_id: str) -> ModelSpec | None:
    model_id = _MODEL_COMPAT_ALIASES.get(model_id, model_id)
    return MODELS.get(model_id)


def find_model_by_litellm(litellm_model: str) -> ModelSpec | None:
    """Find a model by its litellm_model string."""
    for m in MODELS.values():
        if m.litellm_model == litellm_model:
            return m
    return None


def resolve_model(model_str: str) -> ModelSpec | None:
    """Resolve a model string to a ModelSpec.

    Tries in order: exact id match, litellm_model match, partial id match.
    """
    if not model_str:
        return None
    model_str = _MODEL_COMPAT_ALIASES.get(model_str, model_str)
    # Exact id
    spec = MODELS.get(model_str)
    if spec:
        return spec
    # litellm_model match
    spec = find_model_by_litellm(model_str)
    if spec:
        return spec
    # Settings-level aliases (dynamic, from config)
    try:
        from proxym.config import get_settings
        settings_aliases = get_settings().model_aliases
        alias_target = settings_aliases.get(model_str)
        if alias_target and alias_target != model_str:
            spec = MODELS.get(alias_target)
            if spec:
                return spec
    except Exception:
        pass
    # Partial id match
    model_lower = model_str.lower()
    for m in MODELS.values():
        if model_lower in m.id.lower() or model_lower in m.litellm_model.lower():
            return m
    return None



def filter_models(
    *,
    provider: str | None = None,
    is_free: bool | None = None,
    is_local: bool | None = None,
    min_params_b: float | None = None,
    max_params_b: float | None = None,
    max_input_cost: float | None = None,
    max_output_cost: float | None = None,
    capability: Capability | None = None,
    tier: TaskTier | None = None,
    available_providers: set[str] | None = None,
) -> list[ModelSpec]:
    """Filter models by various criteria, sorted by effective cost ascending."""
    results = []
    for m in MODELS.values():
        if provider and m.provider != provider:
            continue
        if is_free is not None and m.is_free != is_free:
            continue
        if is_local is not None and m.is_local != is_local:
            continue
        if min_params_b is not None and m.params_b > 0 and m.params_b < min_params_b:
            continue
        if max_params_b is not None and m.params_b > 0 and m.params_b > max_params_b:
            continue
        if max_input_cost is not None and m.input_cost > max_input_cost:
            continue
        if max_output_cost is not None and m.output_cost > max_output_cost:
            continue
        if capability and capability not in m.capabilities:
            continue
        if tier and (tier.rank < m.min_tier.rank or tier.rank > m.max_tier.rank):
            continue
        if available_providers and m.provider not in available_providers:
            continue
        results.append(m)
    return sorted(results, key=lambda m: m.effective_cost)


def all_model_ids() -> list[str]:
    """Return all model IDs sorted by effective cost."""
    return [m.id for m in sorted(MODELS.values(), key=lambda m: m.effective_cost)]
