"""Pods model configurations — known vLLM model catalog."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from pi_agents.pods.types import GPU


@dataclass
class ModelConfig:
    """A vLLM model configuration."""

    gpu_count: int = 1
    gpu_types: list[str] = field(default_factory=list)
    args: list[str] = field(default_factory=list)
    env: dict[str, str] | None = None
    notes: str | None = None


@dataclass
class ModelInfo:
    """Info about a known model."""

    name: str = ""
    configs: list[ModelConfig] = field(default_factory=list)
    notes: str | None = None


# Static catalog of known vLLM models
_KNOWN_MODELS: dict[str, ModelInfo] = {
    "Qwen/Qwen2.5-Coder-32B-Instruct": ModelInfo(
        name="Qwen2.5-Coder-32B-Instruct",
        configs=[
            ModelConfig(gpu_count=1, args=["--max-model-len", "32768", "--gpu-memory-utilization", "0.95"]),
            ModelConfig(gpu_count=2, args=["--max-model-len", "65536", "--tensor-parallel-size", "2"]),
        ],
    ),
    "Qwen/Qwen3-Coder-30B-A3B-Instruct": ModelInfo(
        name="Qwen3-Coder-30B-A3B-Instruct",
        configs=[
            ModelConfig(gpu_count=1, args=["--max-model-len", "32768"]),
            ModelConfig(gpu_count=2, args=["--max-model-len", "65536", "--tensor-parallel-size", "2"]),
        ],
    ),
    "openai/gpt-oss-20b": ModelInfo(
        name="GPT-OSS 20B",
        configs=[
            ModelConfig(gpu_count=1, gpu_types=["H100", "H200"], args=["--max-model-len", "32768"]),
            ModelConfig(gpu_count=1, gpu_types=["B200"], args=["--max-model-len", "65536"]),
        ],
    ),
    "openai/gpt-oss-120b": ModelInfo(
        name="GPT-OSS 120B",
        configs=[
            ModelConfig(gpu_count=1, args=["--max-model-len", "8192"]),
            ModelConfig(gpu_count=4, args=["--max-model-len", "32768", "--tensor-parallel-size", "4"]),
            ModelConfig(gpu_count=8, args=["--max-model-len", "65536", "--tensor-parallel-size", "8"]),
        ],
    ),
}


def get_known_models() -> list[str]:
    """Get list of known model IDs."""
    return list(_KNOWN_MODELS.keys())


def is_known_model(model_id: str) -> bool:
    """Check if a model ID is in the known catalog."""
    return model_id in _KNOWN_MODELS


def get_model_name(model_id: str) -> str:
    """Get display name for a model. Returns model_id if not known."""
    info = _KNOWN_MODELS.get(model_id)
    return info.name if info else model_id


def get_model_config(
    model_id: str,
    gpus: list[GPU],
    requested_gpu_count: int,
) -> dict[str, Any] | None:
    """Get vLLM config for a model.

    Matches by GPU count first, then GPU type.
    Returns {"args": [...], "env": {...}, "notes": "..."} or None.
    """
    info = _KNOWN_MODELS.get(model_id)
    if not info:
        return None

    gpu_type = ""
    if gpus:
        gpu_type = gpus[0].name.replace("NVIDIA", "").strip().split(" ")[0] if gpus[0].name else ""

    best: ModelConfig | None = None

    # First pass: match GPU count + type
    for config in info.configs:
        if config.gpu_count != requested_gpu_count:
            continue
        if config.gpu_types:
            if any(gpu_type in t or t in gpu_type for t in config.gpu_types):
                best = config
                break
        else:
            best = config
            break

    # Second pass: match GPU count only
    if best is None:
        for config in info.configs:
            if config.gpu_count == requested_gpu_count:
                best = config
                break

    if best is None:
        return None

    result: dict[str, Any] = {"args": list(best.args)}
    if best.env:
        result["env"] = dict(best.env)
    if best.notes or info.notes:
        result["notes"] = best.notes or info.notes
    return result
