"""CPU integration smoke against an owned, real pinned NeMo Gym HTTP process.

Run with .venv-nemo-gym/bin/python examples/agent_gyms/nemo_gym_smoke.py.
Fixtures exercise the server, session isolation and strict verifier. They are
not sampled policy rollouts or evidence that a model learned the tool workflow.
"""
from __future__ import annotations

import argparse
from datetime import datetime, timezone
import hashlib
import importlib.metadata
import json
from pathlib import Path
import socket
import subprocess
import sys
import time
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from hessian_recipes.agent_rl_tasks import generate_tasks  # noqa: E402

PIN = "720b5f318caee339bfd29c554c59a97edcedc846"


def _request(base: str, route: str, body: dict | None = None) -> tuple[int, dict]:
    data = None if body is None else json.dumps(body).encode()
    request = Request(base + route, data=data, headers={"Content-Type": "application/json"})
    try:
        with urlopen(request, timeout=10) as response:
            return response.status, json.load(response)
    except HTTPError as error:
        return error.code, json.load(error)


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


def _completion(task: dict, answer: str | None = None) -> list[dict]:
    facts = task["private_facts"]
    value = str(facts["quantity"] * facts["unit_price"])
    return [_call("lookup", {"record_id": task["record_id"]}),
            {"role": "tool", "name": "lookup", "content": "transcript copy is not authoritative"},
            _call("calculator", {"expression": f"{facts['quantity']}*{facts['unit_price']}"}),
            {"role": "tool", "name": "calculator", "content": "transcript copy is not authoritative"},
            {"role": "assistant", "content": f"<answer>{answer if answer is not None else value}</answer>"}]


def _native_response() -> dict:
    # A declared CPU fixture envelope, not a claim of model inference. The strict
    # verifier scores the separately supplied exact external-policy transcript.
    return {"id": "cpu-fixture", "created_at": 0.0, "error": None,
            "incomplete_details": None, "instructions": None, "metadata": None,
            "model": "cpu-fixture", "object": "response", "output": [],
            "parallel_tool_calls": False, "temperature": 0.0, "tool_choice": "auto",
            "tools": [], "top_p": 1.0, "status": "completed"}


def run(output: Path, startup_timeout: float = 240.0) -> dict:
    output = output.resolve()
    output.mkdir(parents=True, exist_ok=True)
    if (output / "report.json").exists():
        raise ValueError("Use a fresh output directory; completed smoke evidence is immutable")
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        port = sock.getsockname()[1]
    base = f"http://127.0.0.1:{port}"
    server_path = ROOT / "examples/agent_gyms/nemo_gym_server.py"
    started = time.monotonic()
    checks = []
    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"NeMo Gym server exited during startup; inspect {output / 'server.log'}")
                try:
                    status, identity = _request(base, "/hessian_identity")
                    if status == 200:
                        break
                except (URLError, TimeoutError, ConnectionError):
                    pass
                if time.monotonic() - started > startup_timeout:
                    raise TimeoutError("NeMo Gym startup timeout")
                time.sleep(0.2)
            assert identity["backend"] == "nemo_gym" and identity["upstream_revision"] == PIN
            assert identity["protocol"] == "hessian-gym-v1" and identity["inference"] is False
            assert identity["environment_sha256"] == hashlib.sha256(
                (ROOT / "src/hessian_recipes/agent_rl_tasks.py").read_bytes()).hexdigest()
            checks.append("actual_pinned_upstream_identity_and_frozen_environment")

            tasks = [row["task"] for row in generate_tasks("train", 7, seed=31415)]
            sessions = []
            for task in tasks:
                status, result = _request(base, "/reset_episode", {"task": task})
                assert status == 200
                sessions.append(result["session_id"])
            assert len(set(sessions)) == len(tasks)
            checks.append("fresh_opaque_isolated_sessions")

            def tools(index: int):
                task, session_id = tasks[index], sessions[index]
                status, result = _request(base, "/lookup", {"session_id": session_id,
                                                           "record_id": task["record_id"]})
                assert status == 200 and isinstance(result["observation"], str)
                observation = json.loads(result["observation"])
                assert observation["quantity"] == task["private_facts"]["quantity"]
                expression = f"{observation['quantity']}*{observation['unit_price']}"
                status, result = _request(base, "/calculator", {"session_id": session_id,
                                                               "expression": expression})
                assert status == 200
                assert json.loads(result["observation"])["result"] == str(
                    observation["quantity"] * observation["unit_price"])

            # A wrong record in B cannot expose A's facts or invalidate A's state.
            status, result = _request(base, "/lookup", {"session_id": sessions[1],
                                                       "record_id": tasks[0]["record_id"]})
            assert status == 200 and json.loads(result["observation"]) == {"error": "unknown_record"}
            tools(0)
            payload = {"session_id": sessions[0], "completion": _completion(tasks[0])}
            status, result = _request(base, "/verify_episode", payload)
            assert status == 200 and result["score"]["reward"] == 1.0
            assert result["score"]["tool_calls"] == 2
            checks.append("legal_real_lookup_calculator_and_strict_terminal_success")
            status, repeated = _request(base, "/verify_episode", payload)
            assert status == 200 and repeated == result
            status, _ = _request(base, "/verify_episode", payload | {"completion": _completion(tasks[0], "-1")})
            assert status == 409
            status, _ = _request(base, "/lookup", {"session_id": sessions[0],
                                                  "record_id": tasks[0]["record_id"]})
            assert status == 409
            checks.append("verification_is_idempotent_and_freezes_episode")

            tools(1)
            status, result = _request(base, "/verify_episode", {
                "session_id": sessions[1], "completion": _completion(tasks[1])})
            assert status == 200 and result["score"]["reward"] == 0.0
            assert result["score"]["reason"] == "unknown_record"
            checks.append("session_isolation_and_strict_invalid_workflow_rejection")

            tools(2)
            status, result = _request(base, "/verify_episode", {
                "session_id": sessions[2], "completion": _completion(tasks[2], "-1")})
            assert status == 200 and result["score"]["reward"] == 0.0
            checks.append("wrong_final_rejected_after_real_legal_tools")

            status, result = _request(base, "/verify_episode", {
                "session_id": sessions[3], "completion": _completion(tasks[3])})
            assert status == 200 and result["score"]["reward"] == 0.0
            assert result["score"]["trace"] == []
            checks.append("forged_transcript_cannot_replace_real_tool_execution")

            tools(4)
            status, result = _request(base, "/verify", {
                "session_id": sessions[4], "completion": _completion(tasks[4]),
                "responses_create_params": {"input": "CPU fixture; no model was invoked"},
                "response": _native_response()})
            assert status == 200, (status, result)
            assert result["reward"] == result["score"]["reward"] == 1.0
            assert result["mask_sample"] is False
            checks.append("native_gym_BaseVerifyRequest_Response_route")

            status, schema = _request(base, "/openapi.json")
            assert status == 200
            assert {"/seed_session", "/verify", "/close_session", "/aggregate_metrics"} <= set(schema["paths"])
            assert "NativeVerifyRequest" in schema["components"]["schemas"]
            checks.append("native_gym_routes_and_typed_verification_schema")

            status, result = _request(base, "/lookup", {"session_id": sessions[5], "record_id": 42})
            assert status == 200 and json.loads(result["observation"])["error"] == "unknown_record"
            status, result = _request(base, "/verify_episode", {
                "session_id": sessions[5], "completion": "<answer>0</answer>"})
            assert status == 200 and result["score"]["reward"] == 0.0
            assert result["score"]["trace"][0]["ok"] is False
            for expression in (42, "1" * 129):
                # One session may make multiple invalid tool attempts, all scored
                # by the same authoritative bounded frozen environment.
                status, result = _request(base, "/lookup", {
                    "session_id": sessions[6], "record_id": tasks[6]["record_id"]})
                assert status == 200
                status, result = _request(base, "/calculator", {
                    "session_id": sessions[6], "expression": expression})
                assert status == 200 and json.loads(result["observation"])["error"] == "invalid_expression"
            status, result = _request(base, "/verify_episode", {
                "session_id": sessions[6], "completion": "<answer>0</answer>"})
            assert status == 200 and result["score"]["reward"] == 0.0
            checks.append("bad_model_argument_types_and_long_expressions_are_scored_tool_errors")

            for session_id in sessions:
                for _ in range(2):
                    status, result = _request(base, "/close_episode", {"session_id": session_id})
                    assert status == 200 and result == {"closed": True}
            status, _ = _request(base, "/lookup", {"session_id": sessions[0],
                                                  "record_id": tasks[0]["record_id"]})
            assert status == 404
            checks.append("close_is_idempotent_and_removes_session_state")

            report = {"schema_version": 1, "status": "passed", "identity": identity,
                      "checks": checks, "check_count": len(checks),
                      "completed_at": datetime.now(timezone.utc).isoformat(),
                      "wall_seconds": time.monotonic() - started,
                      "smoke_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                      "dependencies": {name: importlib.metadata.version(name) for name in (
                          "nemo-gym", "fastapi", "uvicorn", "pydantic", "openai", "httpx")},
                      "limitations": "Real upstream CPU HTTP resources/verifier integration with scripted fixtures; no GPU inference, RL update, token-capture or policy-quality claim."}
            (output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
            return report
        finally:
            # Manage only the child process created by this smoke.
            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, default=ROOT / "runs/agent-gyms/nemo-native-smoke")
    args = parser.parse_args()
    print(json.dumps(run(args.output), indent=2))


if __name__ == "__main__":
    main()