"""U11-11: `qwen3_parity.run` must not reuse saved tensors that another ckpt/code produced."""
import importlib.util
import json
import os

import pytest

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def _load():
    spec = importlib.util.spec_from_file_location(
        "parity_under_test", os.path.join(REPO, "scripts", "qwen3_parity.py")
    )
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


@pytest.fixture
def env(tmp_path, monkeypatch):
    par = _load()
    ckpt = tmp_path / "ckpt"
    ckpt.mkdir()
    (ckpt / "model-00001.safetensors").write_bytes(b"w" * 10)
    (ckpt / "config.json").write_text("{}")
    out = tmp_path / "out"
    calls = []

    def fake_run(cmd, timeout):
        if "tokenize" in cmd:
            d = cmd[cmd.index("tokenize") + 3]
            os.makedirs(d, exist_ok=True)
            json.dump({"input_ids": [[1, 2]]}, open(os.path.join(d, "ids.json"), "w"))
        elif "worker-align" in cmd:
            calls.append("align")
            d = cmd[cmd.index("--out") + 1]
            os.makedirs(d, exist_ok=True)
            json.dump({}, open(os.path.join(d, "meta.json"), "w"))
        else:
            calls.append("hf")
            d = cmd[cmd.index("hf") + 3]
            os.makedirs(d, exist_ok=True)
            json.dump({}, open(os.path.join(d, "meta.json"), "w"))
        return ""

    monkeypatch.setattr(par, "_run", fake_run)
    monkeypatch.setattr(par, "compare_dirs", lambda *a, **k: {"verdict": {}})
    return par, str(ckpt), str(out), calls


def test_identical_inputs_reuse_cached_tensors(env):
    par, ckpt, out, calls = env
    par.run(ckpt, out, "fp32")
    assert calls == ["align", "hf"]
    par.run(ckpt, out, "fp32")
    assert calls == ["align", "hf"]  # lock-race retry still skips finished workers


def test_changed_checkpoint_forces_rerun(env):
    par, ckpt, out, calls = env
    par.run(ckpt, out, "fp32")
    with open(os.path.join(ckpt, "model-00001.safetensors"), "wb") as f:
        f.write(b"other weights!")
    par.run(ckpt, out, "fp32")
    assert calls == ["align", "hf", "align", "hf"]


def test_changed_align_code_reruns_only_the_align_side(env, monkeypatch):
    par, ckpt, out, calls = env
    par.run(ckpt, out, "fp32")
    monkeypatch.setattr(par, "_align_code_sha", lambda: "edited-qwen3.py")
    par.run(ckpt, out, "fp32")
    assert calls == ["align", "hf", "align"]


def test_unstamped_legacy_output_is_not_trusted(env):
    par, ckpt, out, calls = env
    for side in ("align_fp32", "hf_fp32"):
        os.makedirs(os.path.join(out, side))
        json.dump({}, open(os.path.join(out, side, "meta.json"), "w"))
    par.run(ckpt, out, "fp32")
    assert calls == ["align", "hf"]