"""Pinned native TRL online GRPO through a verified remote NeMo Gym server.

Generation/optimization stay in TRL. The remote client independently replays
server tool observations and rewards; transport failures never become task zeros.
"""

from __future__ import annotations

import argparse
import copy
from datetime import datetime, timezone
import importlib.metadata
import inspect
import json
from pathlib import Path
import statistics
import subprocess
from types import SimpleNamespace

from . import agent_rl_training as native
from .agent_rl_pilot import require_learning_update
from .evaluation import checkpoint_sha256
from .training import read_jsonl_snapshot, require_versions, sha256, write_manifest


def _gym_api():
    from . import gym_client

    return gym_client


def _sources():
    root = Path(__file__).parent
    return {
        name: sha256(root / name)
        for name in (
            "gym_training.py",
            "gym_client.py",
            "agent_rl_training.py",
            "agent_rl_tasks.py",
        )
    }


def _source_commit(sources):
    repo = Path(__file__).resolve().parents[2]
    try:
        commit = subprocess.check_output(
            ["git", "-C", str(repo), "rev-parse", "HEAD"], stderr=subprocess.DEVNULL, text=True
        ).strip()
        import hashlib

        for name, digest in sources.items():
            blob = subprocess.check_output(
                ["git", "-C", str(repo), "show", f"{commit}:src/hessian_recipes/{name}"],
                stderr=subprocess.DEVNULL,
            )
            if hashlib.sha256(blob).hexdigest() != digest:
                return None
        return commit
    except (OSError, subprocess.CalledProcessError):
        return None


def plan(config_path: str | Path, gym_url: str) -> dict:
    """Dry run performs a read-only native-server identity check, never GPU work."""
    report = native.plan(config_path)
    identity = _gym_api().fetch_gym_identity(gym_url)
    sources = _sources()
    report.update(
        method="native_trl_remote_nemo_gym_grpo",
        release_eligible=False,
        validation_only=True,
        qualification="Native gym integration and reward-driven update only; no heldout quality claim",
        gym={"url": gym_url, "identity": identity},
        gym_source_sha256=sources,
        source_commit=_source_commit(sources),
    )
    return report


def _runtime():
    import torch
    from datasets import Dataset
    from peft import LoraConfig
    from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
    from trl import GRPOConfig, GRPOTrainer

    return SimpleNamespace(
        torch=torch,
        Dataset=Dataset,
        LoraConfig=LoraConfig,
        Model=AutoModelForCausalLM,
        Tokenizer=AutoTokenizer,
        set_seed=set_seed,
        Config=GRPOConfig,
        Trainer=GRPOTrainer,
    )


def _gym_trainer_class(base):
    audited = native._audited_trainer_class(base)

    class GymTrainer(audited):
        def _generate(self, prompts):
            result = super()._generate(prompts)
            # Native TRL catches tool exceptions. Sticky client failures must escape
            # before any reward calculation or loss, even if generation continued.
            for environment in self.environments:
                environment._raise_if_failed()
            return result

    return GymTrainer


def _reward_callback(trainer_ref, settings, base_hash, evidence_path, groups):
    def environment_reward(completions, environments, **kwargs):
        trainer = trainer_ref[0]
        if trainer is None or len(completions) != len(environments):
            raise ValueError("Gym environment and rollout batches must align")
        for environment in environments:
            environment._raise_if_failed()
        metadata = native.audit_task_group(
            kwargs.get("task"), environments, settings["num_generations"]
        )
        scores = [
            environment._score(completion)
            for environment, completion in zip(environments, completions, strict=True)
        ]
        for environment in environments:
            environment._raise_if_failed()
        components = [native.milestone_training_reward(score) for score in scores]
        rewards = [item["training_reward"] for item in components]
        audits = trainer._last_policy_audit
        if len(audits) != len(scores):
            raise ValueError("Gym rollouts require one exact sampled-token audit each")
        version = {"global_step": trainer.state.global_step, "base_model_sha256": base_hash}
        groups.append(
            {
                "optimizer_step": trainer.state.global_step,
                "rewards": rewards,
                "reward_std": statistics.pstdev(rewards),
                "strict_successes": sum(score["success"] for score in scores),
                "strict_rewards": [score["reward"] for score in scores],
                "task_ids": [row["task_id"] for row in metadata],
                "episode_seeds": [row["episode_seed"] for row in metadata],
                "distinct_task_count": 1,
                "policy_version": version,
            }
        )
        with evidence_path.open("a", encoding="utf-8") as stream:
            for index, (environment, score, completion, audit) in enumerate(
                zip(environments, scores, completions, audits, strict=True)
            ):
                stream.write(
                    json.dumps(
                        {
                            "optimizer_step": trainer.state.global_step,
                            "rollout_index": index,
                            **metadata[index],
                            "policy_version": version,
                            "score": score,
                            "training_reward": rewards[index],
                            "reward_components": components[index],
                            "completion_messages": completion,
                            "policy_audit": audit,
                            "gym_provenance": copy.deepcopy(environment._provenance),
                        },
                        ensure_ascii=False,
                        allow_nan=False,
                    )
                    + "\n"
                )
        return rewards

    return environment_reward


def _close_sessions(environments):
    errors = []
    for environment in environments:
        try:
            environment._close()
        except Exception as error:
            errors.append({"type": type(error).__name__, "message": str(error)})
    return errors


def run(config_path: str | Path, gym_url: str) -> dict:
    report = plan(config_path, gym_url)
    rows, digest = read_jsonl_snapshot(report["dataset"]["path"])
    if digest != report["dataset"]["sha256"] or sha256(config_path) != report["config_sha256"]:
        raise ValueError("Gym recipe or tasks changed after validation")
    installed = require_versions()
    if importlib.metadata.version("jmespath") != "1.0.1":
        raise RuntimeError("Native Qwen tools require pinned jmespath==1.0.1")
    runtime, gym = _runtime(), _gym_api()
    torch = runtime.torch
    if (
        not torch.cuda.is_available()
        or not torch.cuda.is_bf16_supported()
        or torch.cuda.device_count() != 1
    ):
        raise RuntimeError("Use one visible CUDA GPU allocation with BF16 support")
    output = Path(report["output_dir"])
    output.mkdir(parents=True, exist_ok=False)
    manifest_path, evidence_path = (
        output / "training-manifest.json",
        output / "rollout-evidence.jsonl",
    )
    environments, groups, trainer_ref = [], [], [None]
    report.update(
        status="running",
        dependencies=installed | {"jmespath": "1.0.1"},
        started_at=datetime.now(timezone.utc).isoformat(),
        gpu=torch.cuda.get_device_name(),
        upstream_trainer_sha256=sha256(Path(inspect.getfile(runtime.Trainer))),
        evidence=str(evidence_path),
        reference_policy="input full checkpoint with new RL LoRA disabled",
    )
    write_manifest(manifest_path, report)
    try:
        runtime.set_seed(report["seed"])
        tokenizer = runtime.Tokenizer.from_pretrained(
            report["model"]["path"], local_files_only=True, trust_remote_code=False
        )
        report["policy_format"] = native.configure_policy_tokenizer(tokenizer)
        tokenizer.save_pretrained(str(output / "policy-tokenizer"))
        model = runtime.Model.from_pretrained(
            report["model"]["path"],
            local_files_only=True,
            trust_remote_code=False,
            use_safetensors=True,
            dtype=torch.bfloat16,
        )
        if checkpoint_sha256(report["model"]["path"]) != report["model"]["sha256"]:
            raise ValueError("Gym base checkpoint changed while loading")
        model.config.use_cache = False
        model.generation_config.eos_token_id = tokenizer.eos_token_id
        settings = report["training"]

        def factory():
            environment = gym.RemoteGymEnv(gym_url, identity=report["gym"]["identity"])
            environments.append(environment)
            # Public extras would be silently exposed to the model by native TRL.
            public = {
                name
                for name, _ in inspect.getmembers(environment, inspect.ismethod)
                if not name.startswith("_")
            }
            if public != {"reset", "lookup", "calculator"}:
                raise ValueError("Remote gym must expose only reset, lookup, calculator")
            return environment

        arguments = runtime.Config(
            output_dir=str(output),
            **settings,
            num_iterations=1,
            steps_per_generation=1,
            loss_type="grpo",
            full_determinism=True,
            bf16=True,
            fp16=False,
            gradient_checkpointing=True,
            gradient_checkpointing_kwargs={"use_reentrant": False},
            chat_template_kwargs=report["chat_template_kwargs"],
            use_vllm=False,
            disable_dropout=True,
            report_to="none",
            push_to_hub=False,
            seed=report["seed"],
            data_seed=report["seed"],
            optim="adamw_torch",
            save_strategy="no",
            logging_steps=1,
            dataloader_pin_memory=False,
            mask_truncated_completions=False,
            log_completions=False,
        )
        lora = report["lora"]
        trainer = _gym_trainer_class(runtime.Trainer)(
            model=model,
            processing_class=tokenizer,
            args=arguments,
            train_dataset=runtime.Dataset.from_list(
                [
                    {
                        "id": row["id"],
                        "prompt": [{"role": "user", "content": row["prompt"]}],
                        "task": row["task"],
                    }
                    for row in rows
                ]
            ),
            reward_funcs=_reward_callback(
                trainer_ref, settings, report["model"]["sha256"], evidence_path, groups
            ),
            environment_factory=factory,
            peft_config=runtime.LoraConfig(
                r=lora["r"],
                lora_alpha=lora["alpha"],
                lora_dropout=0.0,
                target_modules="all-linear",
                task_type="CAUSAL_LM",
            ),
        )
        trainer_ref[0] = trainer
        if trainer.ref_model is not None:
            raise ValueError("Gym pilot requires the tracked LoRA-disabled input reference")
        before = {
            name: tensor.detach().cpu().clone()
            for name, tensor in trainer.model.named_parameters()
            if tensor.requires_grad
        }
        if not before:
            raise ValueError("Gym trainer created no trainable adapter parameters")
        trained = trainer.train()
        changed, delta_max = [], 0.0
        for name, tensor in trainer.model.named_parameters():
            if name in before:
                current = tensor.detach().cpu()
                if not torch.isfinite(current).all():
                    raise ValueError("Nonfinite gym adapter parameter")
                delta = float((current.float() - before[name].float()).abs().max())
                if delta > 0:
                    changed.append(name)
                    delta_max = max(delta_max, delta)
        for environment in environments:
            environment._raise_if_failed()
        if (
            _sources() != report["gym_source_sha256"]
            or gym.fetch_gym_identity(gym_url) != report["gym"]["identity"]
        ):
            raise RuntimeError("Gym server or consumed source identity changed during training")
        trainer.save_model(str(output / "checkpoint"))
        tokenizer.save_pretrained(str(output / "checkpoint"))
        report.update(
            status="completed",
            checkpoint=str(output / "checkpoint"),
            checkpoint_sha256=checkpoint_sha256(output / "checkpoint"),
            metrics=trained.metrics,
            train_log_history=trainer.state.log_history,
            reward_groups=groups,
            task_coverage={
                "training_pool_rows": len(rows),
                "consumed_unique_task_ids": sorted(
                    {identity for group in groups for identity in group["task_ids"]}
                ),
                "total_rollouts": sum(len(group["task_ids"]) for group in groups),
                "optimizer_groups": len(groups),
            },
            update_evidence={
                "changed_trainable_tensors": len(changed),
                "max_abs_parameter_delta": delta_max,
                "mixed_reward_groups": sum(group["reward_std"] > 0 for group in groups),
                "nonzero_gradient_steps": sum(
                    float(log.get("grad_norm", 0)) > 0 for log in trainer.state.log_history
                ),
            },
            completed_at=datetime.now(timezone.utc).isoformat(),
        )
        require_learning_update(report)
        return report
    except BaseException as error:
        report.update(status="failed", error={"type": type(error).__name__, "message": str(error)})
        raise
    finally:
        cleanup = _close_sessions(environments)
        cleanup_failed_completed_run = False
        if cleanup:
            report["cleanup_errors"] = cleanup
            if report["status"] == "completed":
                cleanup_failed_completed_run = True
                report.update(
                    status="failed",
                    error={
                        "type": "GymInfrastructureError",
                        "message": "gym session cleanup failed",
                    },
                )
        if evidence_path.exists():
            report["evidence_sha256"] = sha256(evidence_path)
        write_manifest(manifest_path, report)
        if cleanup_failed_completed_run:
            raise gym.GymInfrastructureError("gym session cleanup failed")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("config", type=Path)
    parser.add_argument("--gym-url", required=True)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    result = plan(args.config, args.gym_url) if args.dry_run else run(args.config, args.gym_url)
    print(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False))


if __name__ == "__main__":
    main()