"""Real pinned OpenEnv CPU qualification with a scripted policy, never a model.

The launcher is stdlib-only so a lightweight orchestration process can own the
native service in its separate Python ABI. Native fixture execution imports the
actual installed SDK and exercises its persistent WebSocket sessions.
"""
from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager, ExitStack
import hashlib
import importlib.metadata
import importlib.util
import json
import os
from pathlib import Path
import platform
import socket
import subprocess
import sys
import threading
import time
from urllib.request import build_opener, HTTPRedirectHandler, ProxyHandler, Request

ROOT = Path(__file__).resolve().parents[2]
SERVER = ROOT / "examples/agent_gyms/openenv_server.py"
REVISION = "17c1edd22168fcb76b086def1ed90a7168fb8749"
sys.path.insert(0, str(ROOT / "src"))


def _sha(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


class _NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, *args, **kwargs):
        raise RuntimeError("Loopback identity redirects are forbidden")


def _identity(url: str) -> tuple[dict, int]:
    opener = build_opener(ProxyHandler({}), _NoRedirect())
    with opener.open(Request(url + "/hessian_identity"), timeout=3) as response:
        value = json.loads(response.read(65_536))
        return value, int(response.headers["X-Hessian-Process-Id"])


@contextmanager
def launch_openenv_server(sdk_python: str | Path, run_dir: str | Path, *,
                         max_sessions: int = 256, startup_timeout: float = 120):
    """Yield an owned native loopback service and always reap only that child."""
    output = Path(run_dir)
    output.mkdir(parents=True, exist_ok=False)
    python = Path(sdk_python).absolute()  # Preserve the virtual-environment executable.
    with socket.socket() as reservation:
        reservation.bind(("127.0.0.1", 0))
        port = reservation.getsockname()[1]
    url = f"http://127.0.0.1:{port}"
    log_path = output / "server.log"
    environment = os.environ.copy()
    # The GPU parent may use a Python 3.11 site-packages overlay. Neither its
    # imports nor a foreign stdlib prefix may cross into this Python 3.13 ABI.
    environment.pop("PYTHONPATH", None)
    environment.pop("PYTHONHOME", None)
    environment.update(ENABLE_WEB_INTERFACE="false", OPENENV_MODE="simulation", CUDA_VISIBLE_DEVICES="")
    process = None
    receipt = {"url": url, "pid": None, "log_path": str(log_path), "released": False}
    with log_path.open("x") as log:
        try:
            process = subprocess.Popen([str(python), str(SERVER), "--port", str(port),
                                        "--max-sessions", str(max_sessions)],
                                       cwd=ROOT, env=environment, stdout=log, stderr=log)
            receipt["pid"] = process.pid
            deadline = time.monotonic() + startup_timeout
            while True:
                if process.poll() is not None:
                    raise RuntimeError("Native OpenEnv child exited; inspect its owned log")
                try:
                    value, process_id = _identity(url)
                    if process_id == process.pid:
                        break
                except (OSError, ValueError, KeyError):
                    pass
                if time.monotonic() >= deadline:
                    raise TimeoutError("Native OpenEnv child did not become ready")
                time.sleep(0.1)
            if (value.get("backend") != "openenv" or value.get("upstream_revision") != REVISION
                    or value.get("protocol") != "hessian-openenv-v1"
                    or value.get("server_sha256") != _sha(SERVER)
                    or value.get("environment_sha256") != _sha(ROOT / "src/hessian_recipes/agent_rl_tasks.py")
                    or value.get("requirements_sha256") != _sha(ROOT / "requirements/openenv.txt")
                    or value.get("runtime", {}).get("commit") != REVISION):
                raise RuntimeError("Native OpenEnv service identity differs from this source snapshot")
            receipt["identity"] = value
            yield receipt
        finally:
            if process is not None:
                if process.poll() is None:
                    process.terminate()
                    try:
                        process.wait(timeout=10)
                    except subprocess.TimeoutExpired:
                        process.kill()
                        process.wait(timeout=10)
                receipt.update(released=True, returncode=process.returncode)
            (output / "server-lifecycle.json").write_text(json.dumps(receipt, indent=2) + "\n")


def _native_module():
    spec = importlib.util.spec_from_file_location("hessian_native_openenv", SERVER)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


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


def _solve(client, row: dict) -> tuple[list[dict], dict]:
    reset = client.reset(task=row["task"])
    assert reset.observation["observation"] == "" and reset.observation["score"] is None
    transcript = []
    arguments = {"record_id": row["task"]["record_id"]}
    lookup = client.step({"operation": "tool", "tool_name": "lookup", "arguments": arguments})
    transcript.extend([_call("lookup", arguments), {"role": "tool", "content": lookup.observation["observation"]}])
    facts = json.loads(lookup.observation["observation"])
    arguments = {"expression": f'{facts["quantity"]}*{facts["unit_price"]}'}
    calculation = client.step({"operation": "tool", "tool_name": "calculator", "arguments": arguments})
    transcript.extend([_call("calculator", arguments), {"role": "tool", "content": calculation.observation["observation"]}])
    answer = json.loads(calculation.observation["observation"])["result"]
    transcript.append({"role": "assistant", "content": f"<answer>{answer}</answer>"})
    final = client.step({"operation": "verify", "completion": transcript})
    assert final.done is True and final.reward == 1.0 and final.observation["score"]["success"]
    return transcript, final.observation["score"]


def _wait_empty(url: str) -> None:
    deadline = time.monotonic() + 10
    while _identity(url)[0]["active_environment_instances"]:
        if time.monotonic() > deadline:
            raise AssertionError("Native environment instances leaked after client close")
        time.sleep(0.05)


@contextmanager
def _poisoned_parent_overlay(path: Path):
    path.mkdir()
    for name in ("json.py", "openenv.py"):
        (path / name).write_text('raise RuntimeError("Foreign Python overlay reached the native SDK child")\n')
    previous = {key: os.environ.get(key) for key in ("PYTHONPATH", "PYTHONHOME")}
    os.environ.update(PYTHONPATH=str(path), PYTHONHOME=str(path))
    try:
        yield
    finally:
        for key, value in previous.items():
            if value is None:
                os.environ.pop(key, None)
            else:
                os.environ[key] = value


def run(output: str | Path, sdk_python: str | Path = sys.executable) -> dict:
    from openenv.core.generic_client import GenericEnvClient
    from hessian_recipes.agent_rl_tasks import GermanAgentEnv, generate_tasks, score_episode

    destination = Path(output)
    destination.mkdir(parents=True, exist_ok=False)
    native = _native_module()
    rows = generate_tasks("train", 300, seed=20261020)
    checks, details = {}, {}
    started = time.monotonic()
    with _poisoned_parent_overlay(destination / "foreign-python-overlay"), \
            launch_openenv_server(sdk_python, destination / "service", max_sessions=32) as service:
        url = service["url"]
        checks["inherited_python_overlay_removed"] = True

        def client():
            return GenericEnvClient(base_url=url, connect_timeout_s=5, message_timeout_s=10).sync()

        with native.GermanOpenEnvClient(base_url=url).sync() as typed:
            result = typed.reset(task=rows[0]["task"])
            state = typed.state()
            assert result.observation.observation == "" and result.observation.score is None
            assert state.instance_id and state.episode_id == rows[0]["id"]
            visible = json.dumps(state.model_dump() | result.observation.model_dump())
            assert not any(key in visible for key in ("private_facts", "quantity", "unit_price", "expected_answer"))
        checks["typed_native_reset_state_no_hidden_facts"] = True

        with client() as first, client() as second:
            first.reset(task=rows[0]["task"])
            second.reset(task=rows[1]["task"])
            assert first.state()["instance_id"] != second.state()["instance_id"]
            wrong = first.step({"operation": "tool", "tool_name": "lookup", "arguments": {"record_id": rows[1]["task"]["record_id"]}})
            assert json.loads(wrong.observation["observation"])["error"] == "unknown_record"
            transcript, score = _solve(second, rows[1])
            local = GermanAgentEnv()
            local.reset(task=rows[1]["task"])
            for action in [message["tool_calls"][0]["function"] for message in transcript if message.get("tool_calls")]:
                getattr(local, action["name"])(**action["arguments"])
            assert score == score_episode(local, transcript)
        checks["two_native_sessions_isolated_and_strict_score_parity"] = True

        rejected = []
        cases = [("unknown", {}), ("lookup", {"record_id": rows[0]["task"]["record_id"], "extra": 1}),
                 ("lookup", {"tool": "calculator", "expression": "2*3"}),
                 ("lookup", {"record_id": 12}), ("calculator", {"expression": "2*3"})]
        for name, arguments in cases:
            with client() as active:
                active.reset(task=rows[0]["task"])
                result = active.step({"operation": "tool", "tool_name": name, "arguments": arguments})
                assert "error" in json.loads(result.observation["observation"])
                final = active.step({"operation": "verify", "completion": [_call(name, arguments), {"role": "assistant", "content": "<answer>6</answer>"}]})
                assert final.reward == 0.0 and final.done
                rejected.append(final.observation["score"]["reason"])
        checks["unknown_names_kwargs_reserved_keys_badtypes_scored_zero"] = True
        details["invalid_argument_reasons"] = rejected

        with client() as active:
            transcript, score = _solve(active, rows[2])
            repeated = active.step({"operation": "verify", "completion": transcript})
            assert repeated.observation["score"] == score
            for bad in [{"operation": "verify", "completion": "<answer>0</answer>"},
                        {"operation": "tool", "tool_name": "lookup", "arguments": {"record_id": rows[2]["task"]["record_id"]}}]:
                try:
                    active.step(bad)
                except Exception:
                    pass
                else:
                    raise AssertionError("Verified episode accepted a conflicting mutation")
            _solve(active, rows[3])
        checks["verification_idempotent_terminal_immutable_reset_reusable"] = True

        _wait_empty(url)
        all_ready = threading.Barrier(33)
        release_actions = threading.Event()

        def solve_index(index):
            with client() as active:
                # State proves the native factory has created the session. Hold
                # every connection open until all 32 exist at the same time.
                active.state()
                all_ready.wait(timeout=30)
                if not release_actions.wait(timeout=30):
                    raise TimeoutError("Concurrent native fixture release timed out")
                _, score = _solve(active, rows[index])
                return active.state()["instance_id"], score["reward"]
        with ThreadPoolExecutor(max_workers=32) as pool:
            futures = [pool.submit(solve_index, index) for index in range(32)]
            try:
                all_ready.wait(timeout=30)
                simultaneously_open = _identity(url)[0]["active_environment_instances"]
                assert simultaneously_open == 32
            finally:
                release_actions.set()
            concurrent = [future.result() for future in futures]
        assert len({item[0] for item in concurrent}) == 32
        assert all(item[1] == 1.0 for item in concurrent)
        checks["32_concurrent_native_sessions_unique_and_correct"] = True
        details["concurrent_sessions"] = {"episodes": 32, "workers": 32,
                                          "simultaneously_open_before_actions": simultaneously_open,
                                          "all_native_states_received_before_release": True}

        _wait_empty(url)
        with ExitStack() as sessions:
            held = [sessions.enter_context(client()) for _ in range(32)]
            assert len({active.state()["instance_id"] for active in held}) == 32
            try:
                with client() as excess:
                    excess.state()
            except Exception as error:
                assert "capacity" in str(error).lower() or "session" in str(error).lower()
                details["capacity_error_type"] = type(error).__name__
            else:
                raise AssertionError("Native capacity bound did not reject the 33rd session")
            held.pop().close()
            deadline = time.monotonic() + 10
            while _identity(url)[0]["active_environment_instances"] != 31:
                if time.monotonic() > deadline:
                    raise AssertionError("Native capacity did not recover after session close")
                time.sleep(0.05)
            with client() as recovered:
                _solve(recovered, rows[299])
        checks["native_capacity_rejection_and_recovery"] = True
        details["capacity"] = {"limit": 32, "rejected_session": 33,
                                "active_after_one_close": 31, "replacement_completed": True}

        from websockets.sync.client import connect
        with connect(url.replace("http://", "ws://") + "/ws", proxy=None) as raw:
            for message in ['{"type":"state","type":"reset"}',
                            '{"type":"state","data":NaN}',
                            '{"type":"state","data":Infinity}']:
                raw.send(message)
                rejected = json.loads(raw.recv())
                assert rejected["type"] == "error" and rejected["data"]["code"] == "INVALID_JSON"
                raw.send(json.dumps({"type": "state"}))
                state = json.loads(raw.recv())["data"]
                assert state["episode_id"] is None and state["step_count"] == 0
        checks["native_wire_duplicate_keys_and_nonfinite_json_rejected"] = True

        with client() as active:
            active.reset(task=rows[0]["task"], episode_id="before-invalid-reset")
            try:
                active.reset(task={"schema_version": 1})
            except Exception:
                pass
            else:
                raise AssertionError("Invalid controller reset was accepted")
            assert active.state()["episode_id"] is None
            try:
                active.step({"operation": "tool", "tool_name": "lookup", "arguments": {"record_id": rows[0]["task"]["record_id"]}})
            except Exception:
                pass
            else:
                raise AssertionError("Failed reset retained previous private facts")
            _solve(active, rows[1])
        checks["failed_controller_reset_discards_prior_task_and_recovers"] = True

        with client() as active:
            initial = active.state()["instance_id"]
            for row in rows[32:292]:
                _solve(active, row)
                assert active.state()["instance_id"] == initial
        checks["260_sequential_resets_no_session_limit_leak"] = True
        _wait_empty(url)
        checks["native_close_releases_all_environment_instances"] = True
        frozen_identity = _identity(url)[0]
    checks["owned_service_process_reaped"] = service["released"]
    report = {"schema_version": 1, "status": "passed", "backend": "openenv",
              "mode": "native-sdk-cpu-scripted-fixtures", "native_sdk_executed": True,
              "model_inference": False, "gpu": False, "strict_reward_unchanged": True,
              "server_released": service["released"], "native_identity": frozen_identity,
              "checks": checks, "details": details, "seconds": time.monotonic() - started,
              "python_version": platform.python_version(),
              "source_sha256": {"server": _sha(SERVER), "smoke": _sha(Path(__file__)),
                                "tests": _sha(ROOT / "tests/test_openenv_server.py"),
                                "requirements": _sha(ROOT / "requirements/openenv.txt")},
              "upstream": {"repository": "https://github.com/huggingface/OpenEnv",
                           "revision": REVISION, "version": importlib.metadata.version("openenv"),
                           "pyproject_sha256": (_sha(ROOT / ".cache/upstreams/openenv/pyproject.toml")
                                                if (ROOT / ".cache/upstreams/openenv/pyproject.toml").is_file() else None)},
              "dependencies": dict(sorted((item.metadata["Name"], item.version)
                                          for item in importlib.metadata.distributions())),
              "limits": ["Scripted TRAIN fixtures qualify native environment wiring, not model ability or RL.",
                         "This narrow synthetic task and strict reward are unchanged; no held-out task or MGSM assessment."]}
    (destination / "report.json").write_text(json.dumps(report, indent=2) + "\n")
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--sdk-python", type=Path, default=Path(sys.executable))
    args = parser.parse_args()
    report = run(args.output, args.sdk_python)
    print(json.dumps({key: report[key] for key in ("status", "backend", "checks", "seconds", "server_released")}, indent=2))


if __name__ == "__main__":
    main()