"""Source-bound OpenEnv backend for the frozen native TRL long-run engine.

This executes verified engine bytes in private module namespaces. The original
NeMo modules are never patched. Generation, loss masks, optimizer, schedules,
checkpoint sealing and resume remain the previously qualified engine's code.
Only environment transport and explicit backend/source provenance are injected.
"""
from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
import subprocess
import sys
from types import ModuleType, SimpleNamespace
import uuid

from . import openenv_client
from .training import sha256

FORMAT = "openenv-grpo-long-v1"
FROZEN_ENGINES = {
    "gym_training.py": "c57964e6c9977c09228d9470de281925b4686761fc55b10a5c5d645e71a354a4",
    "gym_training_long.py": "98c7d0ded9fd46126821331d1f97e750ca15140d1e2da01ce6bd3371bf9a58e4",
}


def source_hashes() -> dict:
    """Exact consumed repository paths; new backend identity cannot claim NeMo."""
    root = Path(__file__).resolve().parents[2]
    package = (
        "__init__.py", "openenv_training.py", "openenv_client.py", "gym_training_long.py",
        "gym_training.py", "gym_client.py", "agent_rl_training.py", "agent_rl_tasks.py",
        "agent_rl_pilot.py", "training.py", "evaluation.py", "distillation.py",
        "managed_trial.py", "managed_training.py", "budget_lineage.py",
    )
    names = [f"src/hessian_recipes/{name}" for name in package]
    names += ["examples/agent_gyms/openenv_server.py", "examples/agent_gyms/openenv_smoke.py",
              "requirements/openenv.txt"]
    return {name: sha256(root / name) for name in names}


def _source_commit(sources: dict) -> str | None:
    root = Path(__file__).resolve().parents[2]
    try:
        commit = subprocess.check_output(["git", "-C", str(root), "rev-parse", "HEAD"],
                                         stderr=subprocess.DEVNULL, text=True).strip()
        for name, digest in sources.items():
            blob = subprocess.check_output(["git", "-C", str(root), "show", f"{commit}:{name}"],
                                           stderr=subprocess.DEVNULL)
            if hashlib.sha256(blob).hexdigest() != digest:
                return None
        return commit
    except (OSError, subprocess.CalledProcessError):
        return None


def _load_engine(name: str) -> ModuleType:
    """Compile the very bytes checked here; no read-after-validation source race."""
    path = Path(__file__).parent / name
    source = path.read_bytes()
    if hashlib.sha256(source).hexdigest() != FROZEN_ENGINES[name]:
        raise RuntimeError(f"Unqualified upstream engine source changed: {name}")
    module = ModuleType(f"hessian_recipes._openenv_{path.stem}_{uuid.uuid4().hex}")
    module.__file__, module.__package__ = str(path), "hessian_recipes"
    # Registration makes inspect.getfile and ordinary function references work;
    # these unique names never replace the real NeMo/package modules.
    sys.modules[module.__name__] = module
    try:
        exec(compile(source, str(path), "exec"), module.__dict__)
    except BaseException:
        del sys.modules[module.__name__]
        raise
    return module


def engine(sdk_python: str | Path | None = None) -> ModuleType:
    """Return an independent trainer namespace with the unchanged public engine API.

    ``sdk_python`` is required for plan/run; artifact-only checkpoint validation
    needs neither a running service nor the SDK. No process-global patch occurs.
    The native Trainer classes remain upstream importable classes, and native
    state_dict checkpoints do not pickle these private orchestration namespaces.
    """
    legacy = _load_engine("gym_training.py")
    runtime = _load_engine("gym_training_long.py")

    def identity(url):
        if sdk_python is None:
            raise ValueError("OpenEnv plan/run requires sdk_python for the isolated native SDK")
        server = openenv_client.fetch_gym_identity(url)
        return server | {"client_sdk_runtime": openenv_client.sdk_identity(sdk_python)}

    def environment(url, identity):
        if sdk_python is None:
            raise ValueError("OpenEnv execution requires sdk_python")
        return openenv_client.RemoteGymEnv(url, identity, sdk_python=sdk_python,
                                          expected_sdk_identity=identity["client_sdk_runtime"])

    api = SimpleNamespace(fetch_gym_identity=identity, RemoteGymEnv=environment,
                          GymInfrastructureError=openenv_client.GymInfrastructureError)
    legacy._gym_api = lambda: api
    legacy._sources, legacy._source_commit = source_hashes, _source_commit
    original_legacy_plan = legacy.plan

    def backend_plan(config_path, gym_url):
        report = original_legacy_plan(config_path, gym_url)
        report.update(method="native_trl_remote_openenv_grpo", backend="openenv")
        return report

    legacy.plan = backend_plan
    runtime.legacy = legacy
    runtime.FORMAT, runtime._sources = FORMAT, source_hashes
    original_plan = runtime.plan

    def bound_plan(config_path, gym_url, *, checkpoint_interval=64):
        report = original_plan(config_path, gym_url, checkpoint_interval=checkpoint_interval)
        report.update(method="native_trl_remote_openenv_grpo_long", backend="openenv",
                      openenv={"transport": "native SDK persistent WebSocket reset/step/state",
                               "sdk_runtime": report["gym"]["identity"]["client_sdk_runtime"],
                               "sdk_python": str(Path(sdk_python).absolute()),
                               "upstream_revision": openenv_client.OPENENV_REVISION,
                               "model_execution": "native TRL in allocated CUDA process; SDK worker loads no model"},
                      engine_reuse={"strategy": "isolated source-verified dependency injection",
                                    "frozen_engine_sha256": dict(FROZEN_ENGINES),
                                    "unchanged": ["task schema", "tool schemas", "reward", "generation",
                                                  "token proof", "loss", "optimizer", "schedule", "resume"]})
        return report

    runtime.plan = bound_plan
    runtime.source_hashes = source_hashes
    return runtime


def plan(config_path: str | Path, gym_url: str, *, sdk_python: str | Path,
         checkpoint_interval: int = 64) -> dict:
    return engine(sdk_python).plan(config_path, gym_url, checkpoint_interval=checkpoint_interval)


def run(config_path: str | Path, gym_url: str, *, sdk_python: str | Path,
        resume_from_checkpoint: str | Path | None = None, stop_after_step: int | None = None,
        checkpoint_interval: int = 64) -> dict:
    return engine(sdk_python).run(config_path, gym_url, resume_from_checkpoint=resume_from_checkpoint,
                                stop_after_step=stop_after_step, checkpoint_interval=checkpoint_interval)


def validate_checkpoint(output_dir: str | Path, checkpoint: str | Path, *, expected_step: int | None = None) -> dict:
    return engine().validate_checkpoint(output_dir, checkpoint, expected_step=expected_step)


def discover_resume_checkpoint(output_dir: str | Path) -> Path | None:
    return engine().discover_resume_checkpoint(output_dir)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("config", type=Path)
    parser.add_argument("--gym-url", required=True)
    parser.add_argument("--sdk-python", required=True, type=Path)
    parser.add_argument("--checkpoint-interval", type=int, default=64)
    parser.add_argument("--resume-from-checkpoint", type=Path)
    parser.add_argument("--stop-after-step", type=int)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    result = (plan(args.config, args.gym_url, sdk_python=args.sdk_python,
                   checkpoint_interval=args.checkpoint_interval) if args.dry_run else
              run(args.config, args.gym_url, sdk_python=args.sdk_python,
                  checkpoint_interval=args.checkpoint_interval, resume_from_checkpoint=args.resume_from_checkpoint,
                  stop_after_step=args.stop_after_step))
    print(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False))


if __name__ == "__main__":
    main()