"""A real pinned NeMo Gym resources server for the frozen German tool environment.

Install separately from the Python 3.11 GPU environment:
  uv venv --python 3.13.14 .venv-nemo-gym
  git clone https://github.com/NVIDIA-NeMo/Gym.git .cache/upstreams/nemo-gym
  git -C .cache/upstreams/nemo-gym checkout 720b5f318caee339bfd29c554c59a97edcedc846
  UV_PROJECT_ENVIRONMENT="$PWD/.venv-nemo-gym" uv sync \
    --project .cache/upstreams/nemo-gym --no-dev --locked
  .venv-nemo-gym/bin/python examples/agent_gyms/nemo_gym_server.py --port 18765

Use a recent uv supporting the upstream lock format. This service does not run
inference. The GPU trainer retains the sampled-token audit and injects opaque
session ids into HTTP tool requests; session ids are not model tool arguments.
"""
from __future__ import annotations

import argparse
import copy
import hashlib
import json
from pathlib import Path
import secrets
import subprocess
import sys
from typing import Any

REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT / "src"))

import nemo_gym  # noqa: E402 -- intentionally unavailable outside the isolated environment
from fastapi import FastAPI, HTTPException  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr  # noqa: E402
import uvicorn  # noqa: E402

from nemo_gym.base_resources_server import (  # noqa: E402
    BaseResourcesServerConfig, BaseSeedSessionRequest, BaseSeedSessionResponse,
    BaseVerifyRequest, BaseVerifyResponse, SimpleResourcesServer,
)
from nemo_gym.server_utils import BaseServerConfig, ServerClient  # noqa: E402
from hessian_recipes import agent_rl_tasks  # noqa: E402

UPSTREAM_REVISION = "720b5f318caee339bfd29c554c59a97edcedc846"
PROTOCOL = "hessian-gym-v1"
MAX_SESSIONS = 256


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


def identity() -> dict[str, Any]:
    """Reject an unpinned or modified upstream checkout before accepting episodes."""
    source = Path(nemo_gym.__file__).resolve().parent.parent
    revision = subprocess.check_output(
        ["git", "rev-parse", "HEAD"], cwd=source, text=True, timeout=10,
    ).strip()
    if revision != UPSTREAM_REVISION:
        raise RuntimeError(f"Expected NeMo Gym revision {UPSTREAM_REVISION}, found {revision}")
    subprocess.run(["git", "diff", "--quiet", "HEAD", "--"], cwd=source, check=True, timeout=10)
    return {"backend": "nemo_gym", "protocol": PROTOCOL,
            "upstream_revision": revision, "upstream_version": nemo_gym.__version__,
            "upstream_lock_sha256": _sha256(source / "uv.lock"),
            "environment_sha256": _sha256(Path(agent_rl_tasks.__file__)),
            "server_sha256": _sha256(Path(__file__)), "python": sys.version.split()[0],
            "framework_class": "nemo_gym.base_resources_server.SimpleResourcesServer",
            "inference": False, "max_sessions": MAX_SESSIONS}


class ResetEpisodeRequest(BaseSeedSessionRequest):
    model_config = ConfigDict(extra="forbid")
    task: dict[str, Any]


class ResetEpisodeResponse(BaseSeedSessionResponse):
    session_id: str


class EpisodeRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")
    session_id: str = Field(pattern=r"^[0-9a-f]{32}$")


class LookupRequest(EpisodeRequest):
    # Preserve model-error semantics: the frozen environment scores malformed
    # values as tool errors instead of turning them into HTTP infrastructure errors.
    record_id: Any


class CalculatorRequest(EpisodeRequest):
    expression: Any


class ToolResponse(BaseModel):
    observation: str


class VerifyEpisodeRequest(EpisodeRequest):
    completion: str | list[dict[str, Any]]


class VerifyEpisodeResponse(BaseModel):
    score: dict[str, Any]


class NativeVerifyRequest(BaseVerifyRequest):
    """Native Gym envelope plus exact external-policy transcript for strict scoring."""
    session_id: str = Field(pattern=r"^[0-9a-f]{32}$")
    completion: str | list[dict[str, Any]]


class NativeVerifyResponse(BaseVerifyResponse):
    score: dict[str, Any]


class CloseEpisodeResponse(BaseModel):
    closed: bool


class GermanResourcesServer(SimpleResourcesServer):
    """Native Gym middleware, native verification route and isolated environment state.

    Payload session ids are explicitly supported by this bridge instead of
    requiring the external TRL client to propagate Gym's HTTP session cookies.
    Every session owns an actual GermanAgentEnv. Verification freezes that state;
    repeated identical verification is idempotent, while changed transcripts fail.
    """
    config: BaseResourcesServerConfig
    _sessions: dict[str, agent_rl_tasks.GermanAgentEnv] = PrivateAttr(default_factory=dict)
    _verified: dict[str, tuple[str, dict]] = PrivateAttr(default_factory=dict)
    _identity: dict = PrivateAttr(default_factory=dict)

    def setup_webserver(self) -> FastAPI:
        self._identity = identity()
        app = super().setup_webserver()
        app.get("/hessian_identity")(self.hessian_identity)
        app.post("/reset_episode")(self.seed_session)
        app.post("/lookup")(self.lookup)
        app.post("/calculator")(self.calculator)
        app.post("/verify_episode")(self.verify_episode)
        app.post("/close_episode")(self.close_resources_session)
        return app

    async def hessian_identity(self) -> dict:
        return dict(self._identity)

    def _environment(self, session_id: str, *, tool: bool = False) -> agent_rl_tasks.GermanAgentEnv:
        if session_id not in self._sessions:
            raise HTTPException(status_code=404, detail="unknown_or_closed_session")
        if tool and session_id in self._verified:
            raise HTTPException(status_code=409, detail="episode_already_verified")
        return self._sessions[session_id]

    async def seed_session(self, body: ResetEpisodeRequest) -> ResetEpisodeResponse:
        if len(self._sessions) >= MAX_SESSIONS:
            raise HTTPException(status_code=429, detail="active_session_limit")
        env = agent_rl_tasks.GermanAgentEnv()
        try:
            env.reset(task=body.task)
        except ValueError as error:
            raise HTTPException(status_code=422, detail=str(error)) from error
        session_id = secrets.token_hex(16)
        self._sessions[session_id] = env
        return ResetEpisodeResponse(session_id=session_id)

    async def lookup(self, body: LookupRequest) -> ToolResponse:
        return ToolResponse(observation=self._environment(body.session_id, tool=True).lookup(body.record_id))

    async def calculator(self, body: CalculatorRequest) -> ToolResponse:
        return ToolResponse(observation=self._environment(body.session_id, tool=True).calculator(body.expression))

    def _score(self, session_id: str, completion: str | list[dict]) -> dict:
        env = self._environment(session_id)
        digest = hashlib.sha256(json.dumps(completion, sort_keys=True, ensure_ascii=False).encode()).hexdigest()
        if session_id in self._verified:
            previous_digest, score = self._verified[session_id]
            if digest != previous_digest:
                raise HTTPException(status_code=409, detail="verification_transcript_changed")
            return copy.deepcopy(score)
        score = agent_rl_tasks.score_episode(env, completion)
        self._verified[session_id] = (digest, copy.deepcopy(score))
        return score

    async def verify_episode(self, body: VerifyEpisodeRequest) -> VerifyEpisodeResponse:
        return VerifyEpisodeResponse(score=self._score(body.session_id, body.completion))

    async def verify(self, body: NativeVerifyRequest) -> NativeVerifyResponse:
        score = self._score(body.session_id, body.completion)
        return NativeVerifyResponse(**body.model_dump(), reward=score["reward"], score=score)

    async def close_resources_session(self, body: EpisodeRequest) -> CloseEpisodeResponse:
        self._sessions.pop(body.session_id, None)
        self._verified.pop(body.session_id, None)
        return CloseEpisodeResponse(closed=True)


def create_app(port: int = 18765) -> FastAPI:
    config = BaseResourcesServerConfig(host="127.0.0.1", port=port,
                                      entrypoint=str(Path(__file__)), name="hessian_german_tools")
    client = ServerClient(head_server_config=BaseServerConfig(host="127.0.0.1", port=port),
                          global_config_dict=OmegaConf.create({}))
    return GermanResourcesServer(config=config, server_client=client).setup_webserver()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--port", type=int, default=18765)
    args = parser.parse_args()
    if not 1024 <= args.port <= 65535:
        parser.error("Use a loopback service port between 1024 and 65535")
    uvicorn.run(create_app(args.port), host="127.0.0.1", port=args.port, workers=1,
                access_log=False, log_level="info")


if __name__ == "__main__":
    main()