"""Mocked native trainer orchestration; these fixtures are not GPU qualification."""

from __future__ import annotations

import copy
import json
from pathlib import Path
from types import SimpleNamespace

import pytest

from hessian_recipes import gym_training as gym
from hessian_recipes.agent_rl_tasks import generate_tasks
from hessian_recipes.evaluation import checkpoint_sha256, sha256_file


class RPCFailure(RuntimeError):
    pass


class Tensor:
    requires_grad = True

    def __init__(self, value=0.0):
        self.value = value

    def detach(self):
        return self

    def cpu(self):
        return self

    def clone(self):
        return Tensor(self.value)

    def float(self):
        return self

    def __sub__(self, other):
        return Tensor(self.value - other.value)

    def abs(self):
        return Tensor(abs(self.value))

    def max(self):
        return self.value


class Model:
    def __init__(self):
        self.config = SimpleNamespace(use_cache=True)
        self.generation_config = SimpleNamespace(eos_token_id=0)
        self.weight = Tensor()

    def named_parameters(self):
        return [("adapter.weight", self.weight)]


class Tokenizer:
    eos_token_id = 7

    def save_pretrained(self, path):
        Path(path).mkdir(parents=True, exist_ok=True)
        (Path(path) / "tokenizer.json").write_text("fixture")


@pytest.fixture
def runtime_setup(tmp_path, monkeypatch):
    model = tmp_path / "base"
    model.mkdir()
    (model / "config.json").write_text("{}")
    tasks = generate_tasks("train", 2)
    data = tmp_path / "train.jsonl"
    data.write_text("".join(json.dumps(row) + "\n" for row in tasks))
    config = tmp_path / "training.json"
    config.write_text(
        json.dumps(
            {
                "schema_version": 1,
                "model": {"path": str(model), "sha256": checkpoint_sha256(model)},
                "dataset": {"path": str(data), "sha256": sha256_file(data)},
                "output_dir": "run",
                "training": {
                    "max_steps": 1,
                    "num_generations": 2,
                    "per_device_train_batch_size": 2,
                },
            }
        )
    )
    state = SimpleNamespace(mode="success", cleanup_failure=False, environments=[], heavy_calls=0)
    identity = {
        "protocol": "hessian-gym-v1",
        "backend": "nemo_gym",
        "upstream_revision": "pinned-mock",
    }

    class Remote:
        def __init__(self, base_url, identity):
            self.index = len(state.environments)
            self._provenance = copy.deepcopy(identity)
            self._task = None
            self.failed = None
            self.closed = False
            state.environments.append(self)

        def reset(self, task=None, **row):
            self._task = copy.deepcopy(task)

        def lookup(self, record_id: str) -> str:
            if state.mode == "transport":
                self.failed = RPCFailure("primary RPC disconnect")
                raise self.failed
            return "tool fixture"

        def calculator(self, expression: str) -> str:
            return "tool fixture"

        def _raise_if_failed(self):
            if self.failed:
                raise self.failed

        def _score(self, completion):
            self._raise_if_failed()
            success = self.index % 2 == 0 and state.mode != "uniform"
            return {
                "success": success,
                "reward": float(success),
                "lookup_success": True,
                "calculation_success": success,
                "trace": [{"tool": "lookup", "ok": True}],
                "reason": "fixture",
            }

        def _close(self):
            self.closed = True
            if state.cleanup_failure:
                raise RPCFailure("secondary close failure")

    class NativeTrainer:
        def __init__(self, *, model, environment_factory, reward_funcs, train_dataset, **kwargs):
            self.model, self.reward = model, reward_funcs
            self.environments = [environment_factory() for _ in range(2)]
            self.task = train_dataset[0]["task"]
            self.state = SimpleNamespace(global_step=0, log_history=[])
            self.ref_model = None

        def _generate_single_turn(self, prompts):
            return [[1], [1]], [[2], [3]], None, {}

        def _generate(self, prompts):
            self._generate_single_turn(prompts)
            # Model tool calls can fail inside TRL's catch-and-continue loop.
            for environment in self.environments:
                try:
                    environment.lookup("record")
                except RPCFailure:
                    pass
            return [[1], [1]], [[2], [3]], [[1], [1]], [], 2, None, {}

        def train(self):
            for environment in self.environments:
                environment.reset(task=self.task)
            self._generate([])
            self.reward(
                [[{"role": "assistant", "content": "fixture"}]] * 2,
                self.environments,
                task=[self.task, self.task],
            )
            if state.mode != "zero_update":
                self.model.weight.value += 0.01
            self.state.log_history = [{"grad_norm": 0.0 if state.mode == "zero_update" else 1.0}]
            return SimpleNamespace(metrics={"train_loss": 0.1})

        def save_model(self, path):
            Path(path).mkdir(parents=True)
            (Path(path) / "adapter.bin").write_text(str(self.model.weight.value))

    def load_model(*args, **kwargs):
        state.heavy_calls += 1
        assert kwargs["local_files_only"] and not kwargs["trust_remote_code"]
        return Model()

    runtime = SimpleNamespace(
        torch=SimpleNamespace(
            cuda=SimpleNamespace(
                is_available=lambda: True,
                is_bf16_supported=lambda: True,
                device_count=lambda: 1,
                get_device_name=lambda: "mock GPU",
            ),
            bfloat16="bf16",
            isfinite=lambda tensor: SimpleNamespace(all=lambda: True),
        ),
        Dataset=SimpleNamespace(from_list=lambda rows: rows),
        LoraConfig=lambda **kwargs: kwargs,
        Model=SimpleNamespace(from_pretrained=load_model),
        Tokenizer=SimpleNamespace(from_pretrained=lambda *args, **kwargs: Tokenizer()),
        set_seed=lambda seed: None,
        Config=lambda **kwargs: SimpleNamespace(**kwargs),
        Trainer=NativeTrainer,
    )
    monkeypatch.setattr(gym, "_runtime", lambda: runtime)
    monkeypatch.setattr(
        gym,
        "_gym_api",
        lambda: SimpleNamespace(
            fetch_gym_identity=lambda url: copy.deepcopy(identity),
            RemoteGymEnv=Remote,
            GymInfrastructureError=RPCFailure,
        ),
    )
    monkeypatch.setattr(gym, "require_versions", lambda: {"mode": "mock"})
    monkeypatch.setattr(gym.importlib.metadata, "version", lambda name: "1.0.1")
    monkeypatch.setattr(
        gym.native,
        "configure_policy_tokenizer",
        lambda tokenizer: {"chat_template_sha256": "f" * 64},
    )
    return state, config, tmp_path / "run"


def test_dry_run_checks_gym_identity_without_loading_ml(runtime_setup):
    state, config, output = runtime_setup
    report = gym.plan(config, "http://127.0.0.1:9999")
    assert report["gym"]["identity"]["backend"] == "nemo_gym"
    assert report["validation_only"] and not report["release_eligible"]
    assert state.heavy_calls == 0 and state.environments == [] and not output.exists()


def test_native_orchestration_uses_remote_environments_and_captures_masks(runtime_setup):
    state, config, output = runtime_setup
    result = gym.run(config, "http://127.0.0.1:9999")
    assert result["status"] == "completed" and result["update_evidence"]["mixed_reward_groups"] == 1
    assert result["update_evidence"]["changed_trainable_tensors"] == 1
    assert result["task_coverage"]["total_rollouts"] == 2
    assert all(environment.closed for environment in state.environments)
    evidence = [
        json.loads(line) for line in (output / "rollout-evidence.jsonl").read_text().splitlines()
    ]
    assert [row["training_reward"] for row in evidence] == [1.0, 0.15]
    assert all(row["gym_provenance"]["backend"] == "nemo_gym" for row in evidence)
    assert all(row["policy_audit"]["sampled_token_equality_verified"] for row in evidence)
    assert json.loads((output / "training-manifest.json").read_text())[
        "evidence_sha256"
    ] == sha256_file(output / "rollout-evidence.jsonl")


@pytest.mark.parametrize("mode", ["zero_update", "uniform"])
def test_no_reward_driven_update_is_failed_even_if_native_train_returns(runtime_setup, mode):
    state, config, output = runtime_setup
    state.mode = mode
    with pytest.raises(RuntimeError, match="reward-driven"):
        gym.run(config, "http://127.0.0.1:9999")
    assert json.loads((output / "training-manifest.json").read_text())["status"] == "failed"
    assert all(environment.closed for environment in state.environments)


def test_swallowed_native_rpc_failure_escapes_before_scored_reward(runtime_setup):
    state, config, output = runtime_setup
    state.mode = "transport"
    with pytest.raises(RPCFailure, match="primary RPC"):
        gym.run(config, "http://127.0.0.1:9999")
    manifest = json.loads((output / "training-manifest.json").read_text())
    assert manifest["status"] == "failed" and manifest["error"]["type"] == "RPCFailure"
    assert not (output / "rollout-evidence.jsonl").exists()
    assert all(environment.closed for environment in state.environments)


def test_session_cleanup_cannot_mask_primary_rpc_error(runtime_setup):
    state, config, output = runtime_setup
    state.mode, state.cleanup_failure = "transport", True
    with pytest.raises(RPCFailure, match="primary RPC"):
        gym.run(config, "http://127.0.0.1:9999")
    manifest = json.loads((output / "training-manifest.json").read_text())
    assert "primary RPC" in manifest["error"]["message"] and len(manifest["cleanup_errors"]) == 2
    assert all(environment.closed for environment in state.environments)


def test_cleanup_failure_never_leaves_false_completed_artifact(runtime_setup):
    state, config, output = runtime_setup
    state.cleanup_failure = True
    with pytest.raises(RPCFailure, match="session cleanup failed"):
        gym.run(config, "http://127.0.0.1:9999")
    assert json.loads((output / "training-manifest.json").read_text())["status"] == "failed"