#!/usr/bin/env python3
"""pi0.5 (openpi, in-process) -> GR00T-WBC-Bridge ZMQ adapter server (G1 piston).

Serves the *exact* GR00T ``PolicyServer`` ZMQ REP wire contract so the existing
piston bridge (``gr00t_n1_wbc_bridge_deploy.py --server_codec custom
--transport zmq_isaac --server_port 5555``) can drive our pi0.5 LoRA checkpoint
with ZERO bridge changes -- a drop-in replacement for ``run_gr00t_server.py``.

Unlike ``pi05_bridge_adapter.py`` (the June 2-hop ZMQ<->websocket version), this
runs pi0.5 IN-PROCESS: it loads the trained openpi ``Policy`` once at startup via
``create_trained_policy(get_config(<config>), <ckpt_dir>)`` and calls
``policy.infer(...)`` directly -- no websocket server hop.

Wire contract (vendored from Isaac-GR00T ``gr00t/policy/server_client.py`` --
this file has ZERO Isaac-GR00T imports):

  * ZMQ REP bound to ``tcp://0.0.0.0:<port>``.
  * msgpack payloads with a custom ndarray ext (``__ndarray_class__`` + np.save
    bytes), decoded via an ``object_hook`` (see ``MsgSerializer``).
  * Request envelope ``{"endpoint": str, "data": {...}}``. Endpoints:
      - ``ping``  -> ``{"status": "ok", "message": "Server is running"}``
      - ``kill``  -> shut the loop down, reply ``{}``
      - ``reset`` (``data={"options": None}``) -> ``{}``
      - ``get_action`` (``data={"observation": obs, "options": None}``)
      - ``get_modality_config`` -> ``{}``  (GR00T's default; kept for parity)
  * REP/REQ lock-step: every ``recv`` sends exactly one reply (even on error),
    otherwise the socket wedges. Malformed requests are logged and answered
    with ``{"error": <msg>}`` -- the server never crashes on bad input.

Observation contract (piston profile, ``--camera_passthrough``):
  * ``obs["video"]["ego_view"]`` : uint8 ``(1, 1, H, W, 3)`` at the camera's
    RAW resolution (e.g. 240x424) -- squeezed to HWC RGB and handed straight to
    openpi as ``observation/image`` (openpi's G1Inspire transform owns the 224
    resize; we pass native res).
  * ``obs["state"]`` : dict of 5 groups -- USED (pi0.5 was trained with state).
    CONCATENATE left_arm(7), right_arm(7), left_hand(6), right_hand(6),
    waist(3) -> float32 (29,) -> ``observation/state``.
  * ``obs["language"]["annotation.human.task_description"]`` -> ``prompt``.

Action response: a 2-element list ``[action_dict, info_dict]`` (the bridge
unwraps ``result[0]``). ``policy.infer(...)["actions"]`` is (horizon, 30)
ABSOLUTE physical-unit joint targets. Sliced by fixed columns into the bridge's
per-group dict, keyed WITHOUT the ``action.`` prefix:
  left_arm  [0:7]   (1,H,7)
  right_arm [7:14]  (1,H,7)
  left_hand [14:20] (1,H,6)
  right_hand[20:26] (1,H,6)
  base_height[26:27](1,H,1)
  navigate_command[27:30] (1,H,3)
H (horizon) is read from the returned shape -- NOT hardcoded (it is 10 here).

Run from the openpi repo with its .venv, GPU-free for CPU smoke:
    cd <openpi> && CUDA_VISIBLE_DEVICES="" JAX_PLATFORMS=cpu .venv/bin/python \
        scripts/zmq_adapter_pi05.py --ckpt_dir <ckpt> --port 5544
"""

from __future__ import annotations

import argparse
import io
import logging
import os
import time
import traceback
from typing import Any, Dict, List, Tuple

import numpy as np
import msgpack
import zmq


# ---------------------------------------------------------------------------
# Vendored codec -- byte-for-byte compatible with Isaac-GR00T
# gr00t/policy/server_client.py :: MsgSerializer (ndarray path only; the
# bridge's get_action payloads never carry ModalityConfig objects).
# ---------------------------------------------------------------------------
class MsgSerializer:
    @staticmethod
    def to_bytes(data: Any) -> bytes:
        return msgpack.packb(data, default=MsgSerializer.encode_custom_classes)

    @staticmethod
    def from_bytes(data: bytes) -> Any:
        return msgpack.unpackb(data, object_hook=MsgSerializer.decode_custom_classes)

    @staticmethod
    def decode_custom_classes(obj):
        if not isinstance(obj, dict):
            return obj
        if "__ndarray_class__" in obj:
            return np.load(io.BytesIO(obj["as_npy"]), allow_pickle=False)
        return obj

    @staticmethod
    def encode_custom_classes(obj):
        if isinstance(obj, np.ndarray):
            output = io.BytesIO()
            np.save(output, obj, allow_pickle=False)
            return {"__ndarray_class__": True, "as_npy": output.getvalue()}
        return obj


# ---------------------------------------------------------------------------
# Fixed slice plan for the pi0.5 flat (horizon, 30) chunk -> per-group dict.
# These columns are the G1-Inspire dataset's action layout (verified in
# openpi/src/openpi/policies/g1_inspire_policy.py docstring): left_arm 0:7,
# right_arm 7:14, left_hand 14:20, right_hand 20:26, base_height 26:27,
# navigate 27:30. Keys match the bridge's n17_redball_inspire profile groups.
# ---------------------------------------------------------------------------
_SLICE_PLAN: List[Tuple[str, slice]] = [
    ("left_arm", slice(0, 7)),
    ("right_arm", slice(7, 14)),
    ("left_hand", slice(14, 20)),
    ("right_hand", slice(20, 26)),
    ("base_height", slice(26, 27)),
    ("navigate_command", slice(27, 30)),
]
# proprio state layout the bridge sends (5 groups -> 29-dim concat).
_STATE_PLAN: List[Tuple[str, int]] = [
    ("left_arm", 7), ("right_arm", 7), ("left_hand", 6),
    ("right_hand", 6), ("waist", 3),
]
_DEFAULT_INSTRUCTION = "pick up the piston."

log = logging.getLogger("zmq_adapter_pi05")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _flat(x) -> np.ndarray:
    return np.asarray(x, dtype=np.float32).reshape(-1)


def _obs_to_state(obs: Dict[str, Any]) -> np.ndarray:
    """obs["state"] dict of 5 groups -> float32 (29,) proprio concat, in the
    order left_arm, right_arm, left_hand, right_hand, waist (each sliced to its
    canonical dim). pi0.5's G1Inspire transform keeps only the first 29 dims."""
    st = obs.get("state")
    if not isinstance(st, dict):
        raise ValueError("observation missing state dict")
    parts = []
    for key, dim in _STATE_PLAN:
        if key not in st:
            raise ValueError(f"observation.state missing group {key!r}")
        parts.append(_flat(st[key])[:dim])
    state = np.concatenate(parts).astype(np.float32)
    if state.shape[0] != 29:
        raise ValueError(f"state concat is {state.shape[0]} dims, expected 29")
    return state


def _obs_to_hwc_rgb(obs: Dict[str, Any]) -> np.ndarray:
    """``obs["video"]["ego_view"]`` (1,1,H,W,3) uint8 -> contiguous HWC uint8."""
    video = obs.get("video")
    if not isinstance(video, dict) or "ego_view" not in video:
        raise ValueError("observation missing video.ego_view")
    arr = np.asarray(video["ego_view"])
    # Peel leading singleton batch/time axes -> (H, W, C).
    while arr.ndim > 3 and arr.shape[0] == 1:
        arr = arr[0]
    if arr.ndim != 3 or arr.shape[-1] not in (1, 3, 4):
        raise ValueError(f"unexpected ego_view shape "
                         f"{np.asarray(video['ego_view']).shape} -> squeezed "
                         f"{arr.shape}; expected (...,H,W,C)")
    if arr.dtype != np.uint8:
        arr = np.clip(arr, 0, 255).astype(np.uint8)
    return np.ascontiguousarray(arr)


def _extract_instruction(obs: Dict[str, Any], override: str) -> str:
    """CLI override wins; else the obs language field; else the piston default.

    Bridge sends obs["language"]["annotation.human.task_description"] = [["<i>"]].
    """
    if override:
        return override
    try:
        raw = obs.get("language", {}).get("annotation.human.task_description")
    except AttributeError:
        raw = None
    node = raw
    for _ in range(6):
        if isinstance(node, bytes):
            node = node.decode("utf-8", "ignore")
        if isinstance(node, str):
            s = node.strip()
            if s:
                return s
            break
        if isinstance(node, (list, tuple)) and node:
            node = node[0]
            continue
        break
    return _DEFAULT_INSTRUCTION


def _latest_step_dir(path: str) -> str:
    """If ``path`` is a run dir whose children are integer step dirs (openpi
    orbax layout, e.g. .../piston_v1/{5000,10000,...}), return the newest step
    subdir. If ``path`` itself already contains ``params/`` (it IS a step dir),
    return it unchanged. Falls back to the given path."""
    try:
        if os.path.isdir(os.path.join(path, "params")):
            return path
        steps = []
        for name in os.listdir(path):
            full = os.path.join(path, name)
            if name.isdigit() and os.path.isdir(full):
                steps.append((int(name), full))
        if steps:
            return max(steps)[1]
    except OSError:
        pass
    return path


# ---------------------------------------------------------------------------
# Adapter
# ---------------------------------------------------------------------------
class Pi05ZMQAdapter:
    def __init__(self, config_name: str, ckpt_dir: str, port: int,
                 instruction: str):
        self.port = int(port)
        self.instruction_override = instruction or ""
        self.running = True
        self._reqs = 0

        # Import here so import errors surface with a clear message and after
        # argparse (cheap --help). Requires cwd/env = openpi repo + its .venv.
        from openpi.training.config import get_config
        from openpi.policies import policy_config

        ckpt_dir = _latest_step_dir(ckpt_dir)
        log.info("loading pi0.5 policy: config=%s ckpt_dir=%s ...",
                 config_name, ckpt_dir)
        t0 = time.time()
        self.policy = policy_config.create_trained_policy(
            get_config(config_name), ckpt_dir,
        )
        self.load_seconds = time.time() - t0
        self.ckpt_dir = ckpt_dir
        log.info("model loaded in %.1fs (metadata=%s)",
                 self.load_seconds, self.policy.metadata)

    # ---- endpoint handlers ------------------------------------------------
    def handle_ping(self) -> dict:
        return {"status": "ok", "message": "Server is running"}

    def handle_kill(self) -> dict:
        log.info("kill endpoint received; shutting down after reply")
        self.running = False
        return {}

    def handle_reset(self, options=None) -> dict:
        # pi0.5 inference is stateless per call (the bridge owns the chunk
        # queue), so reset is a no-op -- but the endpoint must exist.
        return {}

    def handle_get_modality_config(self) -> dict:
        return {}

    def handle_get_action(self, observation: Dict[str, Any], options=None):
        img = _obs_to_hwc_rgb(observation)
        state = _obs_to_state(observation)
        instruction = _extract_instruction(observation, self.instruction_override)

        pi0_obs = {
            "observation/image": img,
            "observation/state": state,
            "prompt": instruction,
        }
        out = self.policy.infer(pi0_obs)
        a = np.asarray(out["actions"], dtype=np.float32)
        # Policy.infer strips the batch dim -> (horizon, 30); accept a stray
        # leading batch dim defensively too.
        if a.ndim == 3 and a.shape[0] == 1:
            a = a[0]
        if a.ndim != 2:
            raise ValueError(f"pi0.5 actions ndim={a.ndim}, expected 2 "
                             f"(horizon, dim); shape={a.shape}")
        horizon, dim = a.shape
        if dim < 30:
            raise ValueError(f"pi0.5 action_dim={dim} < 30; cannot slice")

        action_dict: Dict[str, np.ndarray] = {}
        for key, sl in _SLICE_PLAN:
            chunk = a[:, sl]                                   # (horizon, D)
            action_dict[key] = np.ascontiguousarray(
                chunk[None, ...], dtype=np.float32)            # (1, horizon, D)

        self._reqs += 1
        if self._reqs <= 3 or self._reqs % 50 == 0:
            log.info("get_action #%d: img%s state%s '%s' -> keys=%s H=%d",
                     self._reqs, tuple(img.shape), tuple(state.shape),
                     instruction,
                     {k: tuple(v.shape) for k, v in action_dict.items()},
                     horizon)
        return [action_dict, {}]

    # ---- REP loop ---------------------------------------------------------
    def serve(self):
        ctx = zmq.Context()
        socket = ctx.socket(zmq.REP)
        socket.bind(f"tcp://0.0.0.0:{self.port}")
        dispatch = {
            "ping": (self.handle_ping, False),
            "kill": (self.handle_kill, False),
            "reset": (self.handle_reset, True),
            "get_action": (self.handle_get_action, True),
            "get_modality_config": (self.handle_get_modality_config, False),
        }
        # Watched by the console bring-up (TCP port) AND humans (log line).
        print(f"ADAPTER READY on :{self.port}", flush=True)
        log.info("ADAPTER READY on :%d (pi0.5 LoRA in-process, GR00T ZMQ wire)",
                 self.port)
        try:
            while self.running:
                message = socket.recv()
                try:
                    request = MsgSerializer.from_bytes(message)
                    if not isinstance(request, dict):
                        raise ValueError(f"request is not a dict: {type(request)}")
                    endpoint = request.get("endpoint", "get_action")
                    if endpoint not in dispatch:
                        raise ValueError(f"Unknown endpoint: {endpoint}")
                    handler, requires_input = dispatch[endpoint]
                    if requires_input:
                        result = handler(**(request.get("data") or {}))
                    else:
                        result = handler()
                    socket.send(MsgSerializer.to_bytes(result))
                except Exception as e:  # one reply per recv -- never wedge REP
                    log.error("request error: %s\n%s", e, traceback.format_exc())
                    try:
                        socket.send(MsgSerializer.to_bytes({"error": str(e)}))
                    except Exception:
                        log.error("failed to send error reply; REP may be wedged")
        finally:
            socket.close(linger=0)
            ctx.term()
            log.info("adapter stopped (served %d get_action requests)", self._reqs)


def build_argparser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument(
        "--config", type=str, default="pi05_g1_inspire_piston_lora",
        help="openpi TrainConfig name (get_config).")
    p.add_argument(
        "--ckpt_dir", type=str,
        default="/home/bir/Projects/openpi/checkpoints/"
                "pi05_g1_inspire_piston_lora/piston_v1",
        help="Checkpoint step dir (contains params/). If a run dir with "
             "integer step subdirs is given, the newest step is used.")
    p.add_argument("--port", type=int, default=5555, help="ZMQ REP bind port.")
    p.add_argument("--instruction", type=str, default="",
                   help="Override the language instruction (else obs language "
                        f"field, else {_DEFAULT_INSTRUCTION!r}).")
    return p


def main():
    logging.basicConfig(
        level=logging.INFO, force=True,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    args = build_argparser().parse_args()
    adapter = Pi05ZMQAdapter(
        config_name=args.config, ckpt_dir=args.ckpt_dir, port=args.port,
        instruction=args.instruction,
    )
    adapter.serve()


if __name__ == "__main__":
    main()