# Copyright 2025 starVLA community. All rights reserved.
# Licensed under the MIT License.
"""Policy server wrapper.

Encapsulates a `baseframework` instance plus a :class:`PolicyNormProcessor`
that reuses the *training-time* :class:`ComposedModalityTransform` for action
un-normalization (no hand-rolled math). The websocket server returns
already-unnormalized actions.

Client-side responsibilities that REMAIN on the client:
  - environment-specific adapters (image_history, gripper sticky, action
    ensembling)
  - chunk-cache scheduling (`step % chunk_size == 0` triggers a new infer)

Embodiments whose DataConfig declares a ``robot_command_layout`` (currently
xtrainer2) additionally get the raw-robot conversions here, so the client never
has to reproduce training-time normalization or the quaternion order:

  - ``state`` arrives in robot coordinates (16-D, xyzw) and is reordered to the
    training layout, then normalized with the training transforms (20-D)
  - returned actions are un-normalized *and* reordered back to xyzw

Exposed API:
  - ``metadata`` (dict, sent at handshake): ``action_chunk_size``,
    ``available_unnorm_keys``, ``action_keys``, ``state_keys``, plus the
    embodiment's raw-robot contract.
  - ``predict_action(examples, unnorm_key=None, **kwargs)`` returns
    ``{"actions": np.ndarray[B, T, action_dim]}``.
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

import numpy as np
import torch

from starVLA.model.framework.base_framework import baseframework
from starVLA.model.framework.share_tools import read_mode_config

from deployment.model_server.policy_norm_processor import PolicyNormProcessor
from deployment.model_server import xtrainer2_command


class PolicyServerWrapper:
    """Wraps a `baseframework` for use as a websocket-server policy."""

    def __init__(
        self,
        ckpt_path: str,
        device: str = "cuda",
        use_bf16: bool = False,
        unnorm_key: Optional[str] = None,
        normalized_action_clip: Optional[float] = None,
    ) -> None:
        self._ckpt_path = str(ckpt_path)
        self._normalized_action_clip = normalized_action_clip

        logging.info("PolicyServerWrapper: loading framework from %s", self._ckpt_path)
        framework = baseframework.from_pretrained(self._ckpt_path)
        if use_bf16:
            framework = framework.to(torch.bfloat16)
        framework = framework.to(device).eval()
        self._framework = framework

        # Co-located metadata.
        model_cfg, _ = read_mode_config(self._ckpt_path)
        self._model_cfg = model_cfg
        vla_data_cfg = model_cfg.get("datasets", {}).get("vla_data", {})
        include_state = vla_data_cfg.get("include_state", False)
        self._uses_state = include_state not in (False, None, 0, "False", "false", "0")

        # action_chunk_size = future_action_window_size + 1 (matches old client).
        action_model_cfg = model_cfg["framework"]["action_model"]
        
        if "action_horizon" in action_model_cfg:
            self._action_chunk_size = int(action_model_cfg["action_horizon"])
        elif "future_action_window_size" in action_model_cfg:
            self._action_chunk_size = int(action_model_cfg["future_action_window_size"]) + 1
        else:
            raise ValueError(
                f"PolicyServerWrapper: no action_horizon or future_action_window_size found in model config for {self._ckpt_path}"
            )
        # Cache of PolicyNormProcessor instances per unnorm_key.
        # For single-dataset ckpts unnorm_key is auto-selected; for multi-dataset
        # ckpts clients must pass unnorm_key per request.
        self._default_unnorm_key = unnorm_key
        self._norm_processors: Dict[str, PolicyNormProcessor] = {}

        # Peek at available keys without building a full processor.
        _, _ns = read_mode_config(self._ckpt_path)
        self._available_unnorm_keys: List[str] = list(_ns.keys())

        # Eagerly build when unambiguous; defer for multi-key / no explicit key.
        if unnorm_key is not None or len(self._available_unnorm_keys) == 1:
            default_proc = self._get_processor(unnorm_key)
            self._default_unnorm_key = default_proc.unnorm_key
            logging.info(
                "PolicyServerWrapper ready: action_chunk_size=%d, default_unnorm_key=%s, "
                "available_unnorm_keys=%s, action_keys=%s, state_keys=%s",
                self._action_chunk_size,
                default_proc.unnorm_key,
                default_proc.available_unnorm_keys,
                default_proc.action_keys,
                default_proc.state_keys,
            )
        else:
            logging.info(
                "PolicyServerWrapper ready (multi-key): action_chunk_size=%d, "
                "available_unnorm_keys=%s — clients must pass unnorm_key per request.",
                self._action_chunk_size,
                self._available_unnorm_keys,
            )

    def _get_processor(self, unnorm_key: Optional[str]) -> PolicyNormProcessor:
        cache_key = unnorm_key if unnorm_key is not None else "__default__"
        if cache_key not in self._norm_processors:
            self._norm_processors[cache_key] = PolicyNormProcessor(
                self._ckpt_path, unnorm_key=unnorm_key
            )
        return self._norm_processors[cache_key]

    @staticmethod
    def _uses_xtrainer2_layout(proc: PolicyNormProcessor) -> bool:
        """Whether this embodiment speaks the raw xtrainer2 robot layout."""
        layout = getattr(proc._data_config, "robot_command_layout", None)
        return layout == xtrainer2_command.COMMAND_LAYOUT

    @property
    def metadata(self) -> Dict[str, Any]:
        """Model-invariant metadata; sent to client at websocket handshake."""
        base = {
            "env": "starvla_policy_server",
            "ckpt_path": self._ckpt_path,
            "action_chunk_size": self._action_chunk_size,
            "uses_state": self._uses_state,
            "available_unnorm_keys": self._available_unnorm_keys,
            "default_unnorm_key": self._default_unnorm_key,
            "normalized_action_clip": self._normalized_action_clip,
            "action_head_attention_mask": False,
            "inference_rng": "continuous_server_stream",
            # Every client sends images this way; the framework does its own
            # resize, so only the layout and channel order are binding.
            "image_format": "HWC_RGB_uint8",
            "image_key": "image",
        }
        # Enrich with per-embodiment keys when a default processor already exists.
        if self._default_unnorm_key is not None:
            proc = self._get_processor(self._default_unnorm_key)
            base["action_keys"] = proc.action_keys
            base["state_keys"] = proc.state_keys
            base["camera_order"] = list(getattr(proc._data_config, "video_keys", []))
            base["model_action_dim"] = proc.action_dim
            base["model_state_dim"] = proc.state_dim
            base["server_normalizes_state"] = self._server_normalizes_state(proc)
            if base["server_normalizes_state"]:
                base["state_dim"] = proc.raw_state_dim
            if self._uses_xtrainer2_layout(proc):
                base["robot_command_layout"] = xtrainer2_command.COMMAND_LAYOUT
                base["action_dim"] = proc.raw_action_dim
                base["model_quaternion_format"] = xtrainer2_command.MODEL_QUATERNION_FORMAT
                base["robot_quaternion_format"] = xtrainer2_command.SOURCE_QUATERNION_FORMAT
                base["quaternion_starts"] = list(xtrainer2_command.QUATERNION_STARTS)
                action_mode = str(
                    getattr(self, "_model_cfg", {})
                    .get("datasets", {})
                    .get("vla_data", {})
                    .get("action_mode", "abs")
                )
                base["action_mode"] = action_mode
                base["server_composes_relative_eef"] = (
                    xtrainer2_command.is_relative_action_mode(action_mode)
                )
        return base

    @staticmethod
    def _server_normalizes_state(proc: PolicyNormProcessor) -> bool:
        """Whether ``state`` arrives raw and must be normalized server-side."""
        return bool(getattr(proc._data_config, "server_normalizes_state", False))

    @classmethod
    def _prepare_examples(cls, examples: List[dict], proc: PolicyNormProcessor) -> List[dict]:
        """Turn raw-robot observations into the tensors the model was trained on.

        Leaves examples untouched for embodiments that do not opt in, so clients
        that already send normalized state keep working.
        """
        if not cls._server_normalizes_state(proc):
            return examples

        to_model_layout = (
            xtrainer2_command.robot_state_to_model
            if cls._uses_xtrainer2_layout(proc)
            else None
        )

        prepared: List[dict] = []
        for example in examples:
            if "state" not in example or example["state"] is None:
                prepared.append(example)
                continue
            raw_state = np.asarray(example["state"], dtype=np.float32)
            if raw_state.shape[-1] != proc.raw_state_dim:
                raise ValueError(
                    f"state must be the robot's raw {proc.raw_state_dim}-D vector "
                    f"(state_keys={proc.state_keys}); got shape {raw_state.shape}. "
                    "The server applies the training-time normalization itself."
                )
            if to_model_layout is not None:
                raw_state = to_model_layout(raw_state)
            example = dict(example)
            # Keep the training-layout (wxyz) pose so rel-mode actions can be
            # composed back to absolute commands after un-normalization.
            example["_raw_state_wxyz"] = raw_state
            example["state"] = proc.apply_state(raw_state)
            prepared.append(example)
        return prepared

    def predict_action(
        self,
        examples: List[dict],
        unnorm_key: Optional[str] = None,
        **kwargs,
    ) -> Dict[str, np.ndarray]:
        """Run the framework, then un-normalize via training-time transforms.

        Args:
            examples: list of dicts (each with ``image`` / ``lang`` / optional
                ``state``). For embodiments that declare
                ``server_normalizes_state``, ``state`` is the robot's raw vector
                and is normalized here.
            unnorm_key: dataset key for un-normalization stats. ``None`` -->
                use the wrapper's default (auto-picked at startup).
            **kwargs: forwarded to the framework's ``predict_action``
                (``do_sample``, ``use_ddim``, ``num_ddim_steps``, ...).

        Returns:
            ``{"actions": np.ndarray[B, T, D]}`` -- un-normalized.
        """
        effective_key = unnorm_key if unnorm_key is not None else self._default_unnorm_key
        if effective_key is None:
            if len(self._available_unnorm_keys) == 1:
                effective_key = self._available_unnorm_keys[0]
            else:
                raise ValueError(
                    f"predict_action: unnorm_key not specified and no default set. "
                    f"Pass one of {self._available_unnorm_keys}."
                )
        proc = self._get_processor(effective_key)
        examples = self._prepare_examples(examples, proc)

        out = self._framework.predict_action(examples=examples, **kwargs)
        normalized = np.asarray(out["normalized_actions"])  # (B, T, D)
        if self._normalized_action_clip is not None:
            limit = float(self._normalized_action_clip)
            normalized = np.clip(normalized, -limit, limit)

        unnorm = np.stack(
            [proc.unapply_actions(normalized[b]) for b in range(normalized.shape[0])],
            axis=0,
        )
        if self._uses_xtrainer2_layout(proc):
            if xtrainer2_command.is_relative_action_mode(
                self._model_cfg.get("datasets", {}).get("vla_data", {}).get("action_mode")
            ):
                composed = []
                for batch_index in range(unnorm.shape[0]):
                    state_wxyz = examples[batch_index].get("_raw_state_wxyz")
                    if state_wxyz is None:
                        raise ValueError(
                            "action_mode=rel requires a raw robot state so relative "
                            "EE commands can be composed back to absolute poses"
                        )
                    composed.append(
                        xtrainer2_command.compose_relative_eef_wxyz(unnorm[batch_index], state_wxyz)
                    )
                unnorm = np.stack(composed, axis=0)
            unnorm = xtrainer2_command.model_actions_to_robot(unnorm)
        return {"actions": unnorm}