"""Qualify the actual trainer client against an owned native NeMo Gym service.

Run in the isolated Gym environment with a fresh evidence directory:
  .venv-nemo-gym/bin/python examples/agent_gyms/nemo_gym_client_smoke.py \
    --output runs/agent-gyms/nemo-client-smoke/new-run

This CPU integration test uses scripted TRAIN fixtures, not sampled policy
rollouts. It never changes the trainer client, environment, or model checkpoint.
"""
from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import socket
import subprocess
import sys
import threading
import time
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.gym_client import (  # noqa: E402
    GymInfrastructureError, RemoteGymEnv, fetch_gym_identity,
)


def _call(name: str, arguments: dict) -> dict:
    return {"role": "assistant", "content": None, "tool_calls": [{
        "type": "function", "function": {"name": name, "arguments": arguments},
    }]}


def _transcript(task: dict, prefix: list[dict] | None = None) -> list[dict]:
    facts = task["private_facts"]
    answer = str(facts["quantity"] * facts["unit_price"])
    return (prefix or []) + [
        _call("lookup", {"record_id": task["record_id"]}),
        {"role": "tool", "content": "copied transcript is not authoritative"},
        _call("calculator", {"expression": f"{facts['quantity']}*{facts['unit_price']}"}),
        {"role": "tool", "content": "copied transcript is not authoritative"},
        {"role": "assistant", "content": f"<answer>{answer}</answer>"},
    ]


def _check_service_loss(environment: RemoteGymEnv, row: dict) -> dict:
    """After actual server shutdown, verify failures cannot become scored mistakes."""
    try:
        environment.lookup(row["task"]["record_id"])
    except GymInfrastructureError as initial:
        assert environment._infrastructure_error is initial
        operations = {
            "verify": lambda: environment._score("<answer>0</answer>"),
            "calculator": lambda: environment.calculator("2*3"),
            "reset": lambda: environment.reset(**row),
            "close": environment._close,
        }
        for name, operation in operations.items():
            try:
                operation()
            except GymInfrastructureError as repeated:
                assert repeated is initial
            else:
                raise AssertionError(f"{name} swallowed the sticky service failure")
        return {"type": type(initial).__name__, "message": str(initial),
                "sticky_same_exception_for": list(operations), "scored_reward_returned": False}
    raise AssertionError("Service loss must raise an infrastructure error")


def run(output: Path, startup_timeout: float = 240.0) -> dict[str, Any]:
    output = output.resolve()
    if output.exists():
        raise ValueError("Use a fresh output directory; existing evidence is never overwritten")
    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}"
    server_path = ROOT / "examples/agent_gyms/nemo_gym_server.py"
    checks = []
    started = time.monotonic()
    with (output / "server.log").open("w") as log:
        process = subprocess.Popen([sys.executable, str(server_path), "--port", str(port)],
                                   cwd=ROOT, stdout=log, stderr=subprocess.STDOUT)
        try:
            while True:
                if process.poll() is not None:
                    raise RuntimeError(f"Native Gym server exited; inspect {output / 'server.log'}")
                try:
                    identity = fetch_gym_identity(base_url)
                    break
                except GymInfrastructureError:
                    if time.monotonic() - started > startup_timeout:
                        raise
                    time.sleep(0.2)
            checks.append("real_native_gym_identity_checked_by_frozen_client")

            rows = generate_tasks("train", 5, seed=54321)
            environments = [RemoteGymEnv(base_url, identity) for _ in rows]
            for environment, row in zip(environments, rows, strict=True):
                environment.reset(**row)
            assert len({env._session_id for env in environments}) == len(environments)
            checks.append("client_reset_creates_distinct_remote_episode_sessions")

            barrier = threading.Barrier(2)

            def worker(index: int) -> dict:
                environment, task = environments[index], rows[index]["task"]
                prefix = []
                barrier.wait(timeout=10)
                if index == 1:
                    wrong = {"record_id": rows[0]["task"]["record_id"]}
                    observation = json.loads(environment.lookup(**wrong))
                    assert observation == {"error": "unknown_record"}
                    assert environment._infrastructure_error is None
                    prefix = [_call("lookup", wrong), {"role": "tool", "content": json.dumps(observation)}]
                observation = json.loads(environment.lookup(task["record_id"]))
                assert observation["quantity"] == task["private_facts"]["quantity"]
                assert observation["unit_price"] == task["private_facts"]["unit_price"]
                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, prefix))
                assert score["reward"] == (1.0 if index == 0 else 0.0)
                assert environment._infrastructure_error is None
                return {"task_id": task["task_id"], "reward": score["reward"],
                        "success": score["success"], "reason": score["reason"],
                        "tool_calls": score["tool_calls"], "independent_replay_verified": True}

            with ThreadPoolExecutor(max_workers=2) as pool:
                results = list(pool.map(worker, [0, 1]))
            assert results[0]["success"] and results[1]["reason"] == "unknown_record"
            checks.extend([
                "two_concurrent_client_episodes_keep_hidden_state_isolated",
                "actual_tools_and_valid_transcript_score_strict_success",
                "wrong_record_remains_zero_scored_model_failure_not_infrastructure",
                "client_independent_replay_accepts_real_server_observations_and_scores",
            ])

            environment = environments[2]
            assert json.loads(environment.lookup(42)) == {"error": "unknown_record"}
            assert environment._infrastructure_error is None
            assert environment._score("<answer>0</answer>")["reward"] == 0
            environment, task = environments[3], rows[3]["task"]
            environment.lookup(task["record_id"])
            assert json.loads(environment.calculator("1" * 129))["error"] == "invalid_expression"
            assert environment._infrastructure_error is None
            assert environment._score("<answer>0</answer>")["reward"] == 0
            checks.append("malformed_model_values_remain_scored_errors_through_actual_client")

            for environment in environments[:4]:
                environment._close()
                assert environment._session_id is None
            checks.append("actual_client_cleanup_closes_completed_native_sessions")

            # Terminate only this qualification's owned server, then use a live
            # previously reset episode through the unchanged trainer client.
            process.terminate()
            process.wait(timeout=10)
            failure = _check_service_loss(environments[4], rows[4])
            checks.append("real_service_loss_is_sticky_infrastructure_failure_never_zero_reward")

            source_paths = [ROOT / "src/hessian_recipes/gym_client.py", server_path,
                            ROOT / "src/hessian_recipes/agent_rl_tasks.py", Path(__file__).resolve()]
            report = {
                "schema_version": 1, "status": "passed", "identity": identity,
                "checks": checks, "check_count": len(checks),
                "concurrent_episode_outcomes": results, "service_loss": failure,
                "source_sha256": {str(path.relative_to(ROOT)): hashlib.sha256(path.read_bytes()).hexdigest()
                                  for path in source_paths},
                "completed_at": datetime.now(timezone.utc).isoformat(),
                "wall_seconds": time.monotonic() - started,
                "limitations": "Real frozen RemoteGymEnv client against actual pinned NeMo Gym CPU service "
                               "with scripted TRAIN fixtures. No policy inference, GPU training, "
                               "sampled-token proof or quality claim.",
            }
            (output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
            return report
        finally:
            if process.poll() is None:
                process.terminate()
                try:
                    process.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait(timeout=10)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, required=True, help="Fresh evidence directory")
    args = parser.parse_args()
    print(json.dumps(run(args.output), indent=2))


if __name__ == "__main__":
    main()