"""CPU orchestration fixtures; actual GPU equivalence remains a separate gate."""
from __future__ import annotations

import copy
import inspect
import json
from pathlib import Path
import pickle

import pytest

from hessian_recipes import gym_training, gym_training_long, openenv_client, openenv_training as training
from test_gym_training_long import setup_run  # noqa: F401 -- reuse explicit scalar/RNG trainer fixture
from test_openenv_client import SDK, sdk_fixture  # noqa: F401 -- stateful transport fixture


URL = "http://127.0.0.1:8000"


@pytest.fixture
def prepared(setup_run, sdk_fixture, monkeypatch):  # noqa: F811 -- imported pytest fixtures
    original = gym_training_long
    setup = setup_run
    server = {"backend": "openenv", "upstream_revision": openenv_client.OPENENV_REVISION, "fixture": True}
    monkeypatch.setattr(openenv_client, "fetch_gym_identity", lambda url: copy.deepcopy(server))
    monkeypatch.setattr(openenv_client, "sdk_identity", lambda python: copy.deepcopy(SDK))
    runtime = training.engine("fixture-python")
    runtime.legacy._runtime = original.legacy._runtime
    runtime.require_versions = original.require_versions
    runtime._response_parser_identity = original._response_parser_identity
    return setup, runtime, server


def test_private_engines_never_mutate_original_modules():
    first, second = training.engine(), training.engine()
    assert first is not second and first.legacy is not second.legacy
    assert gym_training_long.FORMAT == "nemo-grpo-long-v1"
    assert first.FORMAT == "openenv-grpo-long-v1"
    assert first.run.__globals__ is first.__dict__
    assert first.legacy._reward_callback.__globals__ is first.legacy.__dict__
    assert Path(inspect.getfile(first.run)).name == "gym_training_long.py"
    assert pickle.loads(pickle.dumps(first.task_schedule)) is first.task_schedule
    class FixtureTrainer:
        pass
    derived = first._long_trainer_class(FixtureTrainer)
    assert Path(inspect.getfile(derived)).name == "gym_training_long.py"
    assert gym_training._gym_api().__name__ == "hessian_recipes.gym_client"
    first.FORMAT = "changed-in-private-fixture"
    assert second.FORMAT == training.FORMAT and gym_training_long.FORMAT == "nemo-grpo-long-v1"


def test_plan_freezes_native_backend_and_complete_sources(prepared):
    setup, runtime, server = prepared
    report = runtime.plan(setup.config(), URL, checkpoint_interval=2)
    assert report["method"] == "native_trl_remote_openenv_grpo_long"
    assert report["format_version"] == "openenv-grpo-long-v1" and report["backend"] == "openenv"
    assert report["gym"]["identity"] == server | {"client_sdk_runtime": SDK}
    assert report["openenv"]["sdk_runtime"] == SDK
    assert report["engine_reuse"]["frozen_engine_sha256"] == training.FROZEN_ENGINES
    sources = report["long_source_sha256"]
    for name in ("src/hessian_recipes/openenv_client.py", "src/hessian_recipes/openenv_training.py",
                 "examples/agent_gyms/openenv_server.py", "requirements/openenv.txt"):
        assert len(sources[name]) == 64
    assert setup.state.heavy_calls == 0


def test_native_tokens_optimizer_rng_resume_and_openenv_seals(prepared):
    setup, runtime, _ = prepared
    full = runtime.run(setup.config("full"), URL, checkpoint_interval=2)
    config = setup.config("segmented")
    first = runtime.run(config, URL, checkpoint_interval=2, stop_after_step=2)
    seal = training.validate_checkpoint(setup.root / "segmented", first["checkpoint"], expected_step=2)
    assert seal["format_version"] == training.FORMAT
    resumed = runtime.run(config, "http://127.0.0.1:8001", checkpoint_interval=2,
                          resume_from_checkpoint=first["checkpoint"])
    assert resumed["status"] == "completed" and resumed["optimizer_step"] == 4
    assert (Path(full["checkpoint"]) / "adapter_model.safetensors").read_bytes() == (
        Path(resumed["checkpoint"]) / "adapter_model.safetensors").read_bytes()
    assert (setup.root / "full/rollout-evidence.jsonl").read_bytes() == (
        setup.root / "segmented/rollout-evidence.jsonl").read_bytes()
    assert training.discover_resume_checkpoint(setup.root / "segmented") == Path(resumed["checkpoint"])
    assert resumed["learning_evidence"]["cumulative"]["nonzero_gradient_steps"] == 4
    assert resumed["learning_evidence"]["from_initial_adapter"]["changed_trainable_tensors"] == 1
    assert gym_training_long.FORMAT == "nemo-grpo-long-v1"


def test_sdk_identity_change_blocks_resume_before_model_loading(prepared, monkeypatch):
    setup, runtime, _ = prepared
    config = setup.config()
    first = runtime.run(config, URL, checkpoint_interval=2, stop_after_step=2)
    before = setup.state.heavy_calls
    monkeypatch.setattr(openenv_client, "sdk_identity", lambda python: SDK | {"modified": True})
    with pytest.raises(ValueError, match="identity changed"):
        runtime.run(config, URL, checkpoint_interval=2, resume_from_checkpoint=first["checkpoint"])
    assert setup.state.heavy_calls == before


def test_openenv_source_mutation_invalidates_checkpoint(prepared, monkeypatch):
    setup, runtime, _ = prepared
    first = runtime.run(setup.config(), URL, checkpoint_interval=2, stop_after_step=2)
    original = training.source_hashes
    monkeypatch.setattr(training, "source_hashes", lambda: original() | {"changed.py": "0" * 64})
    with pytest.raises(ValueError, match="source"):
        training.validate_checkpoint(setup.root / "run", first["checkpoint"])


def test_unqualified_engine_source_is_rejected(monkeypatch):
    monkeypatch.setitem(training.FROZEN_ENGINES, "gym_training.py", "0" * 64)
    with pytest.raises(RuntimeError, match="Unqualified"):
        training.engine()


def test_sdk_required_for_live_plan_but_not_artifact_validation(prepared, monkeypatch):
    setup, _, _ = prepared
    with pytest.raises(ValueError, match="sdk_python"):
        training.engine().plan(setup.config(), URL, checkpoint_interval=2)


def test_no_cpu_model_execution_even_with_native_sdk(prepared):
    setup, runtime, _ = prepared
    runtime.legacy._runtime().torch.cuda.is_available = lambda: False
    with pytest.raises(RuntimeError, match="allocated CUDA"):
        runtime.run(setup.config(), URL, checkpoint_interval=2)
    assert setup.state.heavy_calls == 0


def test_payload_tamper_is_rejected_without_optimizer_pickle_loading(prepared):
    setup, runtime, _ = prepared
    result = runtime.run(setup.config(), URL, checkpoint_interval=2, stop_after_step=2)
    path = Path(result["checkpoint"])
    (path / "optimizer.pt").write_bytes(b"not a trusted optimizer payload")
    with pytest.raises(ValueError, match="hashes"):
        training.validate_checkpoint(setup.root / "run", path)
    manifest = json.loads((setup.root / "run/training-manifest.json").read_text())
    assert manifest["backend"] == "openenv"