"""Native OpenEnv WebSocket sessions for the frozen German TRL tool contract.

The isolated SDK process never loads a policy. It owns a persistent native SDK
client; the GPU process retains the original assistant tokens and tool masks.
Transport/state failures are sticky infrastructure errors, never task rewards.
"""
from __future__ import annotations

import argparse
import contextlib
import copy
import hashlib
import importlib.metadata
import json
import os
from pathlib import Path
import platform
import selectors
import subprocess
import sys
import time
from typing import Any
import uuid

from . import agent_rl_tasks
from . import gym_client as _http

OPENENV_REVISION = "17c1edd22168fcb76b086def1ed90a7168fb8749"
PROTOCOL = "hessian-openenv-v1"
MAX_BYTES = 2_000_000
TIMEOUT = 60


class GymInfrastructureError(RuntimeError):
    """Native OpenEnv transport, source identity or independent replay failed."""


def _json(raw: bytes | str):
    return json.loads(raw, object_pairs_hook=_http._unique_object,
                      parse_constant=lambda value: (_ for _ in ()).throw(ValueError(value)))


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


def _repo() -> Path:
    return Path(__file__).resolve().parents[2]


def fetch_gym_identity(base_url: str) -> dict:
    """Read-only metadata check; episode execution exclusively uses native /ws."""
    try:
        identity = _http._request(_http._base_url(base_url), "/hessian_identity")
    except _http.GymInfrastructureError as error:
        raise GymInfrastructureError("OpenEnv identity request failed") from error
    expected = {"protocol": PROTOCOL, "backend": "openenv", "upstream_revision": OPENENV_REVISION,
                "environment_sha256": _sha(Path(agent_rl_tasks.__file__)),
                "server_sha256": _sha(_repo() / "examples/agent_gyms/openenv_server.py"),
                "requirements_sha256": _sha(_repo() / "requirements/openenv.txt")}
    for name, value in expected.items():
        if identity.get(name) != value:
            raise GymInfrastructureError(f"OpenEnv identity mismatch for {name}")
    # Session counts are diagnostics, not immutable service/source identity.
    return {key: value for key, value in identity.items() if key != "active_environment_instances"}


def _sdk_identity() -> dict:
    """Hash installed SDK bytes, separately from its declared upstream revision."""
    import openenv
    distribution = importlib.metadata.distribution("openenv")
    root = Path(openenv.__file__).resolve().parent
    sources = {str(path.relative_to(root)): _sha(path) for path in sorted(root.rglob("*.py"))}
    direct = _json(distribution.read_text("direct_url.json") or "{}")
    if (distribution.version != "0.9.0"
            or direct.get("url") != "https://github.com/huggingface/OpenEnv.git"
            or direct.get("vcs_info", {}).get("commit_id") != OPENENV_REVISION):
        raise RuntimeError("SDK must be installed from the pinned OpenEnv Git revision")
    installed = sorted((item.metadata["Name"], item.version) for item in importlib.metadata.distributions())
    return {"distribution": "openenv", "version": distribution.version,
            "direct_url": direct, "installed_python_sources_sha256": sources,
            "installed_distributions_sha256": hashlib.sha256(json.dumps(installed).encode()).hexdigest(),
            "python_version": platform.python_version(),
            "python_implementation": platform.python_implementation()}


class _SDKProcess:
    """One bounded JSON channel and one native persistent WebSocket client."""

    def __init__(self, base_url: str, sdk_python: str | Path):
        self._url = _http._base_url(base_url)
        executable = Path(sdk_python).absolute()  # Preserve venv symlinks.
        if not executable.is_file():
            raise ValueError("sdk_python must name an existing isolated Python executable")
        environment = os.environ.copy()
        environment.pop("PYTHONHOME", None)
        environment.update(PYTHONPATH=str(_repo() / "src"), CUDA_VISIBLE_DEVICES="",
                           HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", PYTHONUNBUFFERED="1")
        self._process = subprocess.Popen(
            [str(executable), "-u", "-m", "hessian_recipes.openenv_client", "--sdk-worker", "--url", self._url],
            stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, env=environment,
            bufsize=0)
        self._sequence, self._buffer, self._closed = 0, b"", False

    def _call(self, operation: str, **arguments) -> dict:
        if self._closed:
            raise GymInfrastructureError("OpenEnv SDK channel is closed")
        self._sequence += 1
        payload = json.dumps({"id": self._sequence, "operation": operation, **arguments},
                             ensure_ascii=False, allow_nan=False).encode() + b"\n"
        if len(payload) > MAX_BYTES:
            raise GymInfrastructureError("OpenEnv request exceeds the bounded channel size")
        try:
            self._process.stdin.write(payload)
            self._process.stdin.flush()
            deadline = time.monotonic() + TIMEOUT
            with selectors.DefaultSelector() as selector:
                selector.register(self._process.stdout, selectors.EVENT_READ)
                while b"\n" not in self._buffer:
                    remaining = deadline - time.monotonic()
                    if remaining <= 0 or not selector.select(remaining):
                        raise TimeoutError
                    block = os.read(self._process.stdout.fileno(), 65536)
                    if not block:
                        raise EOFError
                    self._buffer += block
                    if len(self._buffer) > MAX_BYTES:
                        raise ValueError("Oversized SDK response")
            line, self._buffer = self._buffer.split(b"\n", 1)
            result = _json(line)
            if (not isinstance(result, dict) or result.get("id") != self._sequence
                    or result.get("ok") is not True or not isinstance(result.get("result"), dict)):
                raise ValueError("Invalid SDK response")
            return result["result"]
        except (OSError, EOFError, TimeoutError, ValueError, TypeError) as error:
            raise GymInfrastructureError(f"OpenEnv SDK channel failed during {operation}: {type(error).__name__}") from error

    def _close(self) -> None:
        if self._closed:
            return
        failure = None
        try:
            self._call("close")
        except GymInfrastructureError as error:
            failure = error
        finally:
            self._closed = True
            self._process.stdin.close()
            try:
                self._process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self._process.kill()
                self._process.wait(timeout=5)
            self._process.stdout.close()
        if failure:
            raise failure


def sdk_identity(sdk_python: str | Path) -> dict:
    """Probe installed SDK bytes without connecting to an episode or using GPUs."""
    process = _SDKProcess("http://127.0.0.1:1", sdk_python)
    try:
        return process._call("identity")
    finally:
        process._close()


class RemoteGymEnv:
    """Native OpenEnv-backed environment with the exact original model tool API."""

    def __init__(self, base_url: str, identity: dict, *, sdk_python: str | Path,
                 expected_sdk_identity: dict):
        self._url, self._python = _http._base_url(base_url), sdk_python
        self._provenance, self._expected_sdk = copy.deepcopy(identity), copy.deepcopy(expected_sdk_identity)
        self._process = None
        self._task, self._session_id, self._instance_id = None, None, None
        self._calls, self._infrastructure_error = [], None

    def _raise_if_failed(self) -> None:
        if self._infrastructure_error is not None:
            raise self._infrastructure_error

    def _fail(self, message: str) -> None:
        self._infrastructure_error = GymInfrastructureError(message)
        raise self._infrastructure_error

    def _rpc(self, operation: str, **arguments) -> dict:
        self._raise_if_failed()
        try:
            result = self._process._call(operation, **arguments)
            if not isinstance(result, dict):
                self._fail("OpenEnv SDK returned an invalid result envelope")
            return result
        except GymInfrastructureError as error:
            self._infrastructure_error = error
            raise

    def _check_state(self, state: dict, *, reset: bool = False) -> None:
        if (not isinstance(state, dict) or not isinstance(state.get("instance_id"), str)
                or not state["instance_id"] or not isinstance(state.get("episode_id"), str)
                or not state["episode_id"]):
            self._fail("OpenEnv returned invalid native session state")
        if self._instance_id is not None and state["instance_id"] != self._instance_id:
            self._fail("OpenEnv native WebSocket session changed unexpectedly")
        if not reset and self._session_id != state["episode_id"]:
            self._fail("OpenEnv native episode identity changed unexpectedly")
        self._instance_id, self._session_id = state["instance_id"], state["episode_id"]

    def reset(self, task: dict | None = None, **row: Any) -> None:
        """Reset private task state; the returned value never exposes hidden facts."""
        self._raise_if_failed()
        local = agent_rl_tasks.GermanAgentEnv()
        local.reset(task=task)
        if self._process is None:
            self._process = _SDKProcess(self._url, self._python)
            if self._rpc("identity") != self._expected_sdk:
                self._fail("OpenEnv installed SDK changed since planning")
        result = self._rpc("reset", task=task)
        self._check_state(result.get("state"), reset=True)
        observation = result.get("observation")
        if (not isinstance(observation, dict) or observation.get("score") is not None
                or observation.get("observation") != "" or result.get("done") is not False):
            self._fail("OpenEnv reset returned an invalid private observation")
        self._task, self._calls = copy.deepcopy(task), []

    def _tool(self, name: str, arguments: dict) -> str:
        self._raise_if_failed()
        if self._task is None:
            self._fail("Reset the OpenEnv episode before executing tools")
        result = self._rpc("step", action={"operation": "tool", "tool_name": name, "arguments": arguments})
        self._check_state(result.get("state"))
        envelope = result.get("observation")
        if (not isinstance(envelope, dict) or not isinstance(envelope.get("observation"), str)
                or envelope.get("score") is not None or result.get("done") is not False):
            self._fail("OpenEnv tool returned an invalid observation")
        observation = envelope["observation"]
        self._calls.append({"tool": name, "arguments": copy.deepcopy(arguments), "observation": observation})
        return observation

    def lookup(self, record_id: str) -> str:
        """Look up hidden quantity and unit_price for a record.

        Args:
            record_id: The exact record identifier supplied in the German task.
        """
        return self._tool("lookup", {"record_id": record_id})

    def calculator(self, expression: str) -> str:
        """Evaluate bounded literal arithmetic after looking up the hidden values.

        Args:
            expression: The product of the two retrieved integer values, as literals.
        """
        return self._tool("calculator", {"expression": expression})

    def _score(self, completion: str | list[dict]) -> dict:
        self._raise_if_failed()
        if self._task is None:
            self._fail("Reset the OpenEnv episode before verification")
        result = self._rpc("step", action={"operation": "verify", "completion": completion})
        self._check_state(result.get("state"))
        replay = agent_rl_tasks.GermanAgentEnv()
        replay.reset(task=self._task)
        for call in self._calls:
            if getattr(replay, call["tool"])(**call["arguments"]) != call["observation"]:
                self._fail("OpenEnv observation differs from independent environment replay")
        score = agent_rl_tasks.score_episode(replay, completion)
        envelope = result.get("observation")
        if (not isinstance(envelope, dict) or envelope.get("score") != score
                or result.get("done") is not True or result.get("reward") != score["reward"]):
            self._fail("OpenEnv verifier differs from independent episode replay")
        return score

    def _close(self) -> None:
        if self._process is not None:
            process, self._process = self._process, None
            try:
                process._close()
            except GymInfrastructureError as error:
                if self._infrastructure_error is None:
                    self._infrastructure_error = error
                raise self._infrastructure_error
            finally:
                self._session_id, self._instance_id, self._task = None, None, None


def _worker() -> None:
    """Native SDK only. stdout is an exact JSON channel; logs stay on stderr."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--sdk-worker", action="store_true", required=True)
    parser.add_argument("--url", required=True)
    args = parser.parse_args()
    url, client = _http._base_url(args.url), None
    try:
        for raw in sys.stdin.buffer:
            request, terminate = {}, False
            try:
                if len(raw) > MAX_BYTES:
                    raise ValueError("Oversized command")
                request = _json(raw)
                with contextlib.redirect_stdout(sys.stderr):
                    operation = request["operation"]
                    if operation == "identity":
                        result = _sdk_identity()
                    elif operation == "close":
                        if client is not None:
                            client.close()
                            client = None
                        result, terminate = {"closed": True}, True
                    else:
                        if client is None:
                            from openenv.core.generic_client import GenericEnvClient
                            client = GenericEnvClient(base_url=url, connect_timeout_s=10,
                                                      message_timeout_s=30).sync()
                            client.connect()
                        before = client.state()
                        if operation == "reset":
                            # Task IDs may repeat in the frozen schedule. Give
                            # every reset its own operational episode identity.
                            observation = client.reset(task=request["task"], episode_id=uuid.uuid4().hex)
                        elif operation == "step":
                            observation = client.step(request["action"])
                        else:
                            raise ValueError("Unknown SDK operation")
                        state = client.state()
                        if before.get("instance_id") != state.get("instance_id"):
                            raise RuntimeError("WebSocket session changed during action")
                        result = {"observation": observation.observation, "reward": observation.reward,
                                  "done": observation.done, "state": state}
                response = {"id": request["id"], "ok": True, "result": result}
            except Exception as error:
                response = {"id": request.get("id"), "ok": False, "error_type": type(error).__name__}
            print(json.dumps(response, ensure_ascii=False, allow_nan=False), flush=True)
            if terminate:
                break
    finally:
        if client is not None:
            with contextlib.redirect_stdout(sys.stderr):
                client.close()


if __name__ == "__main__":
    _worker()