"""CPU-only orchestration checks; synthetic proofs never claim native execution."""
from contextlib import contextmanager
import ast
import hashlib
from pathlib import Path

import pytest

from hessian_recipes import gym_experiment as nemo
from hessian_recipes import openenv_client, openenv_experiment as driver
import test_gym_experiment as nemo_tests


@pytest.fixture
def nemo_prepared(tmp_path):
    return nemo_tests.prepared.__wrapped__(tmp_path)


@pytest.fixture
def qualified(tmp_path, monkeypatch):
    # Explicit unit fixture: no SDK or model is executed by these tests.
    identity = {
        "backend": "openenv", "protocol": "hessian-openenv-v1",
        "upstream_revision": driver.NATIVE_REVISION,
        "server_sha256": driver.sha256_file(driver.ROOT / "examples/agent_gyms/openenv_server.py"),
        "environment_sha256": driver.sha256_file(driver.ROOT / "src/hessian_recipes/agent_rl_tasks.py"),
        "requirements_sha256": driver.sha256_file(driver.ROOT / "requirements/openenv.txt"),
        "runtime": {"version": driver.NATIVE_VERSION, "commit": driver.NATIVE_REVISION},
        "transport": "native-sdk-websocket", "active_environment_instances": 0,
    }
    report = tmp_path / "unit-native-proof.json"
    nemo._write(report, {
        "schema_version": 1, "status": "passed", "backend": "openenv",
        "mode": "native-sdk-cpu-scripted-fixtures", "native_sdk_executed": True,
        "model_inference": False, "gpu": False, "strict_reward_unchanged": True,
        "server_released": True, "checks": dict.fromkeys(driver.CPU_CHECKS, True),
        "native_identity": identity,
        "source_sha256": {key: driver.sha256_file(driver.ROOT / relative)
                          for key, relative in {
                              "server": "examples/agent_gyms/openenv_server.py",
                              "smoke": "examples/agent_gyms/openenv_smoke.py",
                              "requirements": "requirements/openenv.txt"}.items()},
        "upstream": {"revision": driver.NATIVE_REVISION, "version": driver.NATIVE_VERSION},
    })
    sdk = tmp_path / "sdk-python"
    sdk.write_text("unit fixture executable; never run")
    sdk.chmod(0o700)
    monkeypatch.setattr(openenv_client, "sdk_identity", lambda path: {"unit_sdk_fixture": True})
    monkeypatch.setattr(driver, "_sources", lambda: {"unit_source_fixture": "1" * 64})
    monkeypatch.setattr(driver.training, "_source_commit", lambda sources: None)
    return report, sdk


@pytest.fixture
def arm(nemo_prepared, qualified):
    proof, sdk = qualified
    output = nemo_prepared.root.parent / "openenv-arm"
    plan = driver.prepare(nemo_prepared.path, output, sdk, proof)
    return output, output / "experiment-plan.json", plan


def test_matched_prepare_copies_exact_bytes_and_only_changes_paths(nemo_prepared, qualified):
    proof, sdk = qualified
    before = {path: path.read_bytes() for path in nemo_prepared.root.rglob("*") if path.is_file()}
    root = nemo_prepared.root.parent / "arm"
    result = driver.prepare(nemo_prepared.path, root, sdk, proof)
    assert (root / "evaluation-plan/plan.json").read_bytes() == Path(nemo_prepared.plan["evaluation_plan"]).read_bytes()
    assert (root / "train.jsonl").read_bytes() == (nemo_prepared.root / "train.jsonl").read_bytes()
    expected = nemo._read(nemo_prepared.plan["training_config"])
    expected.update(output_dir=str(root / "training"), dataset={
        "path": str(root / "train.jsonl"), "sha256": nemo_prepared.plan["train_sha256"]})
    assert nemo._read(root / "training.json") == expected
    assert result["native_pin"]["sha256"] == driver.sha256_file(proof)
    assert result["nemo_origin"]["sha256"] == driver.sha256_file(nemo_prepared.path)
    assert result["baseline_reuse"] == {"enabled": True, "policy": driver.REUSE_POLICY}
    assert result["release_eligible"] is False
    assert {path: path.read_bytes() for path in nemo_prepared.root.rglob("*") if path.is_file()} == before
    assert driver.verify_plan(root / "experiment-plan.json") == result


def test_arm_rejects_existing_or_nested_destination(nemo_prepared, qualified, arm):
    proof, sdk = qualified
    with pytest.raises(ValueError, match="fresh"):
        driver.prepare(nemo_prepared.path, arm[0], sdk, proof)
    with pytest.raises(ValueError, match="separate arm"):
        driver.prepare(nemo_prepared.path, nemo_prepared.root / "nested", sdk, proof)


@pytest.mark.parametrize("field", ["model_inference", "gpu", "native_sdk_executed", "strict_reward_unchanged", "server_released", "checks", "upstream"])
def test_incomplete_native_qualification_fails_before_preparation(nemo_prepared, qualified, field):
    proof, sdk = qualified
    value = nemo._read(proof)
    value[field] = {} if field in {"checks", "upstream"} else not value[field]
    nemo._write(proof, value)
    output = nemo_prepared.root.parent / "must-not-exist"
    with pytest.raises(ValueError, match="CPU qualification"):
        driver.prepare(nemo_prepared.path, output, sdk, proof)
    assert not output.exists()


@pytest.mark.parametrize("target", ["config", "protocol", "train", "proof", "sdk", "origin"])
def test_tampered_inputs_cannot_resume_even_after_local_hash_resealing(arm, target):
    root, path, plan = arm
    if target == "config":
        config = nemo._read(root / "training.json")
        config["training"]["beta"] = 0.2
        nemo._write(root / "training.json", config)
        plan["training_config_sha256"] = driver.sha256_file(root / "training.json")
    elif target == "protocol":
        protocol = root / "evaluation-plan/plan.json"
        protocol.write_bytes(protocol.read_bytes() + b"\n")
        plan["evaluation_plan_sha256"] = driver.sha256_file(protocol)
    elif target == "train":
        train = root / "train.jsonl"
        train.write_bytes(train.read_bytes() + b"\n")
        plan["train_sha256"] = driver.sha256_file(train)
    elif target == "proof":
        proof = Path(plan["native_pin"]["path"])
        proof.write_bytes(proof.read_bytes() + b"\n")
    elif target == "sdk":
        Path(plan["sdk"]["python"]).write_text("changed executable")
    else:
        plan["nemo_origin"]["sha256"] = "0" * 64
    nemo._write(path, plan)
    with pytest.raises(ValueError, match="changed|differs"):
        driver.verify_plan(path)


def test_engine_reuses_unchanged_orchestration_without_shared_mutation():
    source = Path(nemo.__file__).read_bytes()
    original_globals = dict(nemo.__dict__)
    runtime = driver._engine("not-executed")
    assert runtime.adaptation_counts == {"trainer_imports": 6, "trainer_cli": 2}
    assert runtime.FORMAT == driver.FORMAT and runtime.FORMAT != nemo.FORMAT
    assert runtime._Journal is not nemo._Journal
    for name in ("_validate_suite_manifest", "_audit_merged"):
        # The shared evaluation/merge algorithms are unchanged code, not a second implementation.
        assert getattr(runtime, name).__code__.co_code == getattr(nemo, name).__code__.co_code
        assert getattr(runtime, name).__code__.co_consts == getattr(nemo, name).__code__.co_consts
    assert runtime._frozen_shared_suite.__code__.co_code == nemo._suite.__code__.co_code
    assert runtime._frozen_shared_suite.__code__.co_consts == nemo._suite.__code__.co_consts
    assert runtime.run.__globals__["_openenv_training"].FORMAT == "openenv-grpo-long-v1"
    assert "hessian_recipes.openenv_training" in runtime.run.__code__.co_consts
    assert "hessian_recipes.gym_training_long" not in runtime.run.__code__.co_consts
    assert nemo.__dict__ == original_globals and Path(nemo.__file__).read_bytes() == source


def test_frozen_source_and_rewrite_count_guards(tmp_path, monkeypatch):
    source = Path(nemo.__file__).read_bytes()
    copy = tmp_path / "engine.py"
    copy.write_bytes(source + b'\nextra = "hessian_recipes.gym_training_long"\n')
    monkeypatch.setattr(nemo, "__file__", str(copy))
    with pytest.raises(RuntimeError, match="source changed"):
        driver._engine("unused")
    monkeypatch.setattr(driver, "ENGINE_SHA256", hashlib.sha256(copy.read_bytes()).hexdigest())
    with pytest.raises(RuntimeError, match="adaptation boundary"):
        driver._engine("unused")


def test_ast_adapter_refuses_new_unsupported_trainer_import():
    with pytest.raises(RuntimeError, match="Unexpected"):
        driver._Adapt().visit(ast.parse("from .gym_training_long import new_api"))


def test_training_commands_use_openenv_sdk_and_keep_deadline(arm, monkeypatch):
    root, plan_path, plan = arm
    runtime = driver._engine(plan["sdk"]["python"])
    calls = []
    monkeypatch.setattr(runtime.subprocess, "run", lambda args, **kwargs: calls.append((args, kwargs)))
    monkeypatch.setattr(runtime.shutil, "which", lambda name: "/usr/bin/timeout")
    journal = runtime._Journal(plan_path, False)
    journal.command("train-to-64", ["-m", "hessian_recipes.openenv_training", str(root / "training.json"), "--stop-after-step", "64"])
    journal.close()
    args, options = calls[0]
    assert args[-2:] == ["--sdk-python", plan["sdk"]["python"]]
    assert "hessian_recipes.gym_training_long" not in args
    assert args[:2] == ["/usr/bin/timeout", "--signal=KILL"] and options["start_new_session"]


def test_qualification_only_closes_journal_and_never_starts_main_training(arm, monkeypatch):
    root, path, plan = arm
    runtime = driver._engine(plan["sdk"]["python"])
    events = []

    @contextmanager
    def service(*args):
        events.append("start-unit-fixture")
        try:
            yield "http://127.0.0.1:1", {"unit_fixture": True}
        finally:
            events.append("stop-unit-fixture")

    def command(journal, name, args):
        assert name == "gpu-preflight"
        nemo._write(Path(args[args.index("--output") + 1]), {"unit_fixture": True})

    def qualify(journal, url):
        proof = root / "resume-qualification/proof.json"
        nemo._write(proof, {"unit_fixture": True})
        journal.done("resume-qualification", [proof])

    runtime._service = service
    runtime._Journal.command = command
    runtime._resume_qualification = qualify
    monkeypatch.setattr(driver, "_engine", lambda sdk: runtime)
    result = driver.run(path, qualification_only=True)
    assert result["status"] == "qualification_completed" and result["active_attempt"] is None
    assert result == nemo._read(root / "experiment-status.json")
    assert events == ["start-unit-fixture", "stop-unit-fixture"]
    assert not (root / "training").exists() and not (root / "selection").exists()


def test_sdk_runtime_change_fails_before_gpu_engine(arm, monkeypatch):
    monkeypatch.setattr(openenv_client, "sdk_identity", lambda sdk: {"changed": True})
    monkeypatch.setattr(driver, "_engine", lambda sdk: pytest.fail("must not start GPU driver"))
    with pytest.raises(ValueError, match="SDK changed"):
        driver.run(arm[1])


@pytest.fixture
def reuse_source(arm, nemo_prepared, monkeypatch):
    """Test copy/journal contracts; raw scorer replay is qualified in frozen tests."""
    root, path, plan = arm
    runtime = driver._engine(plan["sdk"]["python"])
    journal = runtime._Journal(path, False)
    monkeypatch.setattr(driver.evaluation, "_verified_selection", lambda p, s: (
        {"selected": nemo_prepared.plan and nemo._read(nemo_prepared.plan["evaluation_plan"])["baseline"]},
        driver.sha256_file(s)))

    def load_suite(protocol, digest, manifest, split, *, selection_digest=None):
        entries = nemo._read(manifest)
        assert all(row["plan_sha256"] == digest and row["selection_sha256"] == selection_digest for row in entries)
        reports = {(row["kind"], row["seed"]): nemo._read(row["path"]) for row in entries}
        return {"reports": reports, "model": protocol["baseline"], "evidence": [
            {"kind": row["kind"], "seed": row["seed"], "path": row["path"],
             "sha256": driver.sha256_file(row["path"])} for row in entries]}

    monkeypatch.setattr(driver.evaluation, "_load_suite", load_suite)
    monkeypatch.setattr(driver.evaluation, "_compatible", lambda suites: None)
    monkeypatch.setattr(driver.evaluation.agent, "verify_report", lambda report: None)
    monkeypatch.setattr(driver.evaluation.numeric, "_verify_report", lambda report: None)
    monkeypatch.setattr(driver.importlib.metadata, "version", lambda name: "unit-version")

    def make(split="dev"):
        source_selection = nemo_prepared.root / "selection/selection.json" if split == "test" else None
        own_selection = root / "selection/selection.json" if split == "test" else None
        if source_selection:
            nemo._write(source_selection, {"unit_selection": "NeMo"})
            nemo._write(own_selection, {"unit_selection": "OpenEnv"})
        entries = driver.evaluation.write_specs(nemo_prepared.plan["evaluation_plan"],
                                                nemo_prepared.baseline, "baseline", split,
                                                selection_path=source_selection)
        manifest = Path(entries[0]["spec_path"]).parent / "manifest.json"
        completed = {}
        for row in entries:
            nemo._write(row["path"], {"explicit_unit_report": True,
                                      "provenance": {"packages": {"unit-package": "unit-version"}}})
            completed[f"baseline-{split}-{row['kind']}-{row['seed']}"] = {
                str(Path(file).resolve()): driver.sha256_file(file)
                for file in (row["path"], row["spec_path"], str(manifest))}
        nemo._write(nemo_prepared.root / "experiment-status.json", {
            "plan_sha256": driver.sha256_file(nemo_prepared.path), "completed": completed})
        return runtime, journal, entries, own_selection

    yield make
    journal.close()


@pytest.mark.parametrize("split", ["dev", "test"])
def test_verified_complete_baseline_reuse_copies_bytes_and_rebinds_selection(reuse_source, nemo_prepared, split):
    runtime, journal, entries, selection = reuse_source(split)
    journal.command = lambda *a: pytest.fail("No generation for an entirely reused baseline suite")
    target = runtime._suite(journal, nemo_prepared.baseline, "baseline", split, selection)
    own = nemo._read(target)
    assert len(own) == (2 if split == "dev" else 5)
    for source, destination in zip(entries, own, strict=True):
        assert Path(source["path"]).read_bytes() == Path(destination["path"]).read_bytes()
        assert Path(source["spec_path"]).read_bytes() == Path(destination["spec_path"]).read_bytes()
        assert Path(destination["path"]).is_relative_to(journal.root)
    receipt = nemo._read(journal.root / "baseline-reuse" / f"{split}.json")
    assert receipt["status"] == "reused" and len(receipt["reports"]) == len(own)
    if split == "test":
        assert receipt["source_selection_sha256"] != receipt["destination_selection_sha256"]
        assert all(row["selection_sha256"] == driver.sha256_file(selection) for row in own)
    runtime._suite(journal, nemo_prepared.baseline, "baseline", split, selection)
    Path(own[0]["path"]).write_bytes(b"tampered copied report")
    with pytest.raises(ValueError, match="evidence changed"):
        runtime._suite(journal, nemo_prepared.baseline, "baseline", split, selection)


@pytest.mark.parametrize("fault", ["missing-report", "package", "spec-bytes", "journal"])
def test_unavailable_or_incompatible_source_reruns_entire_suite(reuse_source, nemo_prepared, monkeypatch, fault):
    runtime, journal, entries, selection = reuse_source()
    if fault == "missing-report":
        Path(entries[0]["path"]).unlink()
    elif fault == "package":
        monkeypatch.setattr(driver.importlib.metadata, "version", lambda name: "different")
    elif fault == "spec-bytes":
        source = Path(entries[0]["spec_path"])
        source.write_bytes(source.read_bytes() + b"\n")
        entries[0]["spec_sha256"] = driver.sha256_file(source)
        nemo._write(source.parent / "manifest.json", entries)
        state = nemo._read(nemo_prepared.root / "experiment-status.json")
        for files in state["completed"].values():
            for file in files:
                files[file] = driver.sha256_file(file)
        nemo._write(nemo_prepared.root / "experiment-status.json", state)
    else:
        state = nemo._read(nemo_prepared.root / "experiment-status.json")
        state["completed"].pop(next(iter(state["completed"])))
        nemo._write(nemo_prepared.root / "experiment-status.json", state)
    commands = []

    def command(name, arguments):
        commands.append(name)
        output = Path(arguments[arguments.index("--output") + 1])
        nemo._write(output / "report.json", {"explicit_fresh_unit_report": True,
                                           "provenance": {"packages": {"unit-package": "unit-version"}}})

    journal.command = command
    runtime._suite(journal, nemo_prepared.baseline, "baseline", "dev")
    assert len(commands) == 2  # Never reuse the one compatible source item/report.
    assert nemo._read(journal.root / "baseline-reuse/dev.json")["status"] == "rerun"
    runtime._suite(journal, nemo_prepared.baseline, "baseline", "dev")
    assert len(commands) == 2  # Freeze the original suite-level decision on resume.


@pytest.mark.parametrize("identity_changed", [False, True])
def test_owned_service_checks_qualified_identity_and_closes_on_failure(arm, monkeypatch, identity_changed):
    root, _, plan = arm
    events = []

    @contextmanager
    def launcher(sdk, output):
        assert str(sdk) == plan["sdk"]["python"] and output.is_relative_to(root / "services")
        events.append("owned-start")
        try:
            yield {"url": "http://127.0.0.1:1", "pid": 123, "log_path": "unit-fixture"}
        finally:
            events.append("owned-stop")

    identity = plan["native_pin"]["native_identity"] | {"active_environment_instances": 7}
    if identity_changed:
        identity["upstream_revision"] = "0" * 40
    monkeypatch.setattr(driver, "_launcher", lambda: launcher)
    monkeypatch.setattr(openenv_client, "fetch_gym_identity", lambda url: identity)
    if identity_changed:
        with pytest.raises(RuntimeError, match="CPU-qualified"):
            with driver._service(root, plan["sdk"]["python"]):
                pytest.fail("must not expose mismatched service")
    else:
        with driver._service(root, plan["sdk"]["python"]) as (url, captured):
            assert url == "http://127.0.0.1:1" and captured == identity
    assert events == ["owned-start", "owned-stop"]


def test_source_inventory_is_explicit_and_includes_sequential_controller():
    sources = driver._sources()
    assert sources["src/hessian_recipes/gym_experiment.py"] == driver.ENGINE_SHA256
    for name in ("src/hessian_recipes/openenv_training.py", "src/hessian_recipes/openenv_client.py",
                 "examples/agent_gyms/openenv_smoke.py", "examples/agent_gyms/openenv_server.py",
                 "requirements/openenv.txt", "scripts/run-openenv-after-nemo.py",
                 "cluster/determined/german-qwen-openenv-command.yaml"):
        assert sources[name] == driver.sha256_file(driver.ROOT / name)
    assert not any("hf_recipe_hub" in name for name in sources)


def test_independent_full_arm_selects_after_all_dev_and_assesses_diagnostic256(arm, monkeypatch):
    root, path, plan = arm
    runtime = driver._engine(plan["sdk"]["python"])
    events = []

    @contextmanager
    def service(*args):
        try:
            yield "http://127.0.0.1:1", {"explicit_unit_service": True}
        finally:
            events.append("owned-service-closed")

    def command(journal, name, args):
        if name == "gpu-preflight":
            nemo._write(Path(args[args.index("--output") + 1]), {"explicit_unit_preflight": True})
        else:
            assert name.startswith("train-to-") and args[1] == "hessian_recipes.openenv_training"
            step = int(args[args.index("--stop-after-step") + 1])
            events.append(f"train{step}")
            nemo._write(root / "training" / f"checkpoint-{step}" / runtime._openenv_training.SEAL,
                        {"explicit_unit_seal": True})

    def suite(journal, model, label, split, selection=None):
        if split == "test":
            assert "selected-diagnostic256" in events and selection.is_file()
        if label == "candidate":
            assert Path(model) == root / "merged/step256"
        events.append(f"{label}-{split}")
        target = root / "unit-manifests" / label / split / "manifest.json"
        nemo._write(target, [])
        return target

    def merge(journal, checkpoint, step):
        target = root / "merged" / f"step{step}"
        target.mkdir(parents=True)
        return target

    def select(protocol, baseline, candidates, output):
        assert set(candidates) == {"64", "128", "256"}
        assert not any(event.endswith("-test") for event in events)
        events.append("selected-diagnostic256")
        nemo._write(output / "selection.json", {
            "selected_step": 256, "selected": {"path": str(root / "merged/step256")},
            "dev_acceptance": False, "retain_baseline": True})

    def summarize(protocol, selection, baseline, candidate, output):
        events.append("summary")
        value = {"status": "complete", "release_eligible": False, "overall_acceptance": False,
                 "selected_step": 256, "dev_acceptance": False, "retain_baseline": True}
        nemo._write(output / "summary.json", value)
        (output / "summary.md").write_text("Explicit unit fixture; no inference result.\n")
        return value

    runtime._service, runtime._suite = service, suite
    runtime._Journal.command = command
    runtime._resume_qualification = lambda *args: events.append("own-resume-qualified")
    runtime._export_and_merge = merge
    runtime._final_chain = lambda journal: {"backend": "openenv", "unit_fixture": True}
    runtime._verify_final_summary = lambda journal: nemo._read(root / "summary/summary.json")
    runtime._openenv_training.discover_resume_checkpoint = lambda output: None
    runtime._openenv_training.validate_checkpoint = lambda *args, **kwargs: {"unit_fixture": True}
    monkeypatch.setattr(driver.evaluation, "select", select)
    monkeypatch.setattr(driver.evaluation, "_verified_selection", lambda p, s: (nemo._read(s), driver.sha256_file(s)))
    monkeypatch.setattr(driver.evaluation, "summarize", summarize)
    monkeypatch.setattr(driver, "_engine", lambda sdk: runtime)
    result = driver.run(path)
    assert result["status"] == "completed" and result["active_attempt"] is None
    assert events == ["own-resume-qualified", "baseline-dev", "train64", "step64-dev",
                      "train128", "step128-dev", "train256", "step256-dev", "selected-diagnostic256",
                      "baseline-test", "candidate-test", "summary", "owned-service-closed"]
    assert nemo._read(root / "summary/summary.json")["overall_acceptance"] is False


def test_claimed_git_commit_must_contain_consumed_source_bytes(arm):
    _, path, plan = arm
    plan["source_commit"] = "0" * 40
    nemo._write(path, plan)
    with pytest.raises(ValueError, match="source commit"):
        driver.verify_plan(path)