"""Proxym-specific types for model specifications and task tiers."""

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,
        }
