"""CPU stress qualification of the frozen Gym service, client and native dispatch.

Use the isolated NeMo Python and a fresh output directory. Optionally provide
--native-python .venv-gpu/bin/python to exercise TRL's actual pinned tool-loop
method with scripted responses, without loading a model or allocating a GPU.
No fixture tokens from this check are presented as model captures.
"""
from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor
import copy
from datetime import datetime, timezone
import hashlib
import inspect
import json
from pathlib import Path
import socket
import subprocess
import sys
import time
from types import SimpleNamespace
from typing import Any

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from hessian_recipes.agent_rl_tasks import generate_tasks  # noqa: E402
from hessian_recipes.agent_rl_training import milestone_training_reward  # noqa: E402
from hessian_recipes.gym_client import (  # noqa: E402
    GymInfrastructureError, RemoteGymEnv, fetch_gym_identity,
)
from nemo_gym_client_smoke import _transcript  # noqa: E402


def _stop(process: subprocess.Popen) -> None:
    if process.poll() is None:
        process.terminate()
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=10)


def _start(port: int, log, *, timeout: float = 240.0) -> tuple[subprocess.Popen, dict]:
    process = subprocess.Popen(
        [sys.executable, str(ROOT / "examples/agent_gyms/nemo_gym_server.py"), "--port", str(port)],
        cwd=ROOT, stdout=log, stderr=subprocess.STDOUT,
    )
    started = time.monotonic()
    try:
        while process.poll() is None:
            try:
                return process, fetch_gym_identity(f"http://127.0.0.1:{port}")
            except GymInfrastructureError:
                if time.monotonic() - started > timeout:
                    raise
                time.sleep(0.1)
        raise RuntimeError("Native Gym server exited during stress startup")
    except BaseException:
        _stop(process)
        raise


def _solve(environment: RemoteGymEnv, row: dict) -> dict:
    task = row["task"]
    observation = json.loads(environment.lookup(task["record_id"]))
    expression = f"{observation['quantity']}*{observation['unit_price']}"
    result = json.loads(environment.calculator(expression))
    assert result["result"] == str(observation["quantity"] * observation["unit_price"])
    score = environment._score(_transcript(task))
    assert score["success"] and score["reward"] == 1
    return score


def native_dispatch(base_url: str) -> dict:
    """Execute the actual upstream tool-loop method using a scripted CPU processor."""
    import importlib.metadata
    from trl import GRPOTrainer

    assert importlib.metadata.version("trl") == "0.29.0"
    identity = fetch_gym_identity(base_url)
    cases = []
    for index, mode in enumerate(("unknown_name", "unexpected_keyword", "unexecuted_at_cap")):
        row = generate_tasks("train", 1, seed=9200 + index)[0]
        environment = RemoteGymEnv(base_url, identity)
        environment.reset(**row)

        def assistant(name=None, arguments=None, answer=None):
            if name is None:
                text = f"<answer>{answer}</answer>"
                return {"role": "assistant", "content": text, "raw_response": text}
            payload = {"name": name, "arguments": arguments}
            text = "<tool_call>" + json.dumps(payload) + "</tool_call>"
            return {"role": "assistant", "content": "", "raw_response": text,
                    "tool_calls": [{"type": "function", "function": payload}]}

        parsed = {}

        class ScriptedProcessor:
            eos_token = "<fixture-eos>"

            def apply_chat_template(self, batches, **kwargs):
                result = []
                for messages in batches:
                    text = ""
                    for message in messages:
                        if message["role"] == "assistant":
                            text += "<assistant>" + message["raw_response"] + "<end>"
                        else:
                            text += "<" + message["role"] + ">" + message["content"] + "<end>"
                    result.append(list((text + "<assistant>").encode()))
                return result

            def parse_response(self, ids):
                return copy.deepcopy(parsed[bytes(ids).decode()])

        processor = ScriptedProcessor()
        stage, answer = 0, None

        def generate(prompts):
            nonlocal stage, answer
            stage += 1
            if stage == 1:
                values = json.loads(prompts[0][-1]["content"])
                message = assistant("calculator", {"expression": f"{values['quantity']}*{values['unit_price']}"})
            elif stage == 2:
                answer = json.loads(prompts[0][-1]["content"])["result"]
                if mode == "unexpected_keyword":
                    message = assistant("lookup", {"record_id": row["task"]["record_id"], "unexpected": "value"})
                else:
                    message = assistant("unknown_tool", {})
            else:
                message = assistant(answer=answer)
            parsed[message["raw_response"]] = message
            return processor.apply_chat_template(prompts), [list(message["raw_response"].encode())], None, {}

        harness = SimpleNamespace(
            max_tool_calling_iterations=2 if mode == "unexecuted_at_cap" else 3,
            max_completion_length=10000, use_vllm=False, chat_template=None,
            chat_template_kwargs={}, tools=[], processing_class=processor,
            model=SimpleNamespace(config=SimpleNamespace(max_position_embeddings=32768)),
            _sync_tool_dicts=[{"lookup": environment.lookup, "calculator": environment.calculator}],
            _async_tool_dicts=[{}], _generate_single_turn=generate,
        )
        prompt = [[{"role": "user", "content": row["prompt"]}]]
        first = assistant("lookup", {"record_id": row["task"]["record_id"]})
        result = GRPOTrainer._tool_call_loop(
            harness, prompt, processor.apply_chat_template(prompt),
            [list(first["raw_response"].encode())], [[first]], None,
        )
        completions = result[1][0]
        score = environment._score(completions)
        components = milestone_training_reward(score)
        tool_results = [message for message in completions if message["role"] == "tool"]
        native_only = len(tool_results) - len(environment._calls)
        expected_failures = int(mode != "unexecuted_at_cap")
        assert result[-1] == native_only == expected_failures
        assert score["reward"] == 0 and components["training_reward"] == 0.4
        assert components["trusted_tool_errors"] == 0
        cases.append({"case": mode, "declared_calls": 3, "actual_rpc_calls": len(environment._calls),
                      "tool_result_messages": len(tool_results), "native_dispatch_failures": result[-1],
                      "native_only_error_results": native_only, "strict_success": score["success"],
                      "trusted_tool_errors": components["trusted_tool_errors"],
                      "partial_training_reward": components["training_reward"]})
        environment._close()
    source = Path(inspect.getfile(GRPOTrainer))
    return {"status": "passed", "trl": "0.29.0",
            "upstream_trainer_sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
            "cases": cases, "model_inference": False,
            "limitation": "Actual upstream tool-loop Python dispatch with a scripted byte processor; no model, optimizer, or sampled-token qualification."}


def run(output: Path, *, native_python: Path | None = None, reuse_count: int = 260) -> dict[str, Any]:
    output = output.resolve()
    if output.exists():
        raise ValueError("Use a fresh output directory")
    output.mkdir(parents=True)
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        port = sock.getsockname()[1]
    base_url = f"http://127.0.0.1:{port}"
    checks, detail = [], {}
    started = time.monotonic()
    with (output / "server.log").open("w") as log:
        process, identity = _start(port, log)
        try:
            rows = generate_tasks("train", max(reuse_count, 256), seed=20261020)

            def worker(index):
                environment = RemoteGymEnv(base_url, identity)
                try:
                    environment.reset(**rows[index])
                    _solve(environment, rows[index])
                    return environment._session_id
                finally:
                    environment._close()

            with ThreadPoolExecutor(max_workers=8) as pool:
                identifiers = list(pool.map(worker, range(32)))
            assert len(set(identifiers)) == 32
            checks.append("32_concurrent_isolated_real_client_episodes_verify_and_close")

            reused = RemoteGymEnv(base_url, identity)
            identifiers = set()
            try:
                for row in rows[:reuse_count]:
                    reused.reset(**row)
                    identifiers.add(reused._session_id)
                    _solve(reused, row)
            finally:
                reused._close()
            assert len(identifiers) == reuse_count
            detail["sequential_reused_environment_episodes"] = reuse_count
            checks.append("environment_reset_reuse_beyond_server_capacity_has_no_active_session_leak")

            active = []
            try:
                for row in rows[:256]:
                    environment = RemoteGymEnv(base_url, identity)
                    environment.reset(**row)
                    active.append(environment)
                overflow = RemoteGymEnv(base_url, identity)
                try:
                    overflow.reset(**rows[0])
                except GymInfrastructureError:
                    assert overflow._infrastructure_error is not None
                else:
                    raise AssertionError("257th active session bypassed the service limit")
            finally:
                for environment in active:
                    environment._close()
            healthy = RemoteGymEnv(base_url, identity)
            healthy.reset(**rows[0])
            _solve(healthy, rows[0])
            healthy._close()
            checks.append("256_active_session_limit_fails_closed_then_recovers_after_cleanup")

            if native_python is not None:
                # Preserve the venv launcher path; resolving its symlink selects
                # the base interpreter and silently discards the venv packages.
                command = [str(native_python.absolute()), str(Path(__file__).resolve()),
                           "--native-dispatch-url", base_url, "--output", str(output / "native-dispatch.json")]
                with (output / "native-dispatch.log").open("w") as dispatch_log:
                    subprocess.run(command, cwd=ROOT, stdout=dispatch_log, stderr=subprocess.STDOUT,
                                   check=True, timeout=240)
                detail["native_dispatch"] = json.loads((output / "native-dispatch.json").read_text())
                checks.append("actual_TRL_unknown_name_kwargs_and_unexecuted_cap_diagnostics")

            before_restart = RemoteGymEnv(base_url, identity)
            before_restart.reset(**rows[1])
            before_restart.lookup(rows[1]["task"]["record_id"])
            _stop(process)
            process, restarted_identity = _start(port, log)
            assert restarted_identity == identity
            try:
                before_restart.calculator("2*3")
            except GymInfrastructureError as primary:
                for operation in (lambda: before_restart._score("<answer>6</answer>"),
                                  lambda: before_restart.reset(**rows[1])):
                    try:
                        operation()
                    except GymInfrastructureError as repeated:
                        assert repeated is primary
                    else:
                        raise AssertionError("Lost session became recoverable after service restart")
            else:
                raise AssertionError("Restarted service restored nonexistent session state")
            before_restart._close()
            fresh = RemoteGymEnv(base_url, restarted_identity)
            fresh.reset(**rows[2])
            _solve(fresh, rows[2])
            fresh._close()
            checks.append("real_service_restart_loses_old_session_stickily_and_allows_new_independent_episode")

            source_paths = [Path(__file__).resolve(), ROOT / "examples/agent_gyms/nemo_gym_server.py",
                            ROOT / "examples/agent_gyms/nemo_gym_client_smoke.py",
                            ROOT / "src/hessian_recipes/gym_client.py", ROOT / "src/hessian_recipes/agent_rl_tasks.py",
                            ROOT / "src/hessian_recipes/agent_rl_training.py"]
            report = {"status": "passed", "identity": identity, "checks": checks,
                      "check_count": len(checks), "details": detail,
                      "completed_at": datetime.now(timezone.utc).isoformat(),
                      "wall_seconds": time.monotonic() - started,
                      "source_sha256": {str(path.relative_to(ROOT)): hashlib.sha256(path.read_bytes()).hexdigest()
                                        for path in source_paths},
                      "limitations": "CPU infrastructure and adversarial protocol qualification using scripted TRAIN fixtures. No GPU allocation, model inference, learning, heldout score, fluency or general-agent claim."}
            (output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
            return report
        finally:
            _stop(process)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--native-python", type=Path)
    parser.add_argument("--reuse-count", type=int, default=260)
    parser.add_argument("--native-dispatch-url", help=argparse.SUPPRESS)
    args = parser.parse_args()
    if args.native_dispatch_url:
        report = native_dispatch(args.native_dispatch_url)
        args.output.write_text(json.dumps(report, indent=2) + "\n")
    else:
        if args.reuse_count < 257:
            parser.error("reuse-count must exceed the 256-session capacity")
        report = run(args.output, native_python=args.native_python, reuse_count=args.reuse_count)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()