"""Model alias registry — maps short names to (profile, model_id) tuples.

Built-in defaults cover common Claude, OpenAI, and other models.
Users can add or override aliases via ~/.config/squadron/models.toml.
"""

from __future__ import annotations

import logging
import tomllib
from pathlib import Path
from typing import Any, TypedDict, cast

_logger = logging.getLogger(__name__)


class ModelPricing(TypedDict, total=False):
    """Per-token pricing for a model, all values in USD per 1M tokens."""

    input: float
    output: float
    cache_read: float
    cache_write: float


class _ModelAliasRequired(TypedDict):
    """Required fields for a model alias."""

    profile: str
    model: str


class ModelAlias(_ModelAliasRequired, total=False):
    """A model alias mapping a short name to a profile and full model ID.

    ``profile`` and ``model`` are always required.  The remaining fields
    are optional metadata added by slice 121.
    """

    private: bool
    cost_tier: str
    notes: str
    pricing: ModelPricing


BUILT_IN_ALIASES: dict[str, ModelAlias] = {
    # Claude family — Max subscription, no per-token pricing
    "opus": {
        "profile": "sdk",
        "model": "claude-opus-4-6",
        "private": True,
        "cost_tier": "subscription",
        "notes": "Max sub",
    },
    "sonnet": {
        "profile": "sdk",
        "model": "claude-sonnet-4-6",
        "private": True,
        "cost_tier": "subscription",
        "notes": "Max sub",
    },
    "haiku": {
        "profile": "sdk",
        "model": "claude-haiku-4-5-20251001",
        "private": True,
        "cost_tier": "subscription",
        "notes": "Max sub",
    },
    # OpenAI GPT-5.4 family
    "gpt54": {
        "profile": "openai",
        "model": "gpt-5.4",
        "private": True,
        "cost_tier": "expensive",
        "notes": "Frontier",
        "pricing": {
            "input": 2.50,
            "output": 15.00,
            "cache_read": 0.25,
            "cache_write": 2.50,
        },
    },
    "gpt54-mini": {
        "profile": "openai",
        "model": "gpt-5.4-mini",
        "private": True,
        "cost_tier": "moderate",
        "notes": "Balanced",
        "pricing": {
            "input": 0.75,
            "output": 4.50,
            "cache_read": 0.075,
            "cache_write": 0.75,
        },
    },
    "gpt54-nano": {
        "profile": "openai",
        "model": "gpt-5.4-nano",
        "private": True,
        "cost_tier": "cheap",
        "notes": "Lightweight",
        "pricing": {
            "input": 0.20,
            "output": 1.25,
            "cache_read": 0.02,
            "cache_write": 1.25,
        },
    },
    # OpenAI specialized
    "codex": {
        "profile": "openai",
        "model": "gpt-5.3-codex",
        "private": True,
        "cost_tier": "expensive",
        "notes": "Code-specialized",
        "pricing": {
            "input": 1.75,
            "output": 14.00,
            "cache_read": 0.0175,
            "cache_write": 1.75,
        },
    },
    "codex-agent": {
        "profile": "openai-oauth",
        "model": "gpt-5.3-codex",
        "private": True,
        "cost_tier": "subscription",
        "notes": "Agentic: sandbox, subscription auth",
    },
    "codex-spark": {
        "profile": "openai-oauth",
        "model": "gpt-5.3-codex-spark",
        "private": True,
        "cost_tier": "subscription",
        "notes": "Near-instant, Pro only",
    },
    # Google Gemini — free tier
    "gemini": {
        "profile": "gemini",
        "model": "gemini-3.1-pro-preview-customtools",
        "private": True,
        "cost_tier": "expensive",
        "notes": "Preview",
        "pricing": {
            "input": 2.00,
            "output": 12.00,
            "cache_read": 0.20,
            "cache_write": 2.00,
        },
    },
    "flash3": {
        "profile": "gemini",
        "model": "gemini-3-flash-preview",
        "private": True,
        "cost_tier": "cheap",
        "notes": "Fast, cheap",
        "pricing": {
            "input": 0.50,
            "output": 3.00,
            "cache_read": 0.05,
            "cache_write": 0.50,
        },
    },
    "flash3-lite": {
        "profile": "gemini",
        "model": "gemini-3-flash-lite-preview",
        "private": True,
        "cost_tier": "cheap",
        "notes": "very cheap",
        "pricing": {
            "input": 0.25,
            "output": 1.50,
            "cache_read": 0.025,
            "cache_write": 0.25,
        },
    },
    # OpenRouter — cost-effective, no retention/training
    "kimi25": {
        "profile": "openrouter",
        "model": "moonshotai/kimi-k2.5",
        "private": True,
        "cost_tier": "cheap",
        "notes": "1M context",
        "pricing": {
            "input": 0.45,
            "output": 2.83,
            "cache_read": 0.189,
            "cache_write": 2.83,
        },
    },
    "minimax": {
        "profile": "openrouter",
        "model": "minimax/minimax-m2.7",
        "private": True,
        "cost_tier": "cheap",
        "pricing": {
            "input": 0.30,
            "output": 1.25,
            "cache_read": 0.06,
            "cache_write": 0.30,
        },
    },
    "glm5": {
        "profile": "openrouter",
        "model": "z-ai/glm-5",
        "private": True,
        "cost_tier": "cheap",
        "pricing": {
            "input": 0.72,
            "output": 2.30,
            "cache_read": 0.72,
            "cache_write": 2.30,
        },
    },
    "glm5-turbo": {
        "profile": "openrouter",
        "model": "z-ai/glm-5-turbo",
        "private": True,
        "cost_tier": "moderate",
        "notes": "on sale",
        "pricing": {
            "input": 0.96,
            "output": 3.20,
            "cache_read": 0.192,
            "cache_write": 3.20,
        },
    },
    "mimo-omni": {
        "profile": "openrouter",
        "model": "xiaomi/mimo-v2-omni",
        "private": True,
        "cost_tier": "cheap",
        "pricing": {
            "input": 0.40,
            "output": 2.00,
            "cache_read": 0.08,
            "cache_write": 2.00,
        },
    },
    "grok-fast": {
        "profile": "openrouter",
        "model": "x-ai/grok-4.1-fast",
        "private": True,
        "cost_tier": "cheap",
        "notes": "2M context",
        "pricing": {
            "input": 0.20,
            "output": 0.50,
            "cache_read": 0.05,
            "cache_write": 0.20,
        },
    },
}


def models_toml_path() -> Path:
    """Return the path to the user's models.toml configuration file."""
    return Path.home() / ".config" / "squadron" / "models.toml"


def _extract_metadata(
    alias: ModelAlias,
    table: dict[str, Any],
    name: str,
    path: Path,
) -> None:
    """Extract optional metadata and pricing fields from a TOML alias table."""
    private_val = table.get("private")
    if isinstance(private_val, bool):
        alias["private"] = private_val

    cost_tier_val = table.get("cost_tier")
    if isinstance(cost_tier_val, str):
        alias["cost_tier"] = cost_tier_val

    notes_val = table.get("notes")
    if isinstance(notes_val, str):
        alias["notes"] = notes_val

    pricing_val = table.get("pricing")
    if isinstance(pricing_val, dict):
        pricing = _extract_pricing(cast(dict[str, Any], pricing_val), name, path)
        if pricing:
            alias["pricing"] = pricing


def _extract_pricing(
    raw: dict[str, Any],
    name: str,
    path: Path,
) -> ModelPricing | None:
    """Build a ModelPricing dict from raw TOML pricing values."""
    pricing = ModelPricing()
    for field in ("input", "output", "cache_read", "cache_write"):
        val = raw.get(field)
        if val is None:
            continue
        if isinstance(val, (int, float)):
            pricing[field] = float(val)  # type: ignore[literal-required]
        else:
            _logger.warning(
                "Skipping pricing field '%s' for alias '%s' in %s — expected a number",
                field,
                name,
                path,
            )
    return pricing if pricing else None


def load_user_aliases() -> dict[str, ModelAlias]:
    """Load user-defined aliases from models.toml.

    Returns an empty dict if the file is missing or has no [aliases] section.
    """
    path = models_toml_path()
    if not path.is_file():
        return {}

    try:
        data = tomllib.loads(path.read_text())
    except tomllib.TOMLDecodeError as exc:
        raise ValueError(
            f"Invalid TOML in {path}: {exc}. "
            "Fix the file or remove it to use built-in defaults."
        ) from exc
    except OSError as exc:
        _logger.warning("Failed to read %s: %s", path, exc)
        return {}

    aliases_section = cast(dict[str, Any], data.get("aliases", {}))
    result: dict[str, ModelAlias] = {}
    for name, entry in aliases_section.items():
        if not isinstance(entry, dict):
            _logger.warning(
                "Skipping invalid alias '%s' in %s — expected a table",
                name,
                path,
            )
            continue
        table = cast(dict[str, Any], entry)
        profile_val = table.get("profile")
        model_val = table.get("model")
        if isinstance(profile_val, str) and isinstance(model_val, str):
            alias: ModelAlias = ModelAlias(profile=profile_val, model=model_val)
            _extract_metadata(alias, table, name, path)
            result[name] = alias
        else:
            _logger.warning(
                "Skipping alias '%s' in %s — "
                "expected string values for 'profile' and 'model'",
                name,
                path,
            )
    return result


def get_all_aliases() -> dict[str, ModelAlias]:
    """Return merged aliases: built-in defaults + user overrides.

    User entries override built-in aliases by name.
    """
    merged = dict(BUILT_IN_ALIASES)
    merged.update(load_user_aliases())
    return merged


def resolve_model_alias(name: str) -> tuple[str, str | None]:
    """Resolve a model alias to (full_model_id, profile_or_none).

    If the name matches a known alias (built-in or user), returns the
    alias's (model, profile). If not, returns (name, None) — treating
    the input as a literal model ID with no profile inference.
    """
    aliases = get_all_aliases()
    alias = aliases.get(name)
    if alias is not None:
        return alias["model"], alias["profile"]
    return name, None


def estimate_cost(
    alias_name: str,
    input_tokens: int,
    output_tokens: int,
    cached_tokens: int = 0,
) -> float | None:
    """Estimate USD cost for a model given token counts.

    Returns None if the alias has no pricing data or insufficient
    pricing fields for the calculation.
    """
    aliases = get_all_aliases()
    alias = aliases.get(alias_name)
    if alias is None:
        return None

    pricing = alias.get("pricing")
    if pricing is None:
        return None

    input_price = pricing.get("input")
    output_price = pricing.get("output")
    if input_price is None or output_price is None:
        return None

    total = (
        input_tokens / 1_000_000 * input_price
        + output_tokens / 1_000_000 * output_price
    )

    if cached_tokens > 0:
        cache_read_price = pricing.get("cache_read")
        if cache_read_price is not None:
            total += cached_tokens / 1_000_000 * cache_read_price

    return total
