"""Config + launch banner. Run: uv run pytest tests/test_config_registry.py -v -m "not gpu"

The banner contract: NO GPU run starts without printing every resolved value
it depends on. A run that cannot answer its question does not start.
"""

from __future__ import annotations

import math
import os

import pytest

from align import registry
from align.config import default_full_ft_config, load_config


def _complete_resolved() -> dict:
    """A synthetic fully-resolved value set covering every mandatory key.
    Built by iterating registry.MANDATORY_BANNER_KEYS so the test tracks the
    list as data instead of hardcoding a parallel copy."""
    samples = {
        "attention_backend": "sdpa_flash",
        "loss_type": "dpo_sigmoid",
        "beta": 0.1,
        "lr": 5e-7,
        "scheduler_horizon_optimizer_steps": 1500,
        "dataset_hashes": {"train": "abc123", "eval_gsm8k": "def456"},
        "ceiling_results": {"extraction_roundtrip": "PASS"},
        "rss_budget": {"total_gb": 108.0, "optimizer_state_gb": 54.0},
        "pinned_staging_gb": 1.8,
        "probe_step_time_s": float("nan"),  # explicit NaN, never absent
        "mde": {"n": 1319, "mde_pp": 2.8},
    }
    missing = [k for k in registry.MANDATORY_BANNER_KEYS if k not in samples]
    assert not missing, f"test sample set out of date with registry: {missing}"
    return samples


def test_banner_contains_every_mandatory_key():
    """Rendered banner for a complete value set contains ALL mandatory keys,
    each with its resolved value on the same line."""
    resolved = _complete_resolved()
    banner = registry.render_banner(resolved)
    for key in registry.MANDATORY_BANNER_KEYS:
        assert key in banner, f"banner missing key {key}"
    assert "sdpa_flash" in banner
    assert "nan" in banner.lower()  # explicit NaN is printed, not dropped


def test_banner_missing_key_raises():
    """A value set missing ANY mandatory key -- or carrying None for one --
    RAISES rather than printing a partial banner. Silent omission is the
    root shape of every predecessor silent failure: a missing value replaced
    by a plausible one."""
    for key in registry.MANDATORY_BANNER_KEYS:
        resolved = _complete_resolved()
        del resolved[key]
        with pytest.raises(RuntimeError, match=key):
            registry.render_banner(resolved)

    resolved = _complete_resolved()
    resolved["beta"] = None
    with pytest.raises(RuntimeError, match="beta"):
        registry.render_banner(resolved)


def test_cli_dotted_overrides(tmp_workdir):
    """`scripts/x.py configs/x.toml --optimizer.lr=1e-5`: dotted overrides
    reach the dataclass field with correct type; unknown keys raise under
    struct mode, never silently create new fields."""
    toml_path = os.path.join(tmp_workdir, "t.toml")
    with open(toml_path, "w") as f:
        f.write("[optimizer]\nlr = 1.0e-6\n")

    cfg = load_config(toml_path, ["optimizer.lr=1e-5"])
    assert isinstance(cfg.optimizer.lr, float)
    assert cfg.optimizer.lr == pytest.approx(1e-5)

    with pytest.raises(Exception, match="optimzer|struct|key"):
        load_config(toml_path, ["optimzer.lr=1e-5"])  # typo'd section

    with pytest.raises(Exception, match="not_a_field|struct|key"):
        load_config(toml_path, ["optimizer.not_a_field=3"])


def test_undefined_statistic_is_nan_never_zero(tmp_workdir):
    """Tracker contract: an undefined statistic is logged as NaN, NEVER 0 and
    never omitted. Incident: an EV clamp turned an undefined explained
    variance into -1.8e9 while the critic looked healthy; three analyses were
    later retracted for trusting plausible stand-in numbers."""
    from align.utils.tracking import Tracker

    t = Tracker(run_dir=tmp_workdir, stage="test", mandatory_keys=("loss", "margin"))
    t.log(step=0, metrics={"loss": 1.0, "margin": float("nan")})

    with pytest.raises(RuntimeError, match="margin"):
        t.log(step=1, metrics={"loss": 1.0})  # omission raises

    with pytest.raises(RuntimeError, match="None"):
        t.log(step=2, metrics={"loss": 1.0, "margin": None})  # None raises

    import json

    lines = [json.loads(ln) for ln in open(os.path.join(tmp_workdir, "history.jsonl"))]
    assert len(lines) == 1
    assert math.isnan(lines[0]["margin"])