"""Real-file budget carry tests; synthetic release receipts never attest a server."""
import copy
import hashlib
import json
from pathlib import Path
from types import SimpleNamespace

import pytest
import test_gym_experiment as fixtures

from hessian_recipes import budget_lineage as budget
from hessian_recipes import gym_experiment as driver
from hessian_recipes import managed_training as managed


def write(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, sort_keys=True))


def make_carry(tmp_path, *, backend="nemo", used=1000., active=True):
    directory = tmp_path / ("prior-" + backend)
    now = budget.ORIGINAL_DEADLINE - 10000
    old_plan = {"root": str(directory), "max_gpu_wall_seconds": 28800, "release_eligible": False}
    if backend == "openenv":
        old_plan["backend"] = backend
    plan_path, journal_path = directory / "experiment-plan.json", directory / "experiment-status.json"
    write(plan_path, old_plan)
    plan_ref = budget.reference(plan_path)
    attempt = {"started_at_epoch": now - 100, "accounted_through_epoch": now - 30}
    if not active:
        attempt["finished_at_epoch"] = now - 30
    write(journal_path, {"plan_sha256": plan_ref["sha256"], "release_eligible": False,
        "status": "running" if active else "failed", "gpu_wall_seconds": used,
        "active_attempt": 0 if active else None, "attempts": [attempt]})
    source_receipt = directory / "owner-evidence.json"
    write(source_receipt, {"unit_fixture_only_not_server_attestation": True})
    release_path = directory / "release.json"
    write(release_path, {"format": budget.RELEASE_FORMAT, "backend": backend,
        "project_id": 363, "workspace_id": 123, "experiment_id": 7, "trial_id": 8,
        "released": True, "active_task_ids": [], "source_receipts": [budget.reference(source_receipt)],
        "experiment_plan_sha256": plan_ref["sha256"],
        "journal_sha256": budget.reference(journal_path)["sha256"], "tasks": [
            {"task_id": "7.task", "end_time_epoch": now - 10,
             "allocations": [{"allocation_id": "7.task.1", "state": "STATE_TERMINATED",
                              "start_time_epoch": now - 110, "end_time_epoch": now - 10}]}]})
    binding = {"max_wall_seconds": 72000, "backends": {backend: {
        "plan": str(plan_path), "plan_sha256": plan_ref["sha256"], "root": str(directory),
        "configuration": {"name": "unit-" + backend}}}}
    binding_path, progress_path = directory / "controller-binding.json", directory / "controller-progress.json"
    write(binding_path, binding)
    canonical = hashlib.sha256(json.dumps(binding, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
    child_name = "unit-" + backend + "-training"
    write(progress_path, {"binding_sha256": canonical, "status": "failed",
        "started_at": budget.ORIGINAL_DEADLINE - 72000, "deadline": budget.ORIGINAL_DEADLINE,
        "children": {child_name: {"id": 7, "creation_started": True,
            "stage": {"backend": backend, "operation": "training", "kind": "experiment", "name": child_name}}}})
    value = {"format": budget.FORMAT, "backend": backend, "cap_seconds": 28800,
        "absolute_deadline": budget.ORIGINAL_DEADLINE, "carry_in_seconds": used + (20 if active else 0),
        "prior_plan": plan_ref, "prior_journal": budget.reference(journal_path),
        "release_receipt": budget.reference(release_path), "controller_binding": budget.reference(binding_path),
        "controller_progress": budget.reference(progress_path),
        "owner_controller_binding": budget.reference(binding_path),
        "owner_controller_progress": budget.reference(progress_path)}
    path = directory / "carry.json"
    write(path, value)
    return SimpleNamespace(path=path, value=value, now=now, journal=journal_path,
                           release=release_path, prior=plan_path)


@pytest.fixture

def carried(tmp_path, monkeypatch):
    (tmp_path / "new").mkdir()
    prepared = fixtures.prepared.__wrapped__(tmp_path / "new")
    old = make_carry(tmp_path)
    plan = copy.deepcopy(prepared.plan)
    plan["budget_carry"] = budget.reference(old.path)
    write(prepared.path, plan)
    prepared.plan = plan
    clock = SimpleNamespace(wall=old.now, monotonic=100.)
    monkeypatch.setattr(driver.time, "time", lambda: clock.wall)
    monkeypatch.setattr(driver.time, "monotonic", lambda: clock.monotonic)
    return SimpleNamespace(prepared=prepared, old=old, clock=clock)


def advance(clock, seconds):
    clock.wall += seconds
    clock.monotonic += seconds


def test_abandoned_old_attempt_is_charged_through_release_without_repair(tmp_path):
    old = make_carry(tmp_path)
    before = old.journal.read_bytes()
    result = budget.validate(old.value, backend="nemo")
    assert result["carry_in_seconds"] == 1020 and result["unclosed_seconds"] == 20
    assert old.journal.read_bytes() == before
    assert json.loads(before)["active_attempt"] == 0


@pytest.mark.parametrize("key", ("prior_plan", "prior_journal", "release_receipt", "controller_binding", "controller_progress"))
def test_any_changed_bound_input_rejects_even_before_journal_creation(tmp_path, key):
    old = make_carry(tmp_path)
    path = Path(old.value[key]["path"])
    path.write_bytes(path.read_bytes() + b"\n")
    with pytest.raises(ValueError, match="hash changed"):
        budget.validate(old.value, backend="nemo")


@pytest.mark.parametrize("change", ("lower_charge", "higher_cap", "later_deadline", "wrong_backend"))
def test_carry_in_cannot_be_reduced_or_original_cap_deadline_renewed(tmp_path, change):
    old = make_carry(tmp_path)
    if change == "lower_charge":
        old.value["carry_in_seconds"] -= 1
    elif change == "higher_cap":
        old.value["cap_seconds"] += 1
    elif change == "later_deadline":
        old.value["absolute_deadline"] += 1
    else:
        old.value["backend"] = "openenv"
    with pytest.raises(ValueError):
        budget.validate(old.value, backend="nemo")


@pytest.mark.parametrize("change", ("active", "running", "missing_end", "ownership", "empty_allocations", "wrong_plan", "wrong_journal"))
def test_release_requires_terminal_owned_resources_and_actual_evidence_binding(tmp_path, change):
    old = make_carry(tmp_path)
    release = json.loads(old.release.read_text())
    allocation = release["tasks"][0]["allocations"][0]
    if change == "active":
        release["active_task_ids"] = ["7.task"]
    elif change == "running":
        allocation["state"] = "STATE_RUNNING"
    elif change == "missing_end":
        release["tasks"][0]["end_time_epoch"] = 0
    elif change == "ownership":
        release["experiment_id"] = 99
    elif change == "empty_allocations":
        release["tasks"][0]["allocations"] = []
    else:
        release["experiment_plan_sha256" if change == "wrong_plan" else "journal_sha256"] = "0" * 64
    write(old.release, release)
    old.value["release_receipt"] = budget.reference(old.release)
    with pytest.raises(ValueError):
        budget.validate(old.value, backend="nemo")


def test_new_journal_has_explicit_carry_without_inventing_old_attempts(carried):
    before = carried.old.journal.read_bytes()
    journal = driver._Journal(carried.prepared.path, False)
    assert journal.state["gpu_wall_seconds"] == 1020
    assert journal.state["budget_carry"]["carry_in_seconds"] == 1020
    assert len(journal.state["attempts"]) == 1 and journal.state["attempts"][0]["charged_gpu_wall_seconds"] == 0
    assert carried.old.journal.read_bytes() == before
    journal.close()


def test_every_native_training_and_assessment_command_spends_the_same_carried_cap(carried, monkeypatch):
    old = make_carry(carried.old.path.parent.parent / "cap-limited", used=20000.)
    plan = driver._read(carried.prepared.path)
    plan["budget_carry"] = budget.reference(old.path)
    write(carried.prepared.path, plan)
    journal = driver._Journal(carried.prepared.path, False)
    allowances = []
    def stage(arguments, **kwargs):
        allowances.append(float(arguments[2][:-1]))
        advance(carried.clock, 10)
    monkeypatch.setattr(driver.subprocess, "run", stage)
    for label in ("resume-qualification", "managed-qualification", "training", "baseline-dev", "candidate-dev", "agent-final", "mgsm-final"):
        journal.command(label, ["unit-fixture-not-executed"])
    assert allowances == [8780 - 10 * i for i in range(7)]
    assert journal.state["gpu_wall_seconds"] == 20090
    assert sum(a["charged_gpu_wall_seconds"] for a in journal.state["attempts"]) == 70
    journal.close()


def test_managed_guard_retains_global_deadline_instead_of_new_trial_deadline(carried):
    journal = driver._Journal(carried.prepared.path, False)
    advance(carried.clock, 9995)
    trial = SimpleNamespace(check_deadline=lambda: None, clock=lambda: carried.clock.wall,
                            binding={"deadline": budget.ORIGINAL_DEADLINE + 50000})
    assert managed._remaining(journal, trial) == 5
    advance(carried.clock, 6)
    with pytest.raises(TimeoutError, match="absolute deadline"):
        managed._remaining(journal, trial)
    journal.close()


def test_unclosed_new_allocation_is_conservatively_charged_on_resume(carried):
    journal = driver._Journal(carried.prepared.path, False)
    advance(carried.clock, 10)
    journal.save()
    advance(carried.clock, 300)
    recovered = driver._Journal(carried.prepared.path, True)
    assert recovered.state["gpu_wall_seconds"] == 1330
    assert recovered.state["attempts"][0]["conservative_unclosed_seconds"] == 300
    assert recovered.state["attempts"][0]["charged_gpu_wall_seconds"] == 310
    recovered.close()
    advance(carried.clock, 100)
    orderly = driver._Journal(carried.prepared.path, True)
    assert orderly.state["gpu_wall_seconds"] == 1330
    orderly.close()


@pytest.mark.parametrize("change", ("zero_counter", "remove_lineage", "reduce_local_charge", "erase_journal"))
def test_existing_carry_journal_cannot_restart_with_a_fresh_eight_hours(carried, change):
    journal = driver._Journal(carried.prepared.path, False)
    advance(carried.clock, 20)
    journal.close()
    state = driver._read(journal.path)
    if change == "zero_counter":
        state["gpu_wall_seconds"] = 0
    elif change == "remove_lineage":
        state.pop("budget_carry")
    elif change == "reduce_local_charge":
        state["attempts"][0]["charged_gpu_wall_seconds"] = 0
        state["gpu_wall_seconds"] = 1020
    else:
        journal.path.unlink()
    if change != "erase_journal":
        write(journal.path, state)
    before = journal.path.read_bytes() if journal.path.exists() else None
    with pytest.raises((ValueError, FileExistsError)):
        driver._Journal(carried.prepared.path, change != "erase_journal")
    assert (journal.path.read_bytes() if journal.path.exists() else None) == before


def test_exhausted_nemo_origin_is_cpu_verifiable_but_cannot_start_any_worker(carried):
    old = make_carry(carried.old.path.parent.parent / "exhausted", used=28800.451, active=False)
    plan = driver._read(carried.prepared.path)
    plan["budget_carry"] = budget.reference(old.path)
    write(carried.prepared.path, plan)
    assert driver.verify_plan(carried.prepared.path)["max_gpu_wall_seconds"] == 28800
    with pytest.raises(TimeoutError, match="eight-hour"):
        driver._Journal(carried.prepared.path, False)
    assert not (carried.prepared.root / "experiment-status.json").exists()
    assert not (carried.prepared.root / "budget-journal-started.json").exists()


def test_old_source_qualification_cannot_be_relabelled_for_carried_plan(carried):
    root = carried.prepared.root / "resume-qualification"
    root.mkdir()
    driver._write(root / "proof.json", {"status": "passed", "adapter_optimizer_scheduler_equal": True,
        "exact_rollout_evidence_equal": True, "rng_state_equal": True, "source_sha256": {"old-source": "old"}})
    (root / "train.jsonl").write_text("".join(json.dumps(r) + "\n" for r in driver.generate_tasks("train", 4, seed=43)))
    journal = SimpleNamespace(root=carried.prepared.root, plan=carried.prepared.plan, plan_path=carried.prepared.path)
    with pytest.raises(ValueError, match="another experiment"):
        driver._verify_resume_proof(journal)


def test_recovery_openenv_requires_its_own_carry_and_fresh_baselines(carried, tmp_path, monkeypatch):
    import test_openenv_experiment as arm_fixtures
    from hessian_recipes import openenv_experiment as arm
    proof, sdk = arm_fixtures.qualified.__wrapped__(tmp_path, monkeypatch)
    output = tmp_path / "new-openenv"
    with pytest.raises(ValueError, match="own carried budget"):
        arm.prepare(carried.prepared.path, output, sdk, proof, reuse_baseline=False)
    assert not output.exists()
    old = make_carry(tmp_path, backend="openenv")
    with pytest.raises(ValueError, match="fresh baseline"):
        arm.prepare(carried.prepared.path, output, sdk, proof, budget_carry=old.path)
    assert not output.exists()
    plan = arm.prepare(carried.prepared.path, output, sdk, proof,
                       reuse_baseline=False, budget_carry=old.path)
    assert plan["max_gpu_wall_seconds"] == 28800 and plan["budget_carry"] == budget.reference(old.path)
    assert budget.load(plan)["carry_in_seconds"] == 1020
    plan.pop("budget_carry")
    write(output / "experiment-plan.json", plan)
    with pytest.raises(ValueError, match="lost its carried budget"):
        arm.verify_plan(output / "experiment-plan.json")


def test_remaining_budget_is_not_a_replacement_for_the_original_plan_cap(carried):
    plan = driver._read(carried.prepared.path)
    plan["max_gpu_wall_seconds"] = 27780
    write(carried.prepared.path, plan)
    with pytest.raises(ValueError, match="source or inputs"):
        driver.verify_plan(carried.prepared.path)


def test_abandoned_charge_uses_allocation_end_not_later_task_or_observer_time(tmp_path):
    old = make_carry(tmp_path)
    release = json.loads(old.release.read_text())
    release["tasks"][0]["allocations"][0]["end_time_epoch"] -= 2
    write(old.release, release)
    old.value["release_receipt"] = budget.reference(old.release)
    old.value["carry_in_seconds"] = 1018
    assert budget.validate(old.value, backend="nemo")["unclosed_seconds"] == 18


def test_underlying_owner_receipt_drift_rejects_without_repair(tmp_path):
    old = make_carry(tmp_path)
    release = json.loads(old.release.read_text())
    path = Path(release["source_receipts"][0]["path"])
    path.write_text("{}")
    with pytest.raises(ValueError, match="hash changed"):
        budget.validate(old.value, backend="nemo")


def test_consistent_release_from_different_experiment_cannot_charge_this_owner(tmp_path):
    old = make_carry(tmp_path)
    release = json.loads(old.release.read_text())
    release["experiment_id"] = 99
    release["tasks"][0]["task_id"] = "99.task"
    release["tasks"][0]["allocations"][0]["allocation_id"] = "99.task.1"
    write(old.release, release)
    old.value["release_receipt"] = budget.reference(old.release)
    with pytest.raises(ValueError, match="exact recorded"):
        budget.validate(old.value, backend="nemo")


def test_running_owning_controller_does_not_qualify_as_a_released_recovery(tmp_path):
    old = make_carry(tmp_path)
    owner = json.loads(Path(old.value["owner_controller_progress"]["path"]).read_text())
    owner["status"] = "running"
    path = tmp_path / "different-owning-controller-progress.json"
    write(path, owner)
    old.value["owner_controller_progress"] = budget.reference(path)
    with pytest.raises(ValueError, match="owning controller"):
        budget.validate(old.value, backend="nemo")


def test_journal_rollback_cannot_erase_local_spend_while_high_water_survives(carried):
    journal = driver._Journal(carried.prepared.path, False)
    advance(carried.clock, 60)
    journal.close()
    state = driver._read(journal.path)
    state["attempts"] = []
    state["gpu_wall_seconds"] = state["budget_carry"]["carry_in_seconds"]
    write(journal.path, state)
    before = journal.path.read_bytes()
    with pytest.raises(ValueError, match="high-water"):
        driver._Journal(carried.prepared.path, True)
    assert journal.path.read_bytes() == before