"""Pre-dispatch model-identity normalization for CLI workers.

The catalog (models.py) carries okstra's model identity. Some worker CLIs
(antigravity `agy`, codex) accept their own spelling of that identity, and it
drifts across CLI versions / accounts. This module maps a catalog execution
value + role to the CLI's exact accepted identifier, or hard-fails with the
CLI's real list. It never substitutes a different model — spelling/effort only.
"""
from __future__ import annotations

import subprocess
from functools import lru_cache
from pathlib import Path

# Per-role reasoning-effort policy for CLIs that bake effort into the model
# name (agy). Deterministic — no per-run improvisation.
# `critic` is intentionally absent: an opt-in critic reuses its worker's
# execution value (render.py maps critic_choice -> *_WORKER_MODEL_EXECUTION_VALUE),
# so it inherits the worker's normalized identity + effort. role_effort("critic")
# returns _DEFAULT_EFFORT, matching the worker it mirrors.
ROLE_EFFORT = {"worker": "High", "executor": "High"}
_DEFAULT_EFFORT = "High"


class ModelUnavailableError(RuntimeError):
    """Requested model identity is not offered by the worker CLI on this host."""


def role_effort(role: str) -> str:
    return ROLE_EFFORT.get(role, _DEFAULT_EFFORT)


# agy lists a tier per model, but the listed slug is not always the model it
# serves: on agy 1.1.10 a `gemini-3.1-pro-high` session identifies itself as
# "Gemini 3.6 Flash", while `gemini-3.1-pro-low` correctly identifies as
# "Gemini 3.1 Pro". The substitution is not visible anywhere in the dispatch —
# `agy models` lists the slug and the run exits 0 — so it surfaces only as a
# verifier that never refutes anything: measured over one 63-item plan-body
# prompt, the high slug returned all-AGREE in 64s while the low slug spent 328s
# and raised three DISAGREEs, one of them the defect claude and codex both
# caught that round. Only the `high` tier is affected, so a role asking for any
# other effort still resolves normally (and a tier agy does not offer still
# hard-fails below).
_UNTRUSTED_HIGH_TIER_EXECUTIONS = frozenset({"gemini-3.1-pro"})


def _dispatch_effort(execution: str, role: str) -> str:
    effort = role_effort(role).lower()
    if effort == "high" and execution in _UNTRUSTED_HIGH_TIER_EXECUTIONS:
        return "low"
    return effort


def dispatch_tier_suffixes() -> frozenset[str]:
    """Every tier suffix dispatch can append to an agy execution value.

    The provider serves back the suffixed identity it was given, so reading
    that identity means undoing exactly this set — not guessing at whatever
    trails the last hyphen. `low` is here because `_dispatch_effort` demotes
    the untrusted high tier to it.
    """
    efforts = {value.lower() for value in (*ROLE_EFFORT.values(), _DEFAULT_EFFORT)}
    return frozenset(efforts | {"low"})


@lru_cache(maxsize=1)
def agy_models(agy_bin: str = "agy") -> tuple[str, ...]:
    """Live `agy models` ids, or () when agy is unavailable (non-blocking).

    agy prints one model per line as `<id>\\t<display label>`; only the id is
    comparable to a catalog execution value, so the label column is dropped.
    A line with no tab is already bare and passes through unchanged.
    """
    try:
        proc = subprocess.run(
            [agy_bin, "models"],
            capture_output=True, text=True, timeout=20,
        )
    except (FileNotFoundError, OSError, subprocess.SubprocessError):
        return ()
    if proc.returncode != 0:
        return ()
    ids = (line.split("\t", 1)[0].strip() for line in proc.stdout.splitlines())
    return tuple(model_id for model_id in ids if model_id)


def _codex_config_path() -> Path:
    return Path.home() / ".codex" / "config.toml"


def codex_availability(config_path: Path | None = None) -> tuple[str, ...]:
    """Keys under [tui.model_availability_nux] in codex config, or ()."""
    path = config_path or _codex_config_path()
    try:
        lines = path.read_text(encoding="utf-8").splitlines()
    except OSError:
        return ()
    out: list[str] = []
    in_section = False
    for raw in lines:
        line = raw.strip()
        if line.startswith("[") and line.endswith("]"):
            in_section = line == "[tui.model_availability_nux]"
            continue
        if in_section and "=" in line:
            key = line.split("=", 1)[0].strip().strip('"')
            if key:
                out.append(key)
    return tuple(out)


def normalize_execution_for_dispatch(
    *, provider: str, execution: str, role: str,
) -> str:
    """Return the CLI-accepted spelling of `execution`. When discovery is
    unavailable, fall back to the best-known dispatchable spelling rather than
    blocking. Raise ModelUnavailableError when the identity is definitely not
    offered (never substitute a different model)."""
    if provider == "codex":
        return normalize_codex_execution(
            execution=execution,
            available=codex_availability(),
        )
    if provider != "antigravity":
        return execution  # claude runs via the Agent tool
    return normalize_antigravity_execution(
        execution=execution,
        role=role,
        available=agy_models(),
    )


def normalize_codex_execution(
    *, execution: str, available: tuple[str, ...]
) -> str:
    """Normalize one Codex value against a previously captured config list."""
    if not available or execution in available:
        return execution
    raise ModelUnavailableError(
        f"codex model {execution!r} is not available for this account. "
        f"Available: {', '.join(available)} (no auto-rename — update the catalog)"
    )


def normalize_antigravity_execution(
    *, execution: str, role: str, available: tuple[str, ...]
) -> str:
    """Normalize one agy value against a previously captured model list."""
    # agy bakes the tier into the model identity as a slug suffix
    # (`gemini-3.1-pro-high`) and rejects the bare slug with "requires --effort"
    # — a flag the wrapper never sends. So the suffixed form is the only
    # dispatchable spelling, including on the un-verified fallback below.
    tiered = f"{execution}-{_dispatch_effort(execution, role)}"
    if not available:
        return tiered  # discovery unavailable → dispatch anyway, wrapper still guards
    if execution in available:
        return execution  # identity already carries its tier (agy's Claude models)
    if tiered in available:
        return tiered
    raise ModelUnavailableError(
        f"antigravity model {tiered!r} (from catalog {execution!r}, role {role!r}) "
        f"is not offered by this agy install. Available: {', '.join(available)}"
    )