"""scripts/p2/contrasts.py: declarative K0/K1/K2/K4/K6 driver. Pure CPU, fake runners, nothing heavy, no GPU, no Lean.

Run: CUDA_VISIBLE_DEVICES= .venv/bin/python -m pytest tests/test_contrasts.py -q
"""

from __future__ import annotations

import json
import shlex
from pathlib import Path

import pytest

from test_p2_support import gen_record, problem_record, write_jsonl, load_script

K = load_script("p2_contrasts", "scripts/p2/contrasts.py")
EI = K.EI
MEAN = 5000.0
BUDGET = 8 * 5000 * 20 * 1  # 20 problems x n=8 x 5000 tokens, 1 round


@pytest.fixture
def world(tmp_workdir):
    d = Path(tmp_workdir)
    for name in ("goedel", "kimina"):
        (d / name).mkdir()
        (d / name / "config.json").write_text("{}")
    write_jsonl(d / "pool.jsonl", [problem_record(f"p{i:02d}") for i in range(10)])
    write_jsonl(d / "mutated.jsonl", [problem_record(f"m{i:02d}", sig=f"(x : ℕ) (h : x = {i + 50}) : x = {i + 50}") for i in range(5)])
    write_jsonl(d / "clf.jsonl", [problem_record(f"c{i:02d}", sig=f"(x : ℕ) (h : x = {i + 70}) : x = {i + 70}") for i in range(5)])
    write_jsonl(d / "rnd.jsonl", [problem_record(f"r{i:02d}", sig=f"(x : ℕ) (h : x = {i + 90}) : x = {i + 90}") for i in range(5)])
    write_jsonl(d / "valid.jsonl", [problem_record(f"mf{i}", "minif2f", "valid", f"(y : ℕ) (h : y = {i + 100}) : y = {i + 100}") for i in range(4)])
    write_jsonl(d / "test.jsonl", [problem_record("mt0", "minif2f", "test", "(y : ℕ) (h : y = 777) : y = 777")])
    write_jsonl(d / "mobench.jsonl", [problem_record("mo0", "mobench", "test", "(z : ℕ) (h : z = 7) : z = 7")])
    rl = d / "rl.toml"
    rl.write_text(Path(K.REPO / "configs" / "p4_cispo_lean.toml").read_text())
    return d


def base_argv(d, *contrasts, **kw):
    a = ["--contrast", *contrasts, "--out-dir", str(d / "out"), "--goedel-model", str(d / "goedel"), "--kimina-model", str(d / "kimina"),
         "--pool", str(d / "pool.jsonl"), "--eval-problems", str(d / "valid.jsonl"), "--protected-problems", str(d / "test.jsonl"), str(d / "mobench.jsonl"),
         "--mean-completion-tokens", str(MEAN), "--token-budget", str(BUDGET), "--rl-config", str(d / "rl.toml"), "--n", "8", "--rounds", "1"]
    for k, v in kw.items():
        a += [f"--{k.replace('_', '-')}"] + ([] if v is True else (list(map(str, v)) if isinstance(v, (list, tuple)) else [str(v)]))
    return a


def plan_of(d, *contrasts, **kw):
    args = K.parse_args(base_argv(d, *contrasts, **kw))
    return K.build_plan(args), args


class FakeRunner:
    def __init__(self, fail=None):
        self.calls, self.fail = [], fail

    def run(self, argv, env, log):
        self.calls.append(list(argv))
        if self.fail and self.fail in " ".join(argv):
            return 7
        return 0


# =============================================================================== guards


def test_a_test_split_scored_file_is_refused_and_there_is_no_escape_hatch(world, capsys):
    argv = base_argv(world, "K1")
    argv[argv.index("--eval-problems") + 1] = str(world / "test.jsonl")
    assert K.main(argv + ["--dry-run"]) == K.EXIT_REFUSED
    assert "TEST split" in capsys.readouterr().err
    with pytest.raises(SystemExit):
        K.parse_args(argv + ["--allow-test-selection"])


def test_missing_protected_problems_and_unreadable_scored_file_are_refused(world, capsys):
    argv = base_argv(world, "K1")
    i = argv.index("--protected-problems")
    assert K.main(argv[:i] + argv[i + 3:] + ["--dry-run"]) == K.EXIT_REFUSED  # protected list emptied
    bad = base_argv(world, "K1")
    bad[bad.index("--eval-problems") + 1] = str(world / "nope.jsonl")
    assert K.main(bad + ["--dry-run"]) == K.EXIT_REFUSED
    assert "cannot read the scored benchmark" in capsys.readouterr().err


def test_guard_passes_a_dev_split_and_every_ei_command_carries_protected_problems(world):
    plan, args = plan_of(world, "K1", "K2", "K6", "K0", mutated_pool=world / "mutated.jsonl", classifier_pool=world / "clf.jsonl", random_pool=world / "rnd.jsonl")
    assert plan.guard["splits"] == ["valid"] and plan.guard["allow_test_selection"] is False
    ei_steps = [s for s in plan.steps if s.kind == "ei"]
    assert ei_steps
    for s in ei_steps:
        i = s.argv.index("--protected-problems")
        assert str(world / "test.jsonl") in s.argv[i:] and str(world / "mobench.jsonl") in s.argv[i:]


@pytest.mark.parametrize("kw,msg", [
    (dict(seeds=[0, 1, 2]), "exactly 2 seeds"),   # K0
])
def test_k0_needs_exactly_two_seeds(world, kw, msg):
    with pytest.raises(K.ContrastError, match=msg):
        plan_of(world, "K0", **kw)


def test_missing_mean_tokens_budget_and_pools_are_refused(world):
    a = base_argv(world, "K1")
    i = a.index("--mean-completion-tokens")
    with pytest.raises(K.ContrastError, match="mean-completion-tokens"):
        K.build_plan(K.parse_args(a[:i] + a[i + 2:]))
    with pytest.raises(K.ContrastError, match="--mutated-pool"):
        plan_of(world, "K2")
    with pytest.raises(K.ContrastError, match="K4 needs EI arms|requires"):
        plan_of(world, "K4")


# =============================================================================== budgets


def test_budget_arithmetic():
    assert K.planned_tokens_ei(2, 10, 8, 100) == 16000
    assert K.derive_ei_problems(16000, 2, 8, 100) == 10
    assert K.derive_ei_problems(16999, 2, 8, 100) == 10
    with pytest.raises(K.ContrastError):
        K.derive_ei_problems(10, 1, 8, 100)
    assert K.derive_rl_steps(1_000_000, 16, 16, 1000) == 4  # 3.9 -> nearest
    assert K.derive_rl_steps(1, 16, 16, 1000) == 1
    assert K.token_match(100, 104, 0.05) and not K.token_match(100, 110, 0.05)


def test_k1_arms_are_matched_on_planned_tokens_and_share_the_prompt_set(world):
    plan, args = plan_of(world, "K1")
    ei = [a for a in plan.arms if a["arm"] == "ei_sft"]
    rl = [a for a in plan.arms if a["arm"] == "rl"]
    assert len(ei) == len(rl) == 3 and {a["seed"] for a in ei} == {0, 1, 2}
    for e, r in zip(ei, rl):
        assert e["tokens_planned"] == BUDGET and e["max_problems"] == 20
        # RL knob derived so its planned tokens are within one optimizer step of the budget
        step_tokens = r["prompts_per_step"] * r["group"] * MEAN
        assert abs(r["tokens_planned"] - BUDGET) <= step_tokens / 2
        assert r["band_file"].endswith("round_01/select/problems.jsonl") and f"/s{e['seed']}/" in r["band_file"]
    rows = K.budget_table(plan, 0.5)
    assert rows[0]["contrast"] == "K1" and rows[0]["matched"].startswith("planned")
    cfgs = [s for s in plan.steps if s.kind == "write_config" and s.arm == "rl"]
    text = cfgs[0].writes[next(iter(cfgs[0].writes))]
    import tomllib
    t = tomllib.loads(text)
    assert t["train"]["optimizer_steps"] == rl[0]["optimizer_steps"] and t["train"]["seed"] == 0
    assert t["data"]["band_file"] == rl[0]["band_file"]


def test_set_toml_keys_is_section_aware_keeps_comments_and_refuses_missing_keys():
    src = 'a = 1  # top\n[x]\na = 2  # inner\nb = "s"\n[y]\na = 3\n'
    out = K.set_toml_keys(src, {("x", "a"): 9, ("", "a"): 7, ("y", "a"): "z"})
    assert out == 'a = 7  # top\n[x]\na = 9  # inner\nb = "s"\n[y]\na = "z"\n'
    with pytest.raises(K.ContrastError):
        K.set_toml_keys(src, {("x", "nope"): 1})


def test_sft_config_carries_the_seed_and_the_global_token_mean_setting(world):
    plan, _ = plan_of(world, "K0")
    cfgs = [s for s in plan.steps if s.kind == "write_config"]
    import tomllib
    seeds = sorted(tomllib.loads(next(iter(s.writes.values())))["train"]["seed"] for s in cfgs)
    assert seeds == [0, 1]
    assert tomllib.loads(next(iter(cfgs[0].writes.values())))["train"]["loss_normalization"] == "global_token_mean"


def test_measured_tokens_and_measure_verdict(world):
    plan, args = plan_of(world, "K6", classifier_pool=world / "clf.jsonl", random_pool=world / "rnd.jsonl", seeds=[0])
    arms = {a["arm"]: a for a in plan.arms}
    for name, n_tok in (("random", 100), ("classifier", 104)):
        gf = Path(arms[name]["gen_files"][0])
        write_jsonl(gf, [dict(gen_record(f"p{i}", 0), completion_ids=list(range(n_tok))) for i in range(3)])
    assert K.measured_generated_tokens([arms["random"]["gen_files"][0]]) == 300
    assert K.measure(plan, args, say=lambda *_: None) == K.EXIT_OK
    write_jsonl(Path(arms["classifier"]["gen_files"][0]), [dict(gen_record(f"p{i}", 0), completion_ids=list(range(150))) for i in range(3)])
    assert K.measure(plan, args, say=lambda *_: None) == K.EXIT_UNMATCHED


# =============================================================================== command lines


def test_ei_command_round_trips_through_expert_iterations_own_parser(world):
    plan, args = plan_of(world, "K1")
    s = next(x for x in plan.steps if x.kind == "ei")
    cfg = EI.parse_args(s.argv[2:])  # the real argparse + split guard + informal-overlap check
    assert cfg.max_problems == 20 and cfg.rounds == 1 and cfg.n == 8 and cfg.seed == 0
    assert cfg.protected_problems and cfg.allow_test_selection is False
    assert cfg.train_config.endswith("ei_sft_s0.sft.toml")


def test_k2_is_ready_now_that_the_temperature_flags_exist(world, capsys):
    """K2 was BLOCKED until lean_gen --temperature / expert_iteration --gen-temperature landed (merged 2026-10-19).
    The driver detects the flags from the real argparse definitions, so K2 must now plan as ready."""
    plan, args = plan_of(world, "K2", mutated_pool=world / "mutated.jsonl")
    assert {a["arm"] for a in plan.arms} == {"temp_matched", "mutated"}
    assert all(a["temperature"] == 1.2 for a in plan.arms)
    assert any("--gen-temperature" in p and ": present" in p for p in plan.prerequisites)
    assert not any("MISSING (pending)" in p for p in plan.prerequisites)
    assert not K.blocked_steps(plan)
    K.dry_run(plan, args)
    out = capsys.readouterr().out
    assert "BLOCKED" not in out and "--gen-temperature 1.2" in out


def test_k4_alpha_exports_scale_lora_alpha_and_reuse_the_arms_own_eval_for_k1_0(world):
    plan, args = plan_of(world, "K1", "K4", seeds=[0], k4_from=["K1"])
    ex = [s for s in plan.steps if s.kind == "alpha_export"]
    assert len(ex) == 2  # 0.5 and 0.7; 1.0 is the arm's existing eval
    alphas = sorted(float(s.argv[s.argv.index("--alpha") + 1]) for s in ex)
    assert alphas == [64.0, 89.6]  # 128 x {0.5, 0.7}
    assert all("P2_EXPORT_BASE_SNAPSHOT" in s.env and s.argv[s.argv.index("--ckpt") + 1].startswith(EI.LATEST) for s in ex)
    rep = next(s for s in plan.steps if s.sid.startswith("K4.") and s.kind == "report")
    assert rep.argv.count("--run") == 3 and "@a1.0" in " ".join(rep.argv)
    evs = [s for s in plan.steps if s.kind == "eval" and s.contrast == "K4"]
    assert len(evs) == 2 and all("--model-rev" in " ".join(s.argv) and "-K4-K1-ei_sft-s0-a0." in " ".join(s.argv) for s in evs)
    assert any("RL arms are not LoRA" in p for p in plan.prerequisites)


def test_k0_report_pairs_the_two_seeds_and_k1_reports_per_seed(world):
    plan, _ = plan_of(world, "K0")
    reps = [s for s in plan.steps if s.kind == "report"]
    assert len(reps) == 1 and "ei_s0" in " ".join(reps[0].argv) and "ei_s1" in " ".join(reps[0].argv)
    plan, _ = plan_of(world, "K1", pkpo_config=world / "rl.toml")
    reps = [s for s in plan.steps if s.kind == "report"]
    assert len(reps) == 6  # (rl, rl_pkpo) x 3 seeds, each vs ei_sft
    assert {a["arm"] for a in plan.arms} == {"ei_sft", "rl", "rl_pkpo"}


# =============================================================================== dry run, manifest, running


def test_dry_run_prints_commands_budget_and_touches_nothing(world, capsys):
    out = world / "out"
    assert K.main(base_argv(world, "K1", "K6", classifier_pool=world / "clf.jsonl", random_pool=world / "rnd.jsonl") + ["--dry-run"]) == K.EXIT_OK
    text = capsys.readouterr().out
    assert "expert_iteration.py" in text and "grpo.py" in text and "eval_report.py" in text and "# budget K1" in text
    assert not out.exists()


def test_print_manifest_is_valid_json_with_the_bookkeeping(world, capsys):
    assert K.main(base_argv(world, "K1", seeds=[0]) + ["--print-manifest"]) == K.EXIT_OK
    m = json.loads(capsys.readouterr().out)
    assert m["guard"]["allow_test_selection"] is False and m["budget"][0]["contrast"] == "K1"
    assert {s["kind"] for s in m["steps"]} >= {"write_config", "ei", "rl", "eval", "report"}
    assert m["inputs"]["mean_completion_tokens"] == MEAN and len(m["plan_sha"]) == 64


def test_run_materialises_configs_resumes_per_step_and_refuses_a_changed_step(world):
    plan, args = plan_of(world, "K0")

    class R(FakeRunner):
        def run(self, argv, env, log):
            super().run(argv, env, log)
            if "expert_iteration" in " ".join(argv):
                m = Path(argv[argv.index("--run-dir") + 1]) / "round_01" / "round_manifest.json"
                m.parent.mkdir(parents=True, exist_ok=True)
                m.write_text("{}")
            if "eval_report" in " ".join(argv):
                rd = Path(argv[argv.index("--out-dir") + 1])
                rd.mkdir(parents=True, exist_ok=True)
                (rd / "report.json").write_text("{}")
            return 0

    fr = R()
    assert K.run_plan(plan, args, fr, say=lambda *_: None) == K.EXIT_OK
    kinds = ["ei" if "expert_iteration" in " ".join(c) else "report" for c in fr.calls]
    assert kinds == ["ei", "ei", "report"]
    assert (Path(args.out_dir) / "contrasts_manifest.json").exists()
    assert all(Path(p).exists() for s in plan.steps for p in s.writes)
    fr2 = R()
    assert K.run_plan(plan, args, fr2, say=lambda *_: None) == K.EXIT_OK and fr2.calls == []  # everything done
    plan2, args2 = plan_of(world, "K0", n=4)
    assert K.run_plan(plan2, args2, R(), say=lambda *_: None) == K.EXIT_FAIL  # same markers, different command


def test_a_failing_step_stops_the_chain_and_leaves_no_marker(world):
    plan, args = plan_of(world, "K0")
    fr = FakeRunner(fail="expert_iteration")
    assert K.run_plan(plan, args, fr, say=lambda *_: None) == K.EXIT_FAIL and len(fr.calls) == 1
    assert not list((Path(args.out_dir) / "steps").glob("*.ei.done.json"))