import json
from pathlib import Path
import sys
from types import SimpleNamespace

import pytest

from hessian_recipes import agent_rl_pilot as pilot


def frozen_plan(tmp_path, monkeypatch):
    monkeypatch.setattr(pilot, "checkpoint_sha256", lambda path: "model-hash")
    monkeypatch.setattr(pilot, "source_hashes", lambda: {"implementation": "frozen"})
    for name in ("training.json", "train.jsonl", "dev.jsonl", "test.jsonl",
                 "candidate/report.json", "candidate-spec.json", "evaluation.jsonl"):
        path = tmp_path / name
        path.parent.mkdir(exist_ok=True)
        path.write_text("{}\n")
    return {
        "model": str(tmp_path / "model"), "model_sha256": "model-hash",
        "tokenizer_sha256": "model-hash", "source_sha256": pilot.source_hashes(),
        "training_config_sha256": pilot.sha256_file(tmp_path / "training.json"),
        "split_sha256": {split: pilot.sha256_file(tmp_path / f"{split}.jsonl")
                         for split in ("train", "dev", "test")},
        "retention_run": str(tmp_path),
        "retention_baseline_report_sha256": pilot.sha256_file(tmp_path / "candidate/report.json"),
        "retention_baseline_spec_sha256": pilot.sha256_file(tmp_path / "candidate-spec.json"),
        "retention_tasks_sha256": pilot.sha256_file(tmp_path / "evaluation.jsonl"),
    }


@pytest.mark.parametrize("changed", ["train.jsonl", "test.jsonl", "candidate/report.json",
                                    "candidate-spec.json", "evaluation.jsonl", "training.json"])
def test_frozen_inputs_reject_changed_training_test_and_retention(tmp_path, monkeypatch, changed):
    plan = frozen_plan(tmp_path, monkeypatch)
    pilot.verify_inputs(tmp_path, plan)
    (tmp_path / changed).write_text('{"changed":true}\n')
    with pytest.raises(ValueError, match="Prepared input changed"):
        pilot.verify_inputs(tmp_path, plan)


def test_frozen_inputs_reject_changed_source(tmp_path, monkeypatch):
    plan = frozen_plan(tmp_path, monkeypatch)
    monkeypatch.setattr(pilot, "source_hashes", lambda: {"implementation": "edited"})
    with pytest.raises(ValueError, match="sources"):
        pilot.verify_inputs(tmp_path, plan)


def test_execute_refuses_cpu_before_any_stage(tmp_path, monkeypatch):
    (tmp_path / "plan.json").write_text(json.dumps({"status": "prepared"}))
    monkeypatch.setitem(sys.modules, "torch", SimpleNamespace(
        cuda=SimpleNamespace(is_available=lambda: False)))
    with pytest.raises(RuntimeError, match="CUDA GPU allocation"):
        pilot.execute(tmp_path / "plan.json")
    assert not (tmp_path / "run-manifest.json").exists()
    assert not (tmp_path / "logs").exists()


def test_prepare_rejects_invalid_budget_without_creating_output(tmp_path):
    output = tmp_path / "prepared"
    with pytest.raises(ValueError, match="steps"):
        pilot.prepare(Path("/missing"), output, retention_run=Path("/missing"), steps=0)
    assert not output.exists()


@pytest.mark.parametrize("missing", ["changed_trainable_tensors", "mixed_reward_groups",
                                    "nonzero_gradient_steps"])
def test_optimizer_completion_requires_actual_reward_driven_update(missing):
    report = {"status": "completed", "update_evidence": {
        "changed_trainable_tensors": 392, "mixed_reward_groups": 2, "nonzero_gradient_steps": 2,
    }}
    pilot.require_learning_update(report)
    report["update_evidence"][missing] = 0
    with pytest.raises(RuntimeError, match="reward-driven parameter update"):
        pilot.require_learning_update(report)