"""Curated model registry — the single source of truth for supported models.

Model IDs were web-verified for **June 2026**. Providers ship new model IDs
frequently, so re-verify against each provider's docs and run ``--smoke-test``
after updating. ``--list-models`` and the README table both render from here.

Each :class:`ModelSpec` maps a user-facing model id to a provider and the exact
``api_id`` sent on the wire. Several providers (DeepSeek, Ollama, xAI, Qwen,
Moonshot) speak the OpenAI Chat Completions protocol, so they all route through
the same :class:`~pentestgpt_legacy.llm.providers.openai_compatible.OpenAICompatibleProvider`
via a per-provider ``base_url``.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class ProviderInfo:
    """Connection metadata for one provider."""

    key: str
    label: str
    # Which connector class handles this provider.
    kind: str  # "openai" | "anthropic" | "gemini"
    # Primary env var holding the API key (None => no key required, e.g. Ollama).
    env: str | None = None
    # Additional accepted env vars for the key (e.g. GOOGLE_API_KEY for Gemini).
    env_alt: tuple[str, ...] = ()
    # Default API base URL (None => the SDK's own default).
    base_url: str | None = None
    requires_key: bool = True


@dataclass(frozen=True)
class ModelSpec:
    """One supported model."""

    id: str  # user-facing id (CLI value, registry key)
    provider: str  # key into PROVIDERS
    context_window: int
    tier: str  # "flagship" | "balanced" | "fast" | "reasoning" | "coding" | "legacy"
    api_id: str = ""  # wire id; defaults to ``id``
    legacy: bool = False
    reasoning: bool = False  # reasoning model (metadata; temperature omitted by default)
    # OpenAI reasoning models cap output with ``max_completion_tokens`` instead.
    max_tokens_param: str = "max_tokens"
    # Some OpenAI models are served only via the Responses API (not chat/completions).
    responses_api: bool = False
    aliases: tuple[str, ...] = ()
    notes: str = ""

    def __post_init__(self) -> None:
        if not self.api_id:
            object.__setattr__(self, "api_id", self.id)


# --------------------------------------------------------------------------- #
# Providers
# --------------------------------------------------------------------------- #

PROVIDERS: dict[str, ProviderInfo] = {
    "openai": ProviderInfo(key="openai", label="OpenAI", kind="openai", env="OPENAI_API_KEY"),
    "anthropic": ProviderInfo(
        key="anthropic", label="Anthropic", kind="anthropic", env="ANTHROPIC_API_KEY"
    ),
    "gemini": ProviderInfo(
        key="gemini",
        label="Google Gemini",
        kind="gemini",
        env="GEMINI_API_KEY",
        env_alt=("GOOGLE_API_KEY",),
    ),
    "deepseek": ProviderInfo(
        key="deepseek",
        label="DeepSeek",
        kind="openai",
        env="DEEPSEEK_API_KEY",
        base_url="https://api.deepseek.com",
    ),
    "xai": ProviderInfo(
        key="xai",
        label="xAI Grok",
        kind="openai",
        env="GROK_API_KEY",
        env_alt=("XAI_API_KEY",),
        base_url="https://api.x.ai/v1",
    ),
    "qwen": ProviderInfo(
        key="qwen",
        label="Alibaba Qwen",
        kind="openai",
        env="QWEN_API_KEY",
        env_alt=("DASHSCOPE_API_KEY",),
        base_url="https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
    ),
    "moonshot": ProviderInfo(
        key="moonshot",
        label="Moonshot Kimi",
        kind="openai",
        env="KIMI_API_KEY",
        env_alt=("MOONSHOT_API_KEY",),
        # China platform default; international keys use https://api.moonshot.ai/v1
        # (override with MOONSHOT_BASE_URL).
        base_url="https://api.moonshot.cn/v1",
    ),
    "ollama": ProviderInfo(
        key="ollama",
        label="Ollama (local)",
        kind="openai",
        env=None,
        base_url="http://localhost:11434/v1",
        requires_key=False,
    ),
}

# Marker provider used for ``ollama:<model>`` dynamic ids.
OLLAMA_PREFIX = "ollama:"


# --------------------------------------------------------------------------- #
# Models (web-verified June 2026)
# --------------------------------------------------------------------------- #

_MODEL_LIST: list[ModelSpec] = [
    # ---- OpenAI -----------------------------------------------------------
    ModelSpec(
        "gpt-5.5",
        "openai",
        1_000_000,
        "flagship",
        reasoning=True,
        max_tokens_param="max_completion_tokens",
        notes="GPT-5.5 flagship",
    ),
    ModelSpec(
        "gpt-5.5-pro",
        "openai",
        1_000_000,
        "flagship",
        reasoning=True,
        max_tokens_param="max_completion_tokens",
        responses_api=True,
        notes="Highest-capability GPT-5.5 (Responses API)",
    ),
    ModelSpec(
        "gpt-5.4-mini",
        "openai",
        400_000,
        "fast",
        reasoning=True,
        max_tokens_param="max_completion_tokens",
        notes="Lower-latency/cost",
    ),
    ModelSpec(
        "gpt-5.4-nano",
        "openai",
        400_000,
        "fast",
        reasoning=True,
        max_tokens_param="max_completion_tokens",
        notes="Smallest GPT-5.4",
    ),
    ModelSpec(
        "gpt-5.2",
        "openai",
        400_000,
        "balanced",
        reasoning=True,
        max_tokens_param="max_completion_tokens",
    ),
    ModelSpec(
        "gpt-5.3-codex",
        "openai",
        400_000,
        "coding",
        reasoning=True,
        max_tokens_param="max_completion_tokens",
        responses_api=True,
        notes="Agentic coding model (Responses API)",
    ),
    ModelSpec("gpt-4o", "openai", 128_000, "legacy", legacy=True),
    ModelSpec("gpt-4o-mini", "openai", 128_000, "legacy", legacy=True),
    ModelSpec(
        "o3",
        "openai",
        200_000,
        "legacy",
        legacy=True,
        reasoning=True,
        max_tokens_param="max_completion_tokens",
    ),
    ModelSpec(
        "o4-mini",
        "openai",
        200_000,
        "legacy",
        legacy=True,
        reasoning=True,
        max_tokens_param="max_completion_tokens",
    ),
    # ---- Anthropic --------------------------------------------------------
    ModelSpec(
        "claude-opus-4-8", "anthropic", 1_000_000, "flagship", notes="Claude Opus 4.8 (May 2026)"
    ),
    ModelSpec("claude-sonnet-4-6", "anthropic", 1_000_000, "balanced", notes="Claude Sonnet 4.6"),
    ModelSpec(
        "claude-haiku-4-5-20251001",
        "anthropic",
        200_000,
        "fast",
        aliases=("claude-haiku-4-5",),
        notes="Claude Haiku 4.5",
    ),
    # ---- Google Gemini ----------------------------------------------------
    ModelSpec(
        "gemini-3.1-pro", "gemini", 1_000_000, "flagship", notes="Most advanced reasoning Gemini"
    ),
    ModelSpec(
        "gemini-3.5-flash", "gemini", 1_000_000, "balanced", notes="GA frontier agentic/coding"
    ),
    ModelSpec("gemini-3-pro", "gemini", 1_000_000, "flagship"),
    ModelSpec("gemini-3.1-flash-lite", "gemini", 1_000_000, "fast"),
    ModelSpec("gemini-2.5-pro", "gemini", 1_000_000, "legacy", legacy=True),
    ModelSpec("gemini-2.5-flash", "gemini", 1_000_000, "legacy", legacy=True),
    # ---- DeepSeek (OpenAI-compatible) -------------------------------------
    ModelSpec(
        "deepseek-v4-flash",
        "deepseek",
        1_000_000,
        "balanced",
        notes="DeepSeek V4 (non-thinking + thinking)",
    ),
    ModelSpec("deepseek-v4-pro", "deepseek", 1_000_000, "flagship"),
    ModelSpec(
        "deepseek-chat",
        "deepseek",
        128_000,
        "legacy",
        legacy=True,
        notes="Retires 2026-07-24; -> deepseek-v4-flash",
    ),
    ModelSpec(
        "deepseek-reasoner",
        "deepseek",
        128_000,
        "legacy",
        legacy=True,
        reasoning=True,
        notes="Retires 2026-07-24; -> deepseek-v4-flash thinking",
    ),
    # ---- xAI Grok (OpenAI-compatible) -------------------------------------
    ModelSpec("grok-4.3", "xai", 1_000_000, "flagship", notes="xAI flagship"),
    # ---- Alibaba Qwen (OpenAI-compatible) ---------------------------------
    ModelSpec("qwen3.7-max", "qwen", 262_144, "flagship", notes="Qwen 3.7 Max"),
    ModelSpec("qwen3.5-flash", "qwen", 262_144, "fast"),
    ModelSpec("qwen3-max", "qwen", 262_144, "legacy", legacy=True),
    # ---- Moonshot Kimi (OpenAI-compatible) --------------------------------
    ModelSpec("kimi-k2.6", "moonshot", 256_000, "flagship", notes="Kimi K2.6"),
]

# Registry keyed by canonical id (and by alias) for O(1) lookup.
MODELS: dict[str, ModelSpec] = {}
_ALIASES: dict[str, str] = {}
for _spec in _MODEL_LIST:
    MODELS[_spec.id] = _spec
    for _alias in _spec.aliases:
        _ALIASES[_alias] = _spec.id


def all_model_ids() -> list[str]:
    """All canonical, user-selectable model ids (registry order)."""
    return [spec.id for spec in _MODEL_LIST]


def resolve(name: str) -> ModelSpec | None:
    """Resolve a model id or alias to a :class:`ModelSpec`.

    Supports the dynamic ``ollama:<model>`` form for arbitrary local models.
    Returns ``None`` if unknown.
    """
    if name.startswith(OLLAMA_PREFIX):
        local = name[len(OLLAMA_PREFIX) :].strip()
        if not local:
            return None
        return ModelSpec(
            id=name,
            provider="ollama",
            api_id=local,
            context_window=128_000,
            tier="local",
            notes="User-configured local Ollama model",
        )
    if name in MODELS:
        return MODELS[name]
    if name in _ALIASES:
        return MODELS[_ALIASES[name]]
    return None


def models_by_provider() -> dict[str, list[ModelSpec]]:
    """Group registry models by provider key (registry order preserved)."""
    grouped: dict[str, list[ModelSpec]] = {}
    for spec in _MODEL_LIST:
        grouped.setdefault(spec.provider, []).append(spec)
    return grouped


# Sensible defaults for the three sessions, in preference order. The first one
# whose provider key is configured is used when the user does not pass --*-model.
DEFAULT_REASONING_PREFERENCE: tuple[str, ...] = (
    "claude-opus-4-8",
    "gpt-5.5",
    "gemini-3.1-pro",
    "deepseek-v4-pro",
    "grok-4.3",
)
DEFAULT_PARSING_PREFERENCE: tuple[str, ...] = (
    "claude-haiku-4-5-20251001",
    "gpt-5.4-mini",
    "gemini-3.5-flash",
    "deepseek-v4-flash",
)

# kept for callers that want the raw ordered list
ALL_SPECS: tuple[ModelSpec, ...] = tuple(_MODEL_LIST)
