"""Model profiles — toggle between local models with one env var (MODEL_PROFILE).

Served by oMLX (Apple MLX inference server) on this machine. `served_model_id` is the id
oMLX reports at GET /v1/models; `match_keywords` let the client resolve the profile even if
that id string changes. Both models below are already downloaded in ~/.omlx/models.
"""

from __future__ import annotations

from pydantic import BaseModel, Field


class ModelProfile(BaseModel):
    key: str
    display_name: str
    served_model_id: str  # id as oMLX reports it; client also falls back to keyword match
    match_keywords: list[str] = Field(default_factory=list)
    context_length: int = 16384
    temperature: float = 0.7
    downloaded: bool = False
    notes: str = ""


GPT_OSS_20B = ModelProfile(
    key="gpt_oss",
    display_name="gpt-oss-20b MXFP4 (~11GB)",
    served_model_id="gpt-oss-20b-MXFP4-Q8",
    match_keywords=["gpt-oss"],
    context_length=131072,  # full window; oMLX SSD KV cache absorbs the growth
    temperature=0.7,
    downloaded=True,
    notes="OpenAI open-weight, trained for tool-use/browsing + reasoning. Lighter → more "
    "RAM headroom for long overnight runs. Good default.",
)

QWEN35_27B = ModelProfile(
    key="qwen35_27b",
    display_name="Qwen3.5-27B Claude-Opus-distilled (~14GB)",
    served_model_id="Qwen3.5-27B-Claude-4.6-Opus-Distilled-MLX-4bit",
    match_keywords=["qwen3.5", "opus-distilled", "qwen3.5-27b"],
    context_length=262144,  # Qwen3.5 supports even more; SSD KV cache absorbs the growth
    temperature=0.7,
    downloaded=True,
    notes="Qwen (strong tool-calling) distilled from Claude Opus (strong reasoning). Bigger "
    "→ tighter RAM, but oMLX's SSD KV cache absorbs long-context growth.",
)

MODEL_PROFILES: dict[str, ModelProfile] = {p.key: p for p in [GPT_OSS_20B, QWEN35_27B]}
DEFAULT_PROFILE = "gpt_oss"


def get_profile(key: str) -> ModelProfile:
    if key not in MODEL_PROFILES:
        raise ValueError(f"Unknown MODEL_PROFILE {key!r}. Options: {', '.join(MODEL_PROFILES)}")
    return MODEL_PROFILES[key]


def resolve_profile(value: str) -> ModelProfile:
    """Tolerant lookup: accepts a profile key ('gpt_oss'), a served model id
    ('gpt-oss-20b-MXFP4-Q8'), or a keyword. Falls back to the default profile."""
    if value in MODEL_PROFILES:
        return MODEL_PROFILES[value]
    low = (value or "").lower()
    for p in MODEL_PROFILES.values():
        if low == p.served_model_id.lower() or any(kw in low for kw in p.match_keywords):
            return p
    return MODEL_PROFILES[DEFAULT_PROFILE]