"""Loopback NeMo Gym tools for native TRL, with independent episode replay.

Only reset, lookup and calculator are public methods on the environment: TRL
discovers public methods as model tools. HTTP failures are sticky infrastructure
errors even if the upstream tool loop catches an individual call exception.
"""
from __future__ import annotations

import copy
import hashlib
from http.client import HTTPException
import json
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import urlsplit
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener

from . import agent_rl_tasks


NEMO_GYM_REVISION = "720b5f318caee339bfd29c554c59a97edcedc846"
PROTOCOL = "hessian-gym-v1"
MAX_RESPONSE_BYTES = 2_000_000


class GymInfrastructureError(RuntimeError):
    """A gym service or verification failure, never a scored model mistake."""


def _base_url(value: str) -> str:
    try:
        parsed = urlsplit(value)
        valid = (parsed.scheme == "http" and parsed.hostname in {"127.0.0.1", "localhost", "::1"}
                 and parsed.port is not None and 1 <= parsed.port <= 65535
                 and parsed.path in {"", "/"} and not parsed.query and not parsed.fragment
                 and parsed.username is None and parsed.password is None)
    except (ValueError, TypeError, AttributeError) as error:
        raise ValueError("Gym URL must be a loopback HTTP origin with an explicit port") from error
    if not valid:
        raise ValueError("Gym URL must be a loopback HTTP origin with an explicit port")
    return value.rstrip("/")


class _NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        raise GymInfrastructureError("Gym service redirects are forbidden")


def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError("Duplicate JSON key in gym response")
        result[key] = value
    return result


def _request(base_url: str, endpoint: str, payload: dict | None = None) -> dict:
    data = None if payload is None else json.dumps(payload, ensure_ascii=False, allow_nan=False).encode()
    request = Request(base_url + endpoint, data=data,
                      headers={"Content-Type": "application/json", "Accept": "application/json"})
    try:
        # Local service traffic must not inherit external HTTP proxy configuration.
        with build_opener(ProxyHandler({}), _NoRedirect()).open(request, timeout=15) as response:
            raw = response.read(MAX_RESPONSE_BYTES + 1)
        if len(raw) > MAX_RESPONSE_BYTES:
            raise ValueError("Gym response exceeds the bounded response size")
        result = json.loads(raw, object_pairs_hook=_unique_object,
                            parse_constant=lambda value: (_ for _ in ()).throw(ValueError(value)))
        if not isinstance(result, dict):
            raise ValueError("Gym response must be an object")
        return result
    except (HTTPError, URLError, HTTPException, TimeoutError, OSError, ValueError) as error:
        # Never include bodies or private task facts in an exception message.
        raise GymInfrastructureError(f"Gym request failed at {endpoint}: {type(error).__name__}") from error


def fetch_gym_identity(base_url: str) -> dict:
    """Require our pinned native Gym backend and the consumed environment source."""
    identity = _request(_base_url(base_url), "/hessian_identity")
    expected = {"protocol": PROTOCOL, "backend": "nemo_gym",
                "upstream_revision": NEMO_GYM_REVISION,
                "environment_sha256": hashlib.sha256(Path(agent_rl_tasks.__file__).read_bytes()).hexdigest()}
    for key, value in expected.items():
        if identity.get(key) != value:
            raise GymInfrastructureError(f"Gym identity mismatch for {key}")
    return identity


class RemoteGymEnv:
    """An isolated remote Gym episode; model-visible tools match the local recipe."""

    def __init__(self, base_url: str, identity: dict):
        self._base_url = _base_url(base_url)
        self._provenance = copy.deepcopy(identity)
        self._task: dict | None = None
        self._session_id: str | None = None
        self._calls: list[dict] = []
        self._infrastructure_error: GymInfrastructureError | None = None

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

    def _rpc(self, endpoint: str, payload: dict) -> dict:
        self._raise_if_failed()
        try:
            return _request(self._base_url, endpoint, payload)
        except GymInfrastructureError as error:
            self._infrastructure_error = error
            raise

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

    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)
        self._close()
        result = self._rpc("/reset_episode", {"task": task})
        session_id = result.get("session_id")
        if not isinstance(session_id, str) or not session_id or len(session_id) > 200:
            self._fail("Gym reset returned an invalid episode identity")
        self._session_id = session_id
        self._task = copy.deepcopy(task)
        self._calls = []

    def _tool(self, name: str, arguments: dict) -> str:
        self._raise_if_failed()
        if self._session_id is None:
            self._fail("Reset the gym episode before executing tools")
        result = self._rpc("/" + name, {"session_id": self._session_id, **arguments})
        observation = result.get("observation")
        if not isinstance(observation, str):
            self._fail("Gym tool returned a non-text 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 or self._session_id is None:
            self._fail("Reset the gym episode before verification")
        remote = self._rpc("/verify_episode", {"session_id": self._session_id, "completion": completion})
        replay = agent_rl_tasks.GermanAgentEnv()
        replay.reset(task=self._task)
        for call in self._calls:
            observed = getattr(replay, call["tool"])(**call["arguments"])
            if observed != call["observation"]:
                self._fail("Gym observation differs from independent environment replay")
        score = agent_rl_tasks.score_episode(replay, completion)
        if remote.get("score") != score:
            self._fail("Gym verifier differs from independent episode replay")
        return score

    def _close(self) -> None:
        if self._session_id is not None:
            # Cleanup still runs after a failed request. Preserve its primary error.
            try:
                result = _request(self._base_url, "/close_episode", {"session_id": self._session_id})
                if result.get("closed") is not True:
                    raise GymInfrastructureError("Gym close did not acknowledge session cleanup")
            except GymInfrastructureError as error:
                if self._infrastructure_error is None:
                    self._infrastructure_error = error
                raise self._infrastructure_error
            finally:
                self._session_id = None