# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""WebSocket connection for robot policy inference (OpenPI protocol).

Protocol (compatible with OpenPI policy clients):
    Connect  -> server sends msgpack(PolicyServerConfig fields)
    Infer    -> client sends msgpack(obs), server sends msgpack(ndarray)
    Reset    -> client sends msgpack({endpoint:reset}), server sends msgpack(status)

NumPy values use msgpack-numpy marker mappings. Outbound payloads follow the
openpi-client format so official OpenPI clients can decode action responses:
    ndarray -> {__ndarray__: true, dtype, shape, data}
    scalar  -> {__npgeneric__: true, dtype, data}

Inbound payloads accept both the openpi-client markers above and the legacy
vLLM-native markers:
    ndarray -> {nd: true, type, kind, shape, data}
    scalar  -> {nd: false, type, kind, data}

`kind` is required on the `nd` markers because it is what separates them from a
plain user mapping, but its value is dialect-specific: the vLLM-native markers
carry the dtype kind character while the `msgpack-numpy` package leaves it empty
for plain dtypes and sets "V" for structured ones. Both spellings are accepted.
Scalars packed by `msgpack-numpy` omit `kind` entirely and are therefore left as
mappings.
"""

from __future__ import annotations

import asyncio
from collections import OrderedDict
from typing import Any

import msgspec
import numpy as np
from fastapi import WebSocket
from starlette.websockets import WebSocketDisconnect
from vllm.logger import init_logger

from vllm_omni.entrypoints.openpi.serving import (
    ServingRealtimeRobotOpenPI,
)

logger = init_logger(__name__)
_DEFAULT_IDLE_TIMEOUT = 30.0
MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024
_MISSING = object()
# Upper bound on the per-connection set of seen session ids.
MAX_TRACKED_SESSIONS = 1024


def _pack_numpy(obj: Any) -> Any:
    if isinstance(obj, np.ndarray | np.generic) and obj.dtype.kind in ("V", "O", "c"):
        raise ValueError(f"Unsupported dtype: {obj.dtype}")
    if isinstance(obj, np.ndarray):
        if not obj.flags.c_contiguous:
            obj = np.ascontiguousarray(obj)
        return {
            b"__ndarray__": True,
            b"data": obj.tobytes(),
            b"dtype": obj.dtype.str,
            b"shape": obj.shape,
        }
    if isinstance(obj, np.generic):
        return {
            b"__npgeneric__": True,
            b"data": obj.item(),
            b"dtype": obj.dtype.str,
        }
    raise TypeError(f"Unsupported type: {type(obj)!r}")


def _mapping_get(obj: dict[Any, Any], key: str, default: Any = None) -> Any:
    return obj.get(key, obj.get(key.encode(), default))


def _decode_marker_text(value: Any) -> str:
    if isinstance(value, bytes):
        return value.decode()
    return str(value)


def _is_truthy_marker(value: Any) -> bool:
    return value is True or value == 1


def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any:
    nd = _mapping_get(obj, "nd", _MISSING)
    dtype = _mapping_get(obj, "type", _MISSING)
    kind = _mapping_get(obj, "kind", _MISSING)
    data = _mapping_get(obj, "data", _MISSING)
    if nd is _MISSING or dtype is _MISSING or kind is _MISSING or data is _MISSING:
        return _MISSING

    kind_text = _decode_marker_text(kind)
    # Structured markers carry a dtype descriptor list in `type`, so reject them
    # before `np.dtype()` sees something it cannot parse.
    if kind_text == "V":
        raise ValueError("Unsupported dtype: structured arrays")

    dtype_obj = np.dtype(_decode_marker_text(dtype))
    if kind_text not in ("", dtype_obj.kind):
        raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}")
    if dtype_obj.kind in ("V", "O", "c"):
        raise ValueError(f"Unsupported dtype: {dtype_obj}")

    array = np.frombuffer(data, dtype=dtype_obj).copy()
    if nd:
        shape = _mapping_get(obj, "shape", _MISSING)
        if shape is _MISSING:
            raise ValueError("NumPy ndarray marker is missing shape")
        return array.reshape(tuple(shape))
    return array[0]


def _decode_openpi_numpy_marker(obj: dict[Any, Any]) -> Any:
    if _is_truthy_marker(_mapping_get(obj, "__ndarray__", False)):
        data = _mapping_get(obj, "data", _MISSING)
        dtype = _mapping_get(obj, "dtype", _MISSING)
        shape = _mapping_get(obj, "shape", _MISSING)
        if data is _MISSING or dtype is _MISSING or shape is _MISSING:
            raise ValueError("OpenPI ndarray marker is missing required fields")

        dtype_obj = np.dtype(_decode_marker_text(dtype))
        if dtype_obj.kind in ("V", "O", "c"):
            raise ValueError(f"Unsupported dtype: {dtype_obj}")
        array = np.frombuffer(data, dtype=dtype_obj).copy()
        return array.reshape(tuple(shape))

    if _is_truthy_marker(_mapping_get(obj, "__npgeneric__", False)):
        data = _mapping_get(obj, "data", _MISSING)
        dtype = _mapping_get(obj, "dtype", _MISSING)
        if data is _MISSING or dtype is _MISSING:
            raise ValueError("OpenPI scalar marker is missing required fields")
        dtype_obj = np.dtype(_decode_marker_text(dtype))
        if dtype_obj.kind in ("V", "O", "c"):
            raise ValueError(f"Unsupported dtype: {dtype_obj}")
        return dtype_obj.type(data)

    return _MISSING


def _unpack_numpy(obj: Any) -> Any:
    if isinstance(obj, dict):
        for decoder in (_decode_vllm_numpy_marker, _decode_openpi_numpy_marker):
            decoded = decoder(obj)
            if decoded is not _MISSING:
                return decoded
        return {key: _unpack_numpy(value) for key, value in obj.items()}
    if isinstance(obj, list):
        return [_unpack_numpy(value) for value in obj]
    return obj


def _pack(obj: Any) -> bytes:
    return msgspec.msgpack.encode(obj, enc_hook=_pack_numpy)


def _unpack(data: bytes) -> Any:
    return _unpack_numpy(msgspec.msgpack.decode(data))


class RobotRealtimeConnection:
    """WebSocket connection for robot policy inference."""

    def __init__(
        self,
        websocket: WebSocket,
        serving: ServingRealtimeRobotOpenPI,
        idle_timeout: float | None = _DEFAULT_IDLE_TIMEOUT,
    ) -> None:
        self.websocket = websocket
        self.serving = serving
        self._idle_timeout = idle_timeout
        self._current_session_id: str | None = None
        # Session ids seen on this connection, most-recently-used last.
        self._seen_sessions: OrderedDict[str, None] = OrderedDict()

    def reset(self) -> None:
        self._current_session_id = None
        self._seen_sessions.clear()

    def _mark_session_seen(self, session_id: str) -> bool:
        """Record ``session_id``; return True when it is the first sighting.

        A session's first observation is the one that carries ``reset`` to the
        policy, so this needs per-session memory.
        """
        if session_id in self._seen_sessions:
            self._seen_sessions.move_to_end(session_id)
            return False

        self._seen_sessions[session_id] = None
        if len(self._seen_sessions) > MAX_TRACKED_SESSIONS:
            evicted, _ = self._seen_sessions.popitem(last=False)
            logger.debug(
                "Robot OpenPI connection is tracking more than %d sessions; dropped least-recently-used %s",
                MAX_TRACKED_SESSIONS,
                evicted,
            )
        return True

    async def _send_error(self, message: str) -> None:
        await self.websocket.send_bytes(_pack({"type": "error", "message": message}))

    def _unpack_request(self, data: bytes) -> dict[str, Any]:
        if len(data) > MAX_OPENPI_PAYLOAD_BYTES:
            raise ValueError("OpenPI request payload too large")
        obs = _unpack(data)
        if not isinstance(obs, dict):
            raise ValueError("Invalid request payload")
        return obs

    async def handle_connection(self) -> None:
        """Main loop for OpenPI-compatible policy serving."""
        await self.websocket.accept()

        try:
            # Send model-specific PolicyServerConfig resolved by serving from
            # diffusion od_config.model_config.
            metadata = self.serving.policy_server_config.to_dict()
            await self.websocket.send_bytes(_pack(metadata))

            while True:
                idle_timeout = self._idle_timeout
                if idle_timeout is None:
                    msg = await self.websocket.receive()
                else:
                    try:
                        msg = await asyncio.wait_for(
                            self.websocket.receive(),
                            timeout=idle_timeout,
                        )
                    except asyncio.TimeoutError:
                        logger.info("Robot OpenPI connection idle timeout after %.1f seconds", idle_timeout)
                        try:
                            await self.websocket.close()
                        except Exception:
                            logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True)
                        return

                if msg.get("type") == "websocket.disconnect":
                    break

                if "bytes" not in msg or not msg["bytes"]:
                    continue

                try:
                    obs = self._unpack_request(msg["bytes"])
                except Exception:
                    logger.exception("Invalid robot OpenPI request payload")
                    try:
                        await self._send_error("Invalid request payload")
                    except Exception:
                        break
                    continue

                try:
                    endpoint = obs.pop("endpoint", "infer")

                    if endpoint == "reset":
                        self.reset()
                        self.serving.reset(obs)
                        await self.websocket.send_bytes(_pack({"status": "reset successful"}))
                    else:
                        session_id = str(obs.get("session_id") or self._current_session_id or "default")
                        if session_id != self._current_session_id:
                            if self._current_session_id is not None:
                                logger.info(
                                    "Robot OpenPI session changed %s -> %s",
                                    self._current_session_id,
                                    session_id,
                                )
                            self._current_session_id = session_id

                        reset = self._mark_session_seen(session_id)
                        actions = await self.serving.infer(
                            obs,
                            session_id=session_id,
                            reset=reset,
                        )
                        await self.websocket.send_bytes(_pack(actions))
                except Exception:
                    logger.exception("Error handling request")
                    try:
                        await self._send_error("Internal inference error")
                    except Exception:
                        break

        except WebSocketDisconnect:
            pass
        except Exception:
            logger.exception("Connection error")