"""Apprentissage automatique : journalisation, format d'invite partagé, job LoRA (sur un Qwen2 minuscule, CPU)."""
import importlib.util
import json
from pathlib import Path

import pytest

from mouh.jev.types import JevMessageType as T, TaskPayload
from mouh.learning import TrajectoryRecorder, build_record, chat_messages, render_files, split_for, task_hash
from mouh.learning import prompts as P
from mouh.verification import CallableGenerator, TestTimeVerificationLoop
from mouh.verification.generators import parse_files

ROOT = Path(__file__).resolve().parents[1]
TESTS = {"test_solution.py": "from solution import add\n\ndef test_add():\n    assert add(2, 3) == 5\n"}


def load_job():
    spec = importlib.util.spec_from_file_location("train_unique", ROOT / "jobs" / "train_unique.py")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


JOB = load_job()


# ---------------------------------------------------------------- format partagé Space <-> job
@pytest.mark.parametrize("prev,diags", [(None, None), ({"solution.py": "x=1\n"}, ["boom", "  autre  "])])
def test_job_prompt_copy_is_identical(prev, diags):
    assert JOB.SYSTEM_PROMPT == P.SYSTEM_PROMPT
    assert JOB.chat_messages("t", TESTS, prev, diags) == P.chat_messages("t", TESTS, prev, diags)
    files = {"solution.py": "def add(a, b):\n    return a + b\n", "pkg/util.py": "X = 1\n"}
    assert JOB.render_files(files) == P.render_files(files)


def test_rendered_target_parses_back():
    files = {"solution.py": "def add(a, b):\n    return a + b\n", "pkg/util.py": "X = 1\n"}
    assert parse_files(render_files(files)) == files
    assert JOB.parse_files(render_files(files)) == files


# ---------------------------------------------------------------- bus et enregistrement
def run_with_repair():
    calls = []

    def gen(req):
        calls.append(req)
        return {"solution.py": "def add(a, b):\n    return a + b\n"}

    loop = TestTimeVerificationLoop(CallableGenerator(gen), max_attempts=3)
    task = TaskPayload(instruction="Écrire add", candidate_files={"solution.py": "def add(a, b):\n    return a - b\n"},
                       test_files=TESTS)
    return task, loop.run(task), calls


def test_router_decision_precedes_code_proposal():
    _, out, calls = run_with_repair()
    types = [m.type for m in out.messages]
    for i, t in enumerate(types):
        if t == T.CODE_PROPOSAL:
            assert types[i - 1] == T.ROUTE_DECISION, types
    assert calls and calls[0].experts  # le générateur reçoit les experts décidés par le routeur


def test_build_record_keeps_attempts_and_repair():
    task, out, _ = run_with_repair()
    rec = build_record(task, out, "test")
    assert rec["certified"] and rec["status"] == "certified"
    assert [a["certified"] for a in rec["attempts"]] == [False, True]
    assert rec["attempts"][0]["diagnostics"] and "a - b" in rec["attempts"][0]["files"]["solution.py"]
    assert rec["final_files"]["solution.py"].endswith("a + b\n")
    assert rec["split"] == split_for(rec["task_hash"]) and rec["task_hash"] == task_hash(task.instruction, TESTS)
    json.dumps(rec)  # sérialisable


def test_task_hash_ignores_whitespace_and_split_is_deterministic():
    assert task_hash("a  b", {"t.py": "x =  1"}) == task_hash("a b", {"t.py": "x = 1\n"})
    splits = {split_for(task_hash(f"t{i}", TESTS)) for i in range(200)}
    assert splits == {"train", "heldout"}


class FakeScheduler:
    def __init__(self, **kw):
        import threading

        self.kw, self.lock = kw, threading.Lock()


def test_recorder_local_and_remote(tmp_path):
    local = TrajectoryRecorder(tmp_path / "a")
    assert not local.remote
    local.record({"x": 1})
    assert json.loads(local.path.read_text()) == {"x": 1}
    remote = TrajectoryRecorder(tmp_path / "b", "me/ds", "tok", every_min=3, scheduler_factory=FakeScheduler)
    assert remote.remote and remote._scheduler.kw["private"] and remote._scheduler.kw["repo_type"] == "dataset"
    assert remote._scheduler.kw["allow_patterns"] == ["data/*.jsonl"]


def test_service_records_session(tmp_path, monkeypatch):
    import mouh.learning.recorder as R
    from mouh.service import stream_session

    monkeypatch.setattr(R, "_RECORDER", TrajectoryRecorder(tmp_path))
    snaps = list(stream_session("add", "def add(a, b):\n    return a + b\n", TESTS["test_solution.py"]))
    rep = json.loads(snaps[-1]["report"])
    assert rep["journalisation"].startswith("oui")
    rec = json.loads(R._RECORDER.path.read_text().splitlines()[-1])
    assert rec["certified"] and rec["schema"] == "mouh-session/1"
    snaps = list(stream_session("add", "def add(a, b):\n    return a + b\n", TESTS["test_solution.py"], record=False))
    assert json.loads(snaps[-1]["report"])["journalisation"].startswith("non")
    assert len(R._RECORDER.path.read_text().splitlines()) == 1


# ---------------------------------------------------------------- données du job
def session(i, *, certified=True, tests=None, split=None, failed_first=True):
    tests = tests or {"test_solution.py": f"from solution import f\n\ndef test_f():\n    assert f() == {i}\n"}
    th = task_hash(f"tâche {i}", tests)
    final = {"solution.py": f"def f():\n    return {i}\n"}
    attempts = []
    if failed_first:
        attempts.append({"attempt": 1, "origin": "user", "files": {"solution.py": "def f():\n    return -1\n"},
                         "score": 0.7, "certified": False, "diagnostics": [f"assert -1 == {i}"]})
    attempts.append({"attempt": len(attempts) + 1, "origin": "model", "files": final, "score": 1.0,
                     "certified": certified, "diagnostics": []})
    return {"schema": "mouh-session/1", "session_id": f"s{i}", "task_hash": th, "split": split or split_for(th),
            "instruction": f"tâche {i}", "tests": tests, "proofs": [], "generator": "x",
            "status": "certified" if certified else "exhausted", "certified": certified,
            "final_files": final if certified else None, "attempts": attempts}


def test_build_examples_filters_and_separates_heldout():
    sessions = [session(1, split="train"), session(2, split="heldout"), session(3, certified=False, split="train"),
                session(4, split="train", tests={"test_solution.py": "def test_x():\n    assert True\n"}),
                session(1, split="train")]  # doublon exact
    train, heldout, stats = JOB.build_examples(sessions)
    assert stats["kept"] == 3 and [t["task_hash"] for t in heldout] == [sessions[1]["task_hash"]]
    assert sorted(e["kind"] for e in train) == ["fix", "gen"]  # doublon retiré, tâche 4 rejetée (n'importe pas solution)
    assert all(sessions[1]["instruction"] not in json.dumps(e) for e in train)


def test_run_tests_really_executes():
    t = {"test_solution.py": "from solution import f\n\ndef test_f():\n    assert f() == 3\n"}
    assert JOB.run_tests({"solution.py": "def f():\n    return 3\n"}, t)
    assert not JOB.run_tests({"solution.py": "def f():\n    return 4\n"}, t)
    assert not JOB.run_tests({"solution.py": "import time\ndef f():\n    time.sleep(5)\n"}, t, timeout_s=2)


# ---------------------------------------------------------------- job complet sur un modèle minuscule
@pytest.fixture(scope="module")
def tiny_qwen(tmp_path_factory):
    pytest.importorskip("peft")
    from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
    from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM

    d = tmp_path_factory.mktemp("tiny_qwen")
    tk = Tokenizer(models.BPE())
    tk.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tk.decoder = decoders.ByteLevel()
    corpus = [P.SYSTEM_PROMPT, render_files({"solution.py": "def f():\n    return 1\n"}), "## Tâche ## Tests assert"] * 20
    tk.train_from_iterator(corpus, trainers.BpeTrainer(vocab_size=400, special_tokens=["<|endoftext|>", "<|im_start|>", "<|im_end|>"],
                                                        initial_alphabet=pre_tokenizers.ByteLevel.alphabet()))
    tok = PreTrainedTokenizerFast(tokenizer_object=tk, eos_token="<|im_end|>", pad_token="<|endoftext|>")
    tok.chat_template = ("{% for m in messages %}<|im_start|>{{ m['role'] }}\n{{ m['content'] }}<|im_end|>\n{% endfor %}"
                         "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}")
    tok.save_pretrained(d)
    cfg = Qwen2Config(vocab_size=len(tok), hidden_size=32, intermediate_size=64, num_hidden_layers=2,
                      num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=4096,
                      eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id, tie_word_embeddings=True)
    Qwen2ForCausalLM(cfg).save_pretrained(d)
    return str(d)


def write_sessions(path: Path, n: int, start: int = 0):
    path.mkdir(parents=True, exist_ok=True)
    with open(path / f"s{start}.jsonl", "w") as fh:
        for i in range(start, start + n):
            fh.write(json.dumps(session(i)) + "\n")


def test_unique_generator_runs_in_the_loop(tiny_qwen, monkeypatch):
    import mouh.service as S

    monkeypatch.setenv("UNIQUE_BASE_MODEL", tiny_qwen)
    monkeypatch.setenv("UNIQUE_MAX_NEW_TOKENS", "8")
    monkeypatch.delenv("UNIQUE_ADAPTER_REPO", raising=False)
    S._MODEL_CACHE.pop("unique", None)
    snaps = list(S.stream_session("add", "def add(a, b):\n    return a - b\n", TESTS["test_solution.py"], max_attempts=2,
                                  generator_choice=S.GENERATOR_CHOICES[2], record=False))
    assert snaps[-1]["status"].startswith("AUCUNE SORTIE")  # modèle aléatoire : rien de non prouvé n'est émis
    S._MODEL_CACHE.pop("unique", None)


# ---------------------------------------------------------------- pilote automatique (tout sur HF)


def test_app_builds():
    import importlib

    import app  # noqa: F401 — l'interface se construit sans réseau ni jeton

    importlib.reload(app)


# ---------------------------------------------------------------- alimentation en masse


# ---------------------------------------------------------------- boucle continue (TRL SFT + GRPO)
FAST = ["--min-heldout", "1", "--sft-chunk", "8", "--batch-size", "2", "--grad-accum", "1", "--max-len", "1024",
        "--grpo-steps", "2", "--grpo-batch", "4", "--grpo-generations", "2", "--grpo-max-completion", "16",
        "--max-new-tokens", "8", "--progress-every", "0", "--benchmark", "0", "--memory-ab", "0",
        "--promotion-margin", "1", "--measure-best-of", "0"]


def test_pytest_reward_scores_real_execution():
    r = JOB.make_pytest_reward(workers=2)
    tests = json.dumps({"test_solution.py": "from solution import f\n\ndef test_f():\n    assert f() == 3\n"})
    good = [{"role": "assistant", "content": render_files({"solution.py": "def f():\n    return 3\n"})}]
    bad = render_files({"solution.py": "def f():\n    return 4\n"})
    assert r(completions=[good, bad], tests=[tests] * 2) == [1.0, 0.1]
    assert r(completions=["### FILE: solution.py\n```python\ndef (\n```\n### END"], tests=[tests]) == [0.0]


def test_continuous_loop_runs_sft_and_grpo_with_trl(tiny_qwen, tmp_path):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 60)
    assert JOB.main(["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1",
                     *FAST]) == 0
    st = json.loads((out / "state.json").read_text())
    assert [h["phase"] for h in st["history"]] == ["sft", "grpo"]
    assert st["history"][0]["train"]["steps"] > 0 and st["history"][1]["train"]["steps"] == 2
    assert not any(h["promoted"] for h in st["history"])  # modèle aléatoire : rien n'est promu
    prog = json.loads((out / "progress.json").read_text())
    assert prog["phase"] == "fin" and prog["cycle"] == 1 and len(prog["history"]) == 2
    assert prog["base"]["n"] == len(st["heldout"]) and "rewards/pytest_reward/mean" in prog["metrics"]


def test_gate_promotes_only_strict_improvement_and_heldout_is_frozen(tiny_qwen, tmp_path, monkeypatch):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 60)
    scores = iter([1, 3, 2,   4, 4])  # passage 1 : base 1, SFT 3 (promu), GRPO 2 ; passage 2 : base en cache, SFT 4 (promu), GRPO 4
    monkeypatch.setattr(JOB, "pass_at_1", lambda *a, **k: {"passed": next(scores), "n": 5, "rate": 0, "ci95": [0, 1]})
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1",
            "--max-heldout", "3", *FAST]
    assert JOB.main(args) == 0
    st1 = json.loads((out / "state.json").read_text())
    assert [h["promoted"] for h in st1["history"]] == [True, False] and (out / "adapter" / "adapter_config.json").exists()
    write_sessions(data, 20, start=500)  # nouvelles données : le heldout figé ne change pas
    assert JOB.main(args) == 0
    st2 = json.loads((out / "state.json").read_text())
    assert st2["heldout"] == st1["heldout"] and st2["base"]["passed"] == 1
    assert [h["promoted"] for h in st2["history"]] == [True, False, True, False] and st2["best"]["passed"] == 4
    assert st2["sft_cursor"] == 16  # la 2e phase SFT a pris la tranche suivante


def test_daily_gpu_budget_pauses_the_loop(tiny_qwen, tmp_path):
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 30)
    assert JOB.main(["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen,
                     "--daily-gpu-hours", "0", *FAST]) == 0
    prog = json.loads((out / "progress.json").read_text())
    assert prog["phase"] == "pause" and "budget" in prog["status"] and not (out / "state.json").exists()


def test_waits_when_not_enough_heldout(tiny_qwen, tmp_path):
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 30)
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, *FAST]
    args[args.index("--min-heldout") + 1] = "1000"
    assert JOB.main(args) == 0
    assert json.loads((out / "progress.json").read_text())["phase"] == "attente"


def test_heldout_extension_remeasures_current_adapter(tiny_qwen, tmp_path, monkeypatch):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 40)
    scores = iter([1, 3, 0,   2, 5, 4, 4])  # passage 2 (heldout agrandi) : base 2, actuel 5, SFT 4 et GRPO 4 -> rien promu
    monkeypatch.setattr(JOB, "pass_at_1", lambda *a, **k: {"passed": next(scores), "n": 9, "rate": 0, "ci95": [0, 1]})
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1", *FAST]
    assert JOB.main(args) == 0
    write_sessions(data, 60, start=300)
    assert JOB.main(args) == 0
    st = json.loads((out / "state.json").read_text())
    assert st["best"]["passed"] == 5 and st["best"]["source"] == "actuel"
    assert [h["promoted"] for h in st["history"]] == [True, False, False, False]


# ---------------------------------------------------------------- mémoire à long terme
from mouh.learning.memory import TaskMemory, items_from_records  # noqa: E402


def mem_items():
    base = [("Écrire une fonction qui inverse une chaîne de caractères", "def f(s):\n    return s[::-1]\n"),
            ("Compter les voyelles dans une chaîne", "def f(s):\n    return sum(c in 'aeiou' for c in s)\n"),
            ("Calculer la somme des éléments pairs d'une liste", "def f(l):\n    return sum(x for x in l if x % 2 == 0)\n")]
    return [{"task_hash": f"h{i}", "instruction": t, "files": {"solution.py": c}} for i, (t, c) in enumerate(base)]


def test_memory_finds_similar_but_never_the_task_itself_or_near_duplicates():
    m = TaskMemory(mem_items())
    hits = m.search("Compter le nombre de voyelles d'un texte", k=1)
    assert hits and hits[0]["task_hash"] == "h1"
    assert all(h["task_hash"] != "h1" for h in m.search("Compter les voyelles dans une chaîne", k=3, exclude_hash="h1"))
    assert all(h["task_hash"] != "h0" for h in m.search("Écrire une fonction qui inverse une chaîne de caractères", k=3))
    jm = JOB.TaskMemory(mem_items())
    for q in ("voyelles texte", "somme pairs liste", "inverse chaîne"):
        assert [h["task_hash"] for h in jm.search(q, 2)] == [h["task_hash"] for h in m.search(q, 2)]


def test_memory_only_uses_certified_train_records():
    recs = [session(i) for i in range(30)] + [session(99, certified=False)]
    items = items_from_records(recs)
    assert items and all(r["split"] == "train" for r in recs if r["task_hash"] in {i["task_hash"] for i in items})
    assert session(99)["task_hash"] not in {i["task_hash"] for i in items}


def test_prompt_with_examples_identical_in_job():
    ex = mem_items()[:2]
    assert JOB.chat_messages("t", TESTS, examples=ex) == P.chat_messages("t", TESTS, examples=ex)
    assert "Exemples validés similaires" in P.user_prompt("t", TESTS, examples=ex)
    assert P.user_prompt("t", TESTS) == P.user_prompt("t", TESTS, examples=[])


# ---------------------------------------------------------------- curriculum + vrais échecs
def test_curriculum_prefers_informative_tasks():
    import random as R

    tasks = [{"task_hash": f"t{i}"} for i in range(300)]
    stats = {f"t{i}": ([4, 4] if i < 100 else [0, 10] if i < 200 else [2, 4]) for i in range(300)}
    picked = JOB.curriculum_select(tasks, stats, 60, R.Random(0))
    zone = sum(1 for t in picked if int(t["task_hash"][1:]) >= 200)
    assert len(picked) == 60 and zone >= 45  # surtout des tâches réussies « parfois »


def test_reward_logs_outcomes_and_real_failures():
    log_ = JOB.RolloutLog()
    r = JOB.make_pytest_reward(workers=2, log_to=log_)
    tests = json.dumps({"test_solution.py": "from solution import f\n\ndef test_f():\n    assert f() == 3\n"})
    good = render_files({"solution.py": "def f():\n    return 3\n"})
    bad = render_files({"solution.py": "def f():\n    return 4\n"})
    assert r(completions=[good, bad], tests=[tests] * 2, task_hash=["h", "h"], instruction=["énoncé", "énoncé"]) == [1.0, 0.1]
    assert log_.outcomes["h"] == [1, 2]
    (f,) = log_.failures
    assert "assert 4 == 3" in f["pytest_output"] and f["files"]["solution.py"].strip().endswith("return 4")


def test_decontamination_drops_benchmark_lookalikes():
    bench = [{"instruction": "Complète : def has_close_elements(numbers, threshold): Check if in given list of numbers, "
                             "are any two numbers closer to each other than given threshold."}]
    s1 = session(1)
    s2 = {**session(2), "instruction": "Check if in given list of numbers, are any two numbers closer to each other "
                                       "than given threshold. Return bool."}
    kept, removed = JOB.decontaminate([s1, s2], bench)
    assert removed == 1 and kept == [s1]


def humaneval_rows():
    return [{"task_id": f"HE/{i}", "entry_point": f"plus{i}",
             "prompt": f"def plus{i}(x):\n    \"\"\"Ajoute {i}.\"\"\"\n",
             "test": f"def check(candidate):\n    assert candidate(1) == {1 + i}\n"} for i in range(3)]


def test_humaneval_conversion_runs_real_checks(tmp_path):
    p = tmp_path / "he.jsonl"
    p.write_text("\n".join(json.dumps(r) for r in humaneval_rows()))
    t = JOB.load_humaneval(str(p), None)[1]
    assert JOB.run_tests({"solution.py": "def plus1(x):\n    return x + 1\n"}, t["tests"])
    assert not JOB.run_tests({"solution.py": "def plus1(x):\n    return x\n"}, t["tests"])


def test_loop_with_benchmark_memory_failures_and_repairs(tiny_qwen, tmp_path, monkeypatch):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 60)
    he = tmp_path / "he.jsonl"
    he.write_text("\n".join(json.dumps(r) for r in humaneval_rows()))
    # le modèle minuscule ne réussit rien : on simule des échecs réels pour vérifier le circuit complet
    real = JOB.make_pytest_reward

    def fake_reward(workers=8, timeout_s=20.0, log_to=None):
        def rw(prompts=None, completions=None, tests=None, task_hash=None, instruction=None, **kw):
            for th, ins, t in zip(task_hash, instruction, tests):
                if ins:
                    log_to.add(th, False, {"task_hash": th, "instruction": ins, "tests": json.loads(t),
                                           "files": {"solution.py": "def f():\n    return -1\n"},
                                           "pytest_output": "E   assert -1 == 1"})
            return [0.0] * len(completions)
        return rw

    monkeypatch.setattr(JOB, "make_pytest_reward", fake_reward)
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1", *FAST]
    args[args.index("--benchmark") + 1] = "1"
    args[args.index("--memory-ab") + 1] = "1"
    args += ["--local-benchmark", str(he)]
    assert JOB.main(args) == 0
    st = json.loads((out / "state.json").read_text())
    m = st["measures"]
    assert m["humaneval"]["n"] == 3 and m["humaneval_base"]["n"] == 3 and m["memory"]["with"]["n"] == len(st["heldout"])
    assert st["history"][1]["train"]["curriculum"]["failures"] > 0 and st["task_stats"]
    assert list((out / "failures").glob("*.jsonl"))
    assert JOB.main(args) == 0  # 2e passage : les échecs réels deviennent des prompts de réparation GRPO
    st2 = json.loads((out / "state.json").read_text())
    assert st2["history"][-1]["train"]["repairs"] > 0
    assert json.loads((out / "progress.json").read_text())["measures"]["humaneval"]["n"] == 3
    monkeypatch.setattr(JOB, "make_pytest_reward", real)


# ---------------------------------------------------------------- modèle maître


def test_sft_order_puts_the_teachers_very_hard_tasks_first_and_keeps_the_old_order_without_levels():
    import hashlib

    def rec(i, generator, level):
        attempt = {"attempt": 1, "origin": "dataset", "files": {"solution.py": "x"}, "score": 1.0,
                   "certified": True, "diagnostics": []}
        if level:
            attempt["level"] = level
        return {"schema": "mouh-session/1", "certified": True, "split": "train", "task_hash": f"h{i}",
                "instruction": f"tâche {i}", "generator": generator,
                "tests": {"test_solution.py": "from solution import f\n\n\ndef test_a():\n    assert f(1) == 1\n"},
                "final_files": {"solution.py": f"def f(x):\n    return {i}\n"}, "attempts": [attempt]}

    recs = [rec(i, "ingest:kodcode", None) for i in range(30)] + [
        rec(100, "teacher:m", "extrême"), rec(101, "teacher:m", "difficile"), rec(102, "teacher:m", "facile")]
    train, _, stats = JOB.build_examples(recs)
    order = JOB.sft_order(train, 0)
    tiers = [train[i]["tier"] for i in order]
    assert tiers[0] == 0 and tiers[1] == 1 and tiers.count(0) == 1 and stats["top_teacher"] == 1
    assert [train[i]["tier"] for i in order[2:]] == [2] * (len(train) - 2)
    # sans niveaux, l'ordre est exactement celui d'avant v3.11
    plain = [{k: v for k, v in e.items() if k != "tier"} for e in train[:30]]
    old = sorted(range(len(plain)), key=lambda i: hashlib.sha256(
        f"0:{JOB._sha(plain[i]['target'])}:{plain[i]['kind']}:{i}".encode()).hexdigest())
    assert JOB.sft_order(plain, 0) == old


# ---------------------------------------------------------------- interrupteurs (retour arrière)


def test_phase_error_is_reported_and_loop_continues(tiny_qwen, tmp_path, monkeypatch):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 60)

    def boom(*a, **k):
        raise ValueError("bf16 non pris en charge")

    monkeypatch.setattr(JOB, "run_sft", boom)
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1", *FAST]
    assert JOB.main(args) == 0
    st = json.loads((out / "state.json").read_text())
    assert st["last_error"]["phase"] == "sft" and "bf16" in st["last_error"]["error"]
    assert [h["phase"] for h in st["history"]] == ["grpo"]  # GRPO a quand même tourné
    monkeypatch.setattr(JOB, "run_grpo", boom)
    assert JOB.main(args) == 0  # 2 échecs d'affilée (SFT puis GRPO) : on continue
    assert JOB.main(args) == 1  # 3e échec d'affilée : arrêt plutôt que de payer le GPU pour rien
    assert json.loads((out / "progress.json").read_text())["phase"] == "erreur"


def test_precision_never_requests_bf16_on_t4(monkeypatch):
    import torch

    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda *a, **k: True)  # ce que PyTorch répond sur T4
    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a: (7, 5))
    assert JOB._precision() == {"fp16": True, "bf16": False}
    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a: (8, 9))
    assert JOB._precision() == {"bf16": True, "fp16": False}


def test_gpu_compat_disables_16_to_32_gemm_on_t4(monkeypatch):
    import torch
    from trl.kernels import chunked_logprob

    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a: (7, 5))
    monkeypatch.setattr(chunked_logprob, "_MM_OUT_DTYPE", True)
    JOB._gpu_compat()
    assert chunked_logprob._MM_OUT_DTYPE is False
    acc = torch.zeros(2, 2)
    chunked_logprob._addmm_fp32(acc, torch.ones(2, 3, dtype=torch.float16), torch.ones(3, 2, dtype=torch.float16))
    assert torch.equal(acc, torch.full((2, 2), 3.0))  # chemin de repli fp32


# ---------------------------------------------------------------- v3.5
def test_no_signal_stopper_stops_after_patience():
    import types

    info = {}
    cb = JOB._no_signal_stopper(0.9, 3, 2, info)
    state, control = types.SimpleNamespace(global_step=0), types.SimpleNamespace(should_training_stop=False)
    for step, z in enumerate([1.0, 0.5, 1.0, 1.0, 0.95], 1):
        state.global_step = step
        cb.on_log(None, state, control, logs={"frac_reward_zero_std": z})
    assert control.should_training_stop and info["early_stop"]["step"] == 5


def test_curriculum_prefers_hard_tasks_and_drops_solved_ones():
    import random as R

    tasks = [{"task_hash": f"t{i}", "level": "hard" if i < 50 else "easy"} for i in range(100)]
    picked = JOB.curriculum_select(tasks, {}, 30, R.Random(1))
    assert sum(t["level"] == "hard" for t in picked) >= 24
    solved = {f"t{i}": [2, 2] for i in range(50)}  # les difficiles déjà réussies 2 fois sur 2
    picked = JOB.curriculum_select(tasks, solved, 30, R.Random(1))
    assert sum(t["level"] == "hard" for t in picked) <= 12


def test_grpo_stops_early_when_every_attempt_has_the_same_score(tiny_qwen, tmp_path, monkeypatch):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 60)

    def always_one(workers=8, timeout_s=20.0, log_to=None):
        return lambda completions=None, **kw: [1.0] * len(completions)

    monkeypatch.setattr(JOB, "make_pytest_reward", always_one)
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1", *FAST,
            "--grpo-stop-patience", "2", "--grpo-stop-min-steps", "2"]
    args[args.index("--grpo-steps") + 1] = "8"
    assert JOB.main(args) == 0
    grpo = json.loads((out / "state.json").read_text())["history"][-1]["train"]
    assert grpo["steps"] < 8 and grpo["early_stop"]["step"] == grpo["steps"]


# ---------------------------------------------------------------- v3.6 : snapshot partagé data/ + failures/
def shared_snapshot(tmp_path):
    root = tmp_path / "snapshot"
    (root / "data").mkdir(parents=True)
    (root / "failures").mkdir()
    (root / "state").mkdir()
    (root / "data" / "ingest-x.jsonl").write_text("\n".join(json.dumps(session(i)) for i in range(5)) + "\n")
    fail = {"task_hash": "h", "instruction": "i", "tests": TESTS, "files": {"solution.py": "x=1\n"},
            "pytest_output": "E assert"}
    (root / "failures" / "r-c1.jsonl").write_text(json.dumps(fail) + "\n" + json.dumps({"incomplet": 1}) + "\n")
    (root / "state" / "teacher.json").write_text("{}")
    return root


def test_loaders_read_only_their_own_folder_of_a_shared_snapshot(tmp_path, monkeypatch):
    import huggingface_hub

    root = shared_snapshot(tmp_path)
    monkeypatch.setattr(huggingface_hub, "snapshot_download", lambda *a, **k: str(root))

    class ListApi:
        def __init__(self, token=None):
            pass

        def list_repo_files(self, repo, repo_type=None):
            return [str(f.relative_to(root)) for f in root.rglob("*") if f.is_file()]

    monkeypatch.setattr(huggingface_hub, "HfApi", ListApi)
    monkeypatch.setattr(huggingface_hub, "hf_hub_download", lambda repo, name, **k: str(root / name))
    sessions = JOB.load_sessions(None, "me/ds", "tok")
    assert len(sessions) == 5 and all(r["schema"] == "mouh-session/1" for r in sessions)
    store = types_store = type("S", (), {"local": None})()
    fails = JOB.load_failures(store, "me/ds", "tok", 100)
    assert len(fails) == 1 and fails[0]["files"]  # ni les sessions de data/, ni l'enregistrement incomplet


# ---------------------------------------------------------------- v3.7
def test_select_data_files_keeps_recent_ingest_slices_only():
    names = [f"data/ingest-opencodeinstruct-{i * 4000:09d}.jsonl" for i in range(50)] + \
            [f"data/ingest-kodcode-{i * 2000:09d}.jsonl" for i in range(3)] + \
            ["data/teacher-20261104T090000Z.jsonl", "data/sessions-abc.jsonl"]
    chosen = JOB.select_data_files(names, 30)
    oci = [n for n in chosen if "opencodeinstruct" in n]
    assert len(oci) == 30 and oci[0].endswith(f"{20 * 4000:09d}.jsonl")  # les 30 plus récentes
    assert sum("kodcode" in n for n in chosen) == 3
    assert "data/teacher-20261104T090000Z.jsonl" in chosen and "data/sessions-abc.jsonl" in chosen
    assert len(JOB.select_data_files(names, 0)) == len(names)


def test_frozen_heldout_survives_when_its_files_are_no_longer_loaded(tiny_qwen, tmp_path, monkeypatch):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    data.mkdir()

    def slice_(start, n, off):
        (data / f"ingest-opencodeinstruct-{off:09d}.jsonl").write_text(
            "\n".join(json.dumps(session(i)) for i in range(start, start + n)) + "\n")

    slice_(0, 60, 0)
    scores = iter([5, 6, 6, 7, 7])  # base 5 ; puis seulement des candidats : jamais de re-mesure de la base
    monkeypatch.setattr(JOB, "pass_at_1", lambda *a, **k: {"passed": next(scores), "n": 6, "rate": 0, "ci95": [0, 1]})
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1", *FAST,
            "--max-ingest-files", "1", "--max-heldout", "3"]
    assert JOB.main(args) == 0
    st1 = json.loads((out / "state.json").read_text())
    slice_(500, 60, 4000)  # nouvelle tranche : avec --max-ingest-files 1, la première n'est plus chargée
    assert JOB.main(args) == 0
    st2 = json.loads((out / "state.json").read_text())
    assert st2["heldout"] == st1["heldout"] and st2["heldout_sig"] == st1["heldout_sig"]
    assert set(st2["heldout_tasks"]) == set(st1["heldout"]) and st2["base"]["passed"] == 5


def test_promotion_needs_the_margin(tiny_qwen, tmp_path, monkeypatch):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 60)
    scores = iter([60, 62, 63])  # base 60 ; SFT 62 (+2) écarté ; GRPO 63 (+3) mis en service
    monkeypatch.setattr(JOB, "pass_at_1", lambda *a, **k: {"passed": next(scores), "n": 100, "rate": 0, "ci95": [0, 1]})
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1", *FAST]
    args[args.index("--promotion-margin") + 1] = "3"
    assert JOB.main(args) == 0
    assert [h["promoted"] for h in json.loads((out / "state.json").read_text())["history"]] == [False, True]


# ---------------------------------------------------------------- v3.8


def test_shrink_batches_keeps_effective_batch():
    import types

    a = types.SimpleNamespace(batch_size=2, grad_accum=8, grpo_batch=8, grpo_generations=4, grpo_max_completion=384,
                              eval_batch=16)
    t = JOB.shrink_batches(a, "sft")
    assert (t["batch_size"], t["grad_accum"]) == (1, 16)
    t = JOB.shrink_batches(a, "grpo")
    assert t["grpo_batch"] == 4 and t["eval_batch"] == 4
    t = JOB.shrink_batches(a, "grpo")
    assert t["grpo_batch"] == 4 and t["grpo_max_completion"] == 192


def test_out_of_memory_shrinks_batches_for_next_runs(tiny_qwen, tmp_path, monkeypatch):
    pytest.importorskip("trl")
    import torch

    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 60)
    real = JOB.run_sft

    def oom(*a, **k):
        raise torch.cuda.OutOfMemoryError("CUDA out of memory. Tried to allocate 2.00 GiB")

    monkeypatch.setattr(JOB, "run_sft", oom)
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1", *FAST]
    assert JOB.main(args) == 0
    st = json.loads((out / "state.json").read_text())
    assert st["tuning"]["batch_size"] == 1 and st["tuning"]["grad_accum"] == 2
    seen = {}
    monkeypatch.setattr(JOB, "run_sft", lambda a, *r, **k: seen.update(bs=a.batch_size) or real(a, *r, **k))
    assert JOB.main(args) == 0 and seen["bs"] == 1  # la réduction est reprise au passage suivant


def test_phase_that_cannot_fit_is_left_for_next_run(tiny_qwen, tmp_path):
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 40)
    out.mkdir()
    (out / "state.json").write_text(json.dumps({"phase_seconds": {"sft": 10_000_000}}))
    assert JOB.main(["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, *FAST]) == 0
    assert json.loads((out / "state.json").read_text()).get("history", []) == []


def test_progress_history_is_per_model(tmp_path):
    store = JOB.Store(local_out=str(tmp_path), adapter_repo="me/a", bucket="me/b", token=None)
    store.write_progress({"adapter": "me/ancien", "day": JOB._now().strftime("%Y-%m-%d"), "gpu_hours_today": 3.0,
                          "history": [{"x": 1}], "best": {"passed": 9}}, True)
    p = JOB.Progress(store, run_id="r", max_hours=1, daily_budget_h=24, min_interval_s=0, adapter="me/nouveau")
    assert p.state["history"] == [] and p.state["best"] is None and p.state["gpu_hours_before_run"] == 3.0


def test_best_of_measure_runs_and_is_reported(tiny_qwen, tmp_path):
    pytest.importorskip("trl")
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 40)
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--cycles", "1", *FAST,
            "--measure-best-of", "2"]
    assert JOB.main(args) == 0
    m = json.loads((out / "state.json").read_text())["measures"]
    assert m["best_of"]["k"] == 2 and m["best_of"]["n"] > 0


def test_prompt_notes_identical_in_job():
    notes = ["doc : utiliser heapq", "autre"]
    assert JOB.chat_messages("t", TESTS, notes=notes) == P.chat_messages("t", TESTS, notes=notes)
    assert "Notes de recherche web" in P.user_prompt("t", TESTS, notes=notes)


def test_complexity_estimate_orders_tasks():
    from mouh.learning.unique_generator import estimate_complexity

    easy = estimate_complexity("Écrire add(a, b).", TESTS)
    hard_tests = {"test_solution.py": "from solution import shortest, Graph\n" +
                  "".join(f"def test_{i}():\n    assert shortest(Graph(), {i}) == {i}\n" for i in range(14))}
    hard = estimate_complexity("Implémenter Dijkstra sur un graphe orienté pondéré, complexité O(n log n), "
                               "avec un tas (heap) et une classe Graph. " * 6, hard_tests)
    assert easy["level"] == "simple" and hard["level"] == "complexe" and hard["score"] > easy["score"]


def test_adaptive_session_picks_the_candidate_that_passes(monkeypatch):
    from mouh.learning import unique_generator as U
    from mouh.verification.generators import GenerationRequest
    from mouh.jev.types import JevMessageType as MT

    class Gen:
        def build_messages(self, r):
            return [{"role": "user", "content": "x"}]

        def candidates(self, messages, n):
            Gen.n = n
            return [{"solution.py": "def add(a, b):\n    return a - b\n"},
                    {"solution.py": "def add(a, b):\n    return a + b\n"}][:max(n, 1)]

    monkeypatch.setattr(U, "estimate_complexity", lambda i, t: {"level": "complexe", "score": 0.9})
    task = TaskPayload(instruction="add", test_files=TESTS)
    sess = U.UniqueSession(Gen(), adaptive=True)
    out = sess.generate(GenerationRequest(task=task, attempt=1, message_type=MT.TASK_SUBMISSION, experts=(1,)))
    assert out["solution.py"].endswith("a + b\n") and Gen.n == 6
    assert sess.info["essais"][0]["retenue"] == 2 and "recherche_web" not in sess.info   # v4.1 : aucune recherche web
    assert not hasattr(U, "web_notes")
    calm = U.UniqueSession(Gen(), adaptive=False)
    calm.generate(GenerationRequest(task=task, attempt=1, message_type=MT.TASK_SUBMISSION, experts=(1,)))
    assert Gen.n == 1


def test_reflective_generator_end_to_end(tiny_qwen, monkeypatch):
    import mouh.service as S

    monkeypatch.setenv("UNIQUE_BASE_MODEL", tiny_qwen)
    monkeypatch.setenv("UNIQUE_MAX_NEW_TOKENS", "8")
    monkeypatch.setenv("UNIQUE_MAX_SAMPLES", "2")
    monkeypatch.delenv("UNIQUE_ADAPTER_REPO", raising=False)
    S._MODEL_CACHE.pop("unique", None)
    snaps = list(S.stream_session("add", "def add(a, b):\n    return a - b\n", TESTS["test_solution.py"], max_attempts=2,
                                  generator_choice=S.GENERATOR_CHOICES[3], record=False))
    rep = json.loads(snaps[-1]["report"])
    assert snaps[-1]["status"].startswith("AUCUNE SORTIE") and rep["reflexion"]["complexite"]["level"] == "simple"
    S._MODEL_CACHE.pop("unique", None)


def test_v313_sft_extras_are_off_by_default_and_set_by_env(monkeypatch):
    monkeypatch.delenv("NEFTUNE_ALPHA", raising=False)
    monkeypatch.delenv("GRADIENT_CHECKPOINTING", raising=False)
    assert JOB._sft_extras() == {} and JOB._gradient_checkpointing() is True
    monkeypatch.setenv("NEFTUNE_ALPHA", "5")
    assert JOB._sft_extras() == {"neftune_noise_alpha": 5.0}
    monkeypatch.setenv("NEFTUNE_ALPHA", "abc")
    assert JOB._sft_extras() == {}
    monkeypatch.setenv("GRADIENT_CHECKPOINTING", "0")
    assert JOB._gradient_checkpointing() is False