"""Pinned native OpenEnv service for the unchanged German lookup/calculator task.

Use the SDK's persistent WebSocket client: its HTTP compatibility routes are
stateless. Only the trusted rollout controller supplies reset metadata and calls
verify; the model sees the existing prompt and the two actual tool observations.
"""
from __future__ import annotations

import argparse
import copy
import hashlib
import importlib.metadata
import json
import os
from pathlib import Path
import platform
import sys
import threading
from typing import Any, Literal
import uuid
from fastapi import Response
from openenv.core.client_types import StepResult
from openenv.core.env_client import EnvClient
from openenv.core.env_server import Action, Environment, Observation, State
from openenv.core.env_server.http_server import create_app
from openenv.core.env_server.types import WSErrorResponse
from pydantic import Field

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

UPSTREAM_REVISION = "17c1edd22168fcb76b086def1ed90a7168fb8749"
PROTOCOL = "hessian-openenv-v1"
MAX_PAYLOAD_BYTES = 524_288
_LIFECYCLE_LOCK = threading.Lock()
_ACTIVE_INSTANCES: set[str] = set()


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


def _strict_json(raw: bytes) -> Any:
    if len(raw) > MAX_PAYLOAD_BYTES:
        raise ValueError("Controller payload exceeds the byte bound")

    def pairs(items):
        result = {}
        for key, value in items:
            if key in result:
                raise ValueError("Duplicate JSON key")
            result[key] = value
        return result

    def constant(value):
        raise ValueError("Nonfinite JSON constant")

    return json.loads(raw, object_pairs_hook=pairs, parse_constant=constant)


class StrictControllerJSON:
    """Bound and validate JSON before the unchanged SDK dispatches requests."""

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] == "websocket":
            async def checked_receive():
                while True:
                    message = await receive()
                    if message["type"] != "websocket.receive":
                        return message
                    try:
                        raw = message.get("bytes")
                        if raw is None:
                            raw = message.get("text", "").encode()
                        _strict_json(raw)
                    except (ValueError, TypeError, UnicodeError, RecursionError):
                        error = WSErrorResponse(data={"message": "Invalid bounded JSON message",
                                                      "code": "INVALID_JSON"})
                        await send({"type": "websocket.send", "text": error.model_dump_json()})
                        continue
                    return message
            return await self.app(scope, checked_receive, send)
        if scope["type"] == "http" and scope["method"] in {"POST", "PUT", "PATCH"}:
            body = bytearray()
            try:
                while True:
                    message = await receive()
                    if message["type"] != "http.request":
                        return
                    body.extend(message.get("body", b""))
                    if len(body) > MAX_PAYLOAD_BYTES:
                        raise ValueError("Controller payload exceeds the byte bound")
                    if not message.get("more_body", False):
                        break
                _strict_json(bytes(body))
            except (ValueError, TypeError, UnicodeError, RecursionError):
                await send({"type": "http.response.start", "status": 400,
                            "headers": [(b"content-type", b"application/json")]})
                await send({"type": "http.response.body", "body": b'{"detail":"Invalid bounded JSON body"}'})
                return
            consumed = False

            async def replay_receive():
                nonlocal consumed
                if not consumed:
                    consumed = True
                    return {"type": "http.request", "body": bytes(body), "more_body": False}
                return await receive()
            return await self.app(scope, replay_receive, send)
        return await self.app(scope, receive, send)


def runtime_identity() -> dict:
    distribution = importlib.metadata.distribution("openenv")
    direct = json.loads(distribution.read_text("direct_url.json") or "{}")
    vcs = direct.get("vcs_info", {})
    if distribution.version != "0.9.0" or vcs.get("commit_id") != UPSTREAM_REVISION:
        raise RuntimeError("Use the exact OpenEnv VCS revision in requirements/openenv.txt")
    package = Path(sys.modules[EnvClient.__module__].__file__).parents[1]
    python_sources = {str(path.relative_to(package)): _sha(path)
                      for path in sorted(package.rglob("*.py"))}
    return {"version": distribution.version, "commit": vcs["commit_id"],
            "source_url": direct.get("url"),
            "python_version": platform.python_version(),
            "installed_distributions_sha256": hashlib.sha256(json.dumps(sorted(
                (item.metadata["Name"], item.version) for item in importlib.metadata.distributions()
            )).encode()).hexdigest(),
            "python_source_sha256": python_sources,
            "python_source_tree_sha256": hashlib.sha256(json.dumps(python_sources, sort_keys=True).encode()).hexdigest(),
            "env_client_sha256": _sha(Path(sys.modules[EnvClient.__module__].__file__)),
            "http_server_sha256": _sha(Path(sys.modules[create_app.__module__].__file__))}


def identity() -> dict:
    with _LIFECYCLE_LOCK:
        active = len(_ACTIVE_INSTANCES)
    return {"backend": "openenv", "protocol": PROTOCOL,
            "upstream_revision": UPSTREAM_REVISION,
            "server_sha256": _sha(Path(__file__)),
            "environment_sha256": _sha(ROOT / "src/hessian_recipes/agent_rl_tasks.py"),
            "requirements_sha256": _sha(ROOT / "requirements/openenv.txt"),
            "runtime": runtime_identity(), "transport": "native-sdk-websocket",
            "active_environment_instances": active}


class GermanAction(Action):
    """Controller envelope; model argument values deliberately remain uncoerced."""

    operation: Literal["tool", "verify"]
    tool_name: str | None = None
    arguments: dict[str, Any] = Field(default_factory=dict)
    completion: Any = None


class GermanObservation(Observation):
    observation: str = ""
    score: dict[str, Any] | None = None


class GermanState(State):
    instance_id: str
    done: bool = False


class GermanOpenEnv(Environment[GermanAction, GermanObservation, GermanState]):
    """A fresh frozen task instance per native SDK WebSocket session."""

    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self):
        super().__init__()
        self._instance_id = str(uuid.uuid4())
        self._env = GermanAgentEnv()
        self._score: dict | None = None
        self._completion_digest: str | None = None
        self._state = GermanState(instance_id=self._instance_id)
        self._closed = False
        with _LIFECYCLE_LOCK:
            _ACTIVE_INSTANCES.add(self._instance_id)

    def reset(self, seed: int | None = None, episode_id: str | None = None,
              task: dict | None = None, **kwargs: Any) -> GermanObservation:
        if self._closed:
            raise RuntimeError("The native session has closed")
        # A failed controller reset must not leave the preceding task available.
        self._env, self._score, self._completion_digest = GermanAgentEnv(), None, None
        self._state = GermanState(instance_id=self._instance_id)
        if kwargs:
            raise ValueError("Unexpected controller reset parameters")
        _strict_json(json.dumps(task, ensure_ascii=False, allow_nan=False).encode())
        self._env.reset(task=task)
        self._state = GermanState(instance_id=self._instance_id,
                                  episode_id=episode_id or task.get("task_id"),
                                  step_count=0, done=False)
        return GermanObservation(reward=0.0)

    def step(self, action: GermanAction, timeout_s: float | None = None,
             **kwargs: Any) -> GermanObservation:
        if self._closed or self._env._task is None:
            raise RuntimeError("Reset a live native session before stepping")
        if kwargs:
            raise ValueError("Unexpected controller step parameters")
        _strict_json(json.dumps(action.model_dump(), ensure_ascii=False, allow_nan=False).encode())
        if action.operation == "verify":
            if action.tool_name is not None or action.arguments:
                raise ValueError("Verification cannot contain a tool call")
            digest = hashlib.sha256(json.dumps(action.completion, ensure_ascii=False,
                                               sort_keys=True).encode()).hexdigest()
            if self._score is not None:
                if digest != self._completion_digest:
                    raise ValueError("A verified transcript cannot be replaced")
                return GermanObservation(done=True, reward=self._score["reward"],
                                         score=copy.deepcopy(self._score))
            result = score_episode(self._env, action.completion)
            self._score, self._completion_digest = copy.deepcopy(result), digest
            self._env._terminal = copy.deepcopy(result)
            self._state.done = True
            self._state.step_count += 1
            return GermanObservation(done=True, reward=result["reward"], score=result)
        if self._score is not None:
            raise ValueError("Tools cannot mutate a verified episode")
        if action.completion is not None:
            raise ValueError("Tool actions cannot contain a final transcript")
        required = {"record_id"} if action.tool_name == "lookup" else {"expression"}
        if action.tool_name not in {"lookup", "calculator"} or set(action.arguments) != required:
            result = step_action(self._env, {"tool": None})
        else:
            result = step_action(self._env, {"tool": action.tool_name, **action.arguments})
        self._state.step_count += 1
        # A tool limit is a scored task error. Verification still needs to accept
        # the actual final transcript, including malformed or truncated outputs.
        return GermanObservation(observation=result["observation"], reward=0.0)

    @property
    def state(self) -> GermanState:
        return self._state.model_copy(deep=True)

    def close(self) -> None:
        if not self._closed:
            self._closed = True
            self._env = GermanAgentEnv()
            self._score, self._completion_digest = None, None
            with _LIFECYCLE_LOCK:
                _ACTIVE_INSTANCES.discard(self._instance_id)


class GermanOpenEnvClient(EnvClient[GermanAction, GermanObservation, GermanState]):
    """Typed native SDK client; GenericEnvClient uses the identical wire API."""

    def _step_payload(self, action: GermanAction) -> dict:
        return action.model_dump()

    def _parse_result(self, payload: dict) -> StepResult[GermanObservation]:
        observation = GermanObservation.model_validate(
            payload["observation"] | {"done": payload["done"], "reward": payload["reward"]})
        return StepResult(observation=observation, reward=payload["reward"],
                          done=payload["done"], metadata=payload.get("metadata"))

    def _parse_state(self, payload: dict) -> GermanState:
        return GermanState.model_validate(payload)


def make_app(max_sessions: int = 256):
    if type(max_sessions) is not int or not 1 <= max_sessions <= 256:
        raise ValueError("Use between one and 256 native sessions")
    runtime_identity()
    app = create_app(GermanOpenEnv, GermanAction, GermanObservation,
                     state_cls=GermanState, env_name="hessian_german_tools",
                     max_concurrent_envs=max_sessions, mode="simulation")
    app.add_middleware(StrictControllerJSON)

    @app.get("/hessian_identity")
    def hessian_identity(response: Response):
        response.headers["X-Hessian-Process-Id"] = str(os.getpid())
        return identity()

    return app


def main() -> None:
    import uvicorn

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument("--max-sessions", type=int, default=256)
    args = parser.parse_args()
    if not 1024 <= args.port <= 65535:
        raise ValueError("Use an unprivileged loopback port")
    uvicorn.run(make_app(args.max_sessions), host="127.0.0.1", port=args.port,
                workers=1, log_level="warning", ws_max_size=MAX_PAYLOAD_BYTES)


if __name__ == "__main__":
    main()