#!/usr/bin/env python3
"""TCP <-> WebSocket bridge: legacy Quantum-1 robot TCP protocol to
Wall-OSS-0.5 official WebSocket serving.

Architecture:
    Quantum-1 robot (legacy TCP client, e.g. `infer ip port`)
        <-> legacy TCP on this bridge (default 30123)
    this bridge
        <-> Wall-OSS-0.5 WebSocket serving (default ws://127.0.0.1:32195)

Legacy TCP protocol (from robot_controller.py):
    robot -> bridge : [u32 len][state json], then 3x [u32 len][jpeg bytes]
                      (camera_left, camera_front, camera_right)
    bridge -> robot : [u32 len][action json dict]

Safety: default DRY-RUN. It converts and runs inference but does NOT send
actions back to the robot. Physical sending is gated by BOTH --allow-send
and ACTION_SEMANTICS_CONFIRMED (source-code flag) so a CLI typo alone cannot
enable motion.
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import json
import logging
import struct
import sys

import numpy as np
import cv2
import msgpack
import msgpack_numpy as _m
import websockets
from scipy.spatial.transform import Rotation

_m.patch()

# Legacy wire order of the three cameras sent by the robot.
LEGACY_CAM_ORDER = ("camera_left", "camera_front", "camera_right")
# Forward all three legacy camera streams.  The serving train config decides
# which of these it consumes (X2Robot commonly uses all three).
SERVE_CAM_KEYS = LEGACY_CAM_ORDER

# Keep False until the arm action semantics (absolute vs relative) and the
# camera wire order have been confirmed against a real robot capture.
ACTION_SEMANTICS_CONFIRMED = True

logger = logging.getLogger("bridge")


async def _recvall(reader, n):
    buf = b""
    while len(buf) < n:
        chunk = await reader.read(n - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf


async def _recv_frame(reader):
    size_b = await _recvall(reader, 4)
    if size_b is None:
        return None
    size = struct.unpack("<L", size_b)[0]
    return await _recvall(reader, size)


async def _recv_state(reader):
    raw = await _recv_frame(reader)
    if raw is None:
        return None
    state = json.loads(raw.decode("utf-8"))
    if not isinstance(state, dict):
        raise ValueError("robot state is not a JSON object")
    for key in ("follow1_pos", "follow2_pos"):
        if key not in state:
            raise ValueError(f"robot state missing required key {key!r}")
    return state


async def _recv_image(reader, index):
    raw = await _recv_frame(reader)
    if raw is None:
        raise ConnectionError("robot closed during image stream")
    arr = np.frombuffer(raw, np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(f"failed to decode image #{index}")
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


def _img_to_b64(rgb):
    ok, buf = cv2.imencode(".jpg", cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))
    if not ok:
        raise ValueError("failed to JPEG-encode camera image")
    return base64.b64encode(buf.tobytes()).decode("ascii")


def _as_action_chunk(value, key):
    try:
        arr = np.asarray(value, dtype=np.float64)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{key} must be numeric") from exc
    if arr.ndim != 2 or arr.shape[1] != 7:
        raise ValueError(f"{key} must have shape (T,7), got {arr.shape}")
    if not np.isfinite(arr).all():
        raise ValueError(f"{key} contains NaN/Inf")
    return arr


def _single_arm_action(ws_action, current_right=None):
    """Decode a right-arm-only raw model response into a (T, 7) chunk."""
    if "predict_action" not in ws_action:
        raise ValueError("WebSocket response missing predict_action")
    raw = np.asarray(ws_action["predict_action"], dtype=np.float64)
    if raw.ndim == 3:
        if raw.shape[0] != 1:
            raise ValueError(f"predict_action batch must have size 1, got {raw.shape}")
        raw = raw[0]
    if raw.ndim != 2 or not np.isfinite(raw).all():
        raise ValueError(f"predict_action must be finite with shape (T,D), got {raw.shape}")
    # Training layout: 16 padding values, then position(3), rotation-6D(6),
    # and gripper(1). Convert the rotation representation without importing
    # the LIBERO/MuJoCo package, which is unavailable on headless deployments.
    if raw.shape[1] < 26:
        raise ValueError(f"predict_action has {raw.shape[1]} dims; expected at least 26")
    pos = raw[:, 16:19]
    rot6d = raw[:, 19:25].reshape(-1, 2, 3)
    first = rot6d[:, 0]
    second = rot6d[:, 1]
    first = first / np.maximum(np.linalg.norm(first, axis=1, keepdims=True), 1e-12)
    second = second - np.sum(first * second, axis=1, keepdims=True) * first
    second = second / np.maximum(np.linalg.norm(second, axis=1, keepdims=True), 1e-12)
    third = np.cross(first, second)
    matrices = np.stack((first, second, third), axis=-1)
    rot = Rotation.from_matrix(matrices).as_euler("zyx")
    if current_right is not None:
        current = np.asarray(current_right, dtype=np.float64).reshape(7)
        if not np.isfinite(current).all():
            raise ValueError("state.follow2_pos contains NaN/Inf")
        # Training keys are explicitly *_relative: compose the predicted
        # translation and rotation with the robot's current EE pose.
        pos = current[None, :3] + pos
        from wall_x._vendor.x2robot_utils import geometry as geom
        rot = geom.compose_state_and_delta_to_abs_rpy(
            raw[:, 19:25], current[3:6]
        )
    grip = raw[:, 25:26]
    right = np.concatenate((pos, rot, grip), axis=1)
    return _as_action_chunk(right, "predict_action.right_arm")


def _validate_optional_series(ws_action, key, length, width=None):
    """Validate an optional serialized trajectory and return JSON-safe lists."""
    if key not in ws_action:
        return None
    value = np.asarray(ws_action[key], dtype=np.float64)
    if value.ndim == 1:
        value = value[:, None]
    if value.ndim != 2 or (width is not None and value.shape[1] != width):
        raise ValueError(f"{key} has invalid shape {value.shape}")
    if value.shape[0] != length or not np.isfinite(value).all():
        raise ValueError(f"{key} has invalid length or non-finite values")
    return value.tolist()


def _legacy_resample(values, key, end_ratio, interpolate_multiplier):
    """Apply the legacy infer.py trim + linear interpolation policy."""
    arr = np.asarray(values, dtype=np.float64)
    if arr.ndim == 1:
        arr = arr[:, None]
    if arr.ndim != 2 or arr.shape[0] == 0:
        raise ValueError(f"{key} has invalid shape {arr.shape}")
    end = int(end_ratio * arr.shape[0])
    if end <= 0:
        raise ValueError(
            f"{key} trim is empty: end_ratio={end_ratio} length={arr.shape[0]}"
        )
    trimmed = arr[:end]
    target_length = int(interpolate_multiplier * len(trimmed))
    if target_length <= 0:
        raise ValueError(f"{key} interpolation produced no frames")
    if target_length == len(trimmed):
        return trimmed.tolist()
    source_idx = np.linspace(0, len(trimmed) - 1, len(trimmed))
    target_idx = np.linspace(0, len(trimmed) - 1, target_length)
    out = np.empty((target_length, arr.shape[1]), dtype=np.float64)
    for col in range(arr.shape[1]):
        out[:, col] = np.interp(target_idx, source_idx, trimmed[:, col])
    return out.tolist()


def _limit_action_packet(traj, current, position_limit, rotation_limit):
    """Scale one absolute action packet without changing its path shape.

    The limits apply to both the packet's excursion from the observed pose and
    its largest inter-sample step.  A single scale factor per position/rotation
    group keeps the model's direction and timing intact; independently clipping
    each sample would turn a large target into an artificial diagonal ramp.
    """
    limited = np.asarray(traj, dtype=np.float64).copy()
    current = np.asarray(current, dtype=np.float64).reshape(7)
    if limited.ndim != 2 or limited.shape[1] != 7:
        raise ValueError(f"right-arm trajectory must have shape (T,7), got {limited.shape}")
    if not np.isfinite(limited).all() or not np.isfinite(current).all():
        raise ValueError("right-arm trajectory/current contains NaN/Inf")
    if position_limit <= 0 or rotation_limit <= 0:
        raise ValueError("action delta limits must be positive")

    scales = {}
    for name, columns, limit in (
        ("position", slice(0, 3), float(position_limit)),
        ("rotation", slice(3, 6), float(rotation_limit)),
    ):
        delta = limited[:, columns] - current[columns]
        peak = float(np.max(np.abs(delta))) if delta.size else 0.0
        if len(delta) > 1:
            peak = max(peak, float(np.max(np.abs(np.diff(delta, axis=0)))))
        scale = min(1.0, limit / peak) if peak > 0.0 else 1.0
        if scale < 1.0:
            limited[:, columns] = current[columns] + delta * scale
        scales[name] = (peak, scale)
    return limited, scales


def assess_right_arm_feedback(
    initial_pose,
    observed_pose,
    commanded_pose,
    *,
    min_position_delta=0.0005,
    min_rotation_delta=0.002,
):
    """Compare the next robot state to the preceding right-arm command."""
    initial = np.asarray(initial_pose, dtype=np.float64).reshape(7)
    observed = np.asarray(observed_pose, dtype=np.float64).reshape(7)
    commanded = np.asarray(commanded_pose, dtype=np.float64).reshape(7)
    if not (
        np.isfinite(initial).all()
        and np.isfinite(observed).all()
        and np.isfinite(commanded).all()
    ):
        raise ValueError("right-arm feedback poses must be finite 7D vectors")
    commanded_delta = commanded - initial
    observed_delta = observed - initial
    command_requests_motion = (
        np.max(np.abs(commanded_delta[:3])) >= min_position_delta
        or np.max(np.abs(commanded_delta[3:6])) >= min_rotation_delta
    )
    observed_motion = (
        np.max(np.abs(observed_delta[:3])) >= min_position_delta
        or np.max(np.abs(observed_delta[3:6])) >= min_rotation_delta
    )
    return {
        "commanded_delta": commanded_delta,
        "observed_delta": observed_delta,
        "command_requests_motion": command_requests_motion,
        "observed_motion": observed_motion,
        "missing_feedback": command_requests_motion and not observed_motion,
    }


def prepare_robot_actions(
    ws_action,
    *,
    state_follow1_pos=None,
    state_follow2_pos=None,
    state_head_pos=None,
    state_lift=None,
    state_car_pose=None,
    action_horizon=32,
    action_end_ratio=0.2,
    action_interpolate_multiplier=32,
    allow_base_motion=False,
    hold_left_arm=True,
    fixed_car_pose=None,
    fixed_lift=None,
    allow_constant_fallbacks=False,
    max_position_delta=0.10,
    max_rotation_delta=0.50,
    clip_action_delta=False,
    gripper_min=None,
    gripper_max=None,
):
    """Validate Wall-OSS-0.5 serialized actions for the legacy robot.

    The official Turtle serializer emits head/lift/base trajectories.  The
    bridge preserves head/lift, but holds the base at the robot-reported pose
    by default for safety.  Constant fallbacks are available only for dry-run
    compatibility tests and must be explicitly requested.
    """
    if not isinstance(ws_action, dict):
        raise ValueError("WebSocket action response is not a dict")
    if "follow1_pos" in ws_action and "follow2_pos" in ws_action:
        right = _as_action_chunk(ws_action["follow2_pos"], "follow2_pos")
        if hold_left_arm:
            if state_follow1_pos is None:
                raise ValueError("state_follow1_pos is required when holding the left arm")
            current_left = _as_action_chunk(
                np.asarray(state_follow1_pos, dtype=np.float64).reshape(1, 7),
                "state.follow1_pos",
            )[0]
            # A right-arm-only checkpoint can still receive a synthetic
            # follow1_pos from the generic serializer. Never actuate it.
            left = np.repeat(current_left[None, :], right.shape[0], axis=0)
        else:
            left = _as_action_chunk(ws_action["follow1_pos"], "follow1_pos")
    elif "follow2_pos" in ws_action:
        # Single-arm serving responses contain the right arm reconstructed by
        # Wall-X's official preprocessor. Hold the left arm at its current pose.
        right = _as_action_chunk(ws_action["follow2_pos"], "follow2_pos")
        if state_follow1_pos is None:
            raise ValueError("state_follow1_pos is required for single-arm responses")
        current_left = _as_action_chunk(
            np.asarray(state_follow1_pos, dtype=np.float64).reshape(1, 7),
            "state.follow1_pos",
        )[0]
        left = np.repeat(current_left[None, :], right.shape[0], axis=0)
    else:
        raise ValueError(
            "Turtle2 requires a serialized follow2_pos trajectory; "
            "start serving with --env X2ROBOT --robot-type turtle "
            "--serialize-actions"
        )
    if left.shape[0] != right.shape[0]:
        raise ValueError("follow1_pos/follow2_pos trajectory length mismatch")
    serialized_length = int(left.shape[0])
    source_length = serialized_length
    # TurtleRobotPreprocessor stacks the observed state before the H model
    # actions, then interpolates the result.  Strip that one state row before
    # applying the legacy infer.py timing policy, so its 20% window starts at
    # the first predicted action rather than repeating the current pose.
    if action_horizon <= 0:
        raise ValueError("action_horizon must be positive")
    if source_length == action_horizon + 1:
        left = left[1:]
        right = right[1:]
        source_length -= 1
    elif source_length != action_horizon:
        raise ValueError(
            f"unexpected serialized action length {source_length}; expected "
            f"{action_horizon} or {action_horizon + 1}"
        )
    if not (0 < action_end_ratio <= 1):
        raise ValueError("action_end_ratio must be in (0, 1]")
    selected_length = int(action_end_ratio * source_length)
    if selected_length <= 0:
        raise ValueError(
            f"action trim is empty: end_ratio={action_end_ratio} "
            f"length={source_length}"
        )
    logger.info(
        "executing first %d of %d predicted action steps before interpolation",
        selected_length,
        source_length,
    )
    left = left[:selected_length]
    right = right[:selected_length]
    if action_interpolate_multiplier < 1:
        raise ValueError("action_interpolate_multiplier must be >= 1")
    actions = {
        "follow1_pos": _legacy_resample(
            left, "follow1_pos", 1.0, action_interpolate_multiplier
        ),
        "follow2_pos": _legacy_resample(
            right, "follow2_pos", 1.0, action_interpolate_multiplier
        ),
    }
    for key, width, fallback in (
        ("head_pos", 2, [[0.0, -1.0] for _ in range(source_length)]),
        ("lift", 1, [0.4 for _ in range(source_length)]),
        ("car_pose", 3, [[0.0, 0.0, 0.0] for _ in range(source_length)]),
    ):
        values = _validate_optional_series(ws_action, key, serialized_length, width)
        if values is None:
            state_value = {"head_pos": state_head_pos, "lift": state_lift}.get(key)
            if state_value is not None:
                values = np.repeat(np.asarray(state_value, dtype=np.float64).reshape(1, -1), serialized_length, axis=0).tolist()
            elif key == "car_pose" and state_car_pose is not None:
                values = np.repeat(np.asarray(state_car_pose, dtype=np.float64).reshape(1, -1), serialized_length, axis=0).tolist()
            elif not allow_constant_fallbacks:
                raise ValueError(
                    f"WebSocket response missing {key!r}; start serving with "
                    "--env X2ROBOT --robot-type turtle --serialize-actions"
                )
            values = fallback
        elif len(values) == action_horizon + 1:
            values = values[1:]
        elif len(values) != action_horizon:
            raise ValueError(
                f"{key} has unexpected serialized length {len(values)}; expected "
                f"{action_horizon} or {action_horizon + 1}"
            )
        values = values[:selected_length]
        # Preserve the old timing for non-base trajectories as well.
        actions[key] = _legacy_resample(
            values, key, 1.0, action_interpolate_multiplier
        )

    # The official Turtle serializer represents lift as (T, 1), but the
    # legacy Turtle2 receiver assigns each row directly to ``lift_cmd`` and
    # expects a scalar.  Keep the wire format compatible with that receiver.
    actions["lift"] = [float(row[0]) for row in actions["lift"]]
    if fixed_lift is not None:
        fixed_lift = float(fixed_lift)
        if not np.isfinite(fixed_lift) or not 0.0 <= fixed_lift <= 0.47:
            raise ValueError("fixed_lift must be within Turtle2 range [0.0, 0.47]")
        actions["lift"] = [fixed_lift for _ in actions["lift"]]

    if fixed_car_pose is not None:
        pose = np.asarray(fixed_car_pose, dtype=np.float64).reshape(-1)
        if pose.shape != (3,) or not np.isfinite(pose).all():
            raise ValueError("fixed_car_pose must have shape (3,) and finite values")
        # Turtle2 car_pose is a three-value [x, y, yaw] target.
        actions["car_pose"] = [pose.tolist() for _ in range(len(actions["follow1_pos"]))]
    elif not allow_base_motion:
        # Turtle2 converts pose commands through relative_pose_to_absolute_pose
        # before calling set_target_pose().  Sending the reported absolute pose
        # again would therefore be interpreted as a relative displacement.
        # Zero is the no-motion command; keep the observed pose only in logs.
        actions["car_pose"] = [[0.0, 0.0, 0.0] for _ in actions["follow1_pos"]]
    # The checkpoint-1 normalizer records the gripper in the robot's 0..4.5
    # units. Invalid negative values can make Turtle2 reject the whole command
    # packet, so constrain only this scalar channel before sending.
    gripper = np.asarray(actions["follow2_pos"], dtype=np.float64)
    if gripper_min is not None and gripper_max is not None:
        if not np.isfinite([gripper_min, gripper_max]).all() or gripper_min > gripper_max:
            raise ValueError("invalid gripper limits")
        before = gripper[:, 6].copy()
        gripper[:, 6] = np.clip(gripper[:, 6], float(gripper_min), float(gripper_max))
        if not np.array_equal(before, gripper[:, 6]):
            logger.warning(
                "clipped gripper command range from [%.5g, %.5g] to [%.5g, %.5g]",
                float(np.min(before)), float(np.max(before)),
                float(np.min(gripper[:, 6])), float(np.max(gripper[:, 6])),
            )
        actions["follow2_pos"] = gripper.tolist()

    if state_follow2_pos is not None:
        current = np.asarray(state_follow2_pos, dtype=np.float64).reshape(7)
        traj = np.asarray(actions["follow2_pos"], dtype=np.float64)
        if clip_action_delta:
            # Limit the complete packet relative to the observed pose while
            # preserving the model trajectory's shape.  Per-sample cumulative
            # clipping creates a long synthetic ramp toward an unreachable
            # pose and was the source of the repeatable fixed-pose stop.
            clipped, scales = _limit_action_packet(
                traj, current, max_position_delta, max_rotation_delta
            )
            if scales["position"][1] < 1.0 or scales["rotation"][1] < 1.0:
                logger.debug(
                    "scaled right-arm packet: position peak %.5g scale %.5g; "
                    "rotation peak %.5g scale %.5g",
                    scales["position"][0], scales["position"][1],
                    scales["rotation"][0], scales["rotation"][1],
                )
            actions["follow2_pos"] = clipped.tolist()
            traj = clipped
        first_delta = np.abs(traj[0, :3] - current[:3])
        limit_eps = 1e-6
        if np.any(first_delta > max_position_delta + limit_eps):
            raise ValueError(
                f"right-arm first position delta {first_delta.tolist()} exceeds "
                f"limit {max_position_delta} m"
            )
        first_rot = np.abs(traj[0, 3:6] - current[3:6])
        if np.any(first_rot > max_rotation_delta + limit_eps):
            raise ValueError(
                f"right-arm first rotation delta {first_rot.tolist()} exceeds "
                f"limit {max_rotation_delta} rad"
            )
        if len(traj) > 1:
            step_pos = np.max(np.abs(np.diff(traj[:, :3], axis=0)), axis=1)
            step_rot = np.max(np.abs(np.diff(traj[:, 3:6], axis=0)), axis=1)
            if np.any(step_pos > max_position_delta + limit_eps) or np.any(step_rot > max_rotation_delta + limit_eps):
                if clip_action_delta:
                    # A trajectory can contain Euler-angle wrap discontinuities
                    # after preprocessing. In clip mode keep the connection
                    # alive and clamp the offending samples to the prior pose.
                    logger.warning(
                        "right-arm trajectory still has over-limit step after clipping; "
                        "holding offending samples"
                    )
                    safe = traj.copy()
                    for i in range(1, len(safe)):
                        dp = safe[i, :3] - safe[i - 1, :3]
                        dr = safe[i, 3:6] - safe[i - 1, 3:6]
                        if np.any(np.abs(dp) > max_position_delta) or np.any(
                            np.abs(dr) > max_rotation_delta
                        ):
                            safe[i] = safe[i - 1]
                    actions["follow2_pos"] = safe.tolist()
                else:
                    raise ValueError("right-arm trajectory contains an over-limit step")
    return actions


async def _handle_client(
    reader,
    writer,
    ws_url,
    instruction,
    allow_send,
    allow_constant_fallbacks,
    action_horizon,
    action_end_ratio,
    action_interpolate_multiplier,
    allow_base_motion,
    hold_left_arm,
    fixed_car_pose,
    fixed_lift,
    max_action_cycles,
    require_right_feedback,
    clip_action_delta,
    max_position_delta,
    max_rotation_delta,
    gripper_min=None,
    gripper_max=None,
):
    addr = writer.get_extra_info("peername")
    logger.info("Robot connected: %s", addr)

    async with websockets.connect(ws_url, max_size=None) as ws:
        await ws.recv()  # consume server metadata

        # One observation per connection in dry-run; continuous loop when sending.
        cycles = 0
        previous_right_plan = None
        stop_after_feedback = False
        while True:
            state = await _recv_state(reader)
            if state is None:
                logger.info("Robot closed the connection")
                return
            if previous_right_plan is not None:
                feedback = assess_right_arm_feedback(
                    previous_right_plan["initial"],
                    state["follow2_pos"],
                    previous_right_plan["target"],
                )
                logger.info(
                    "right-arm feedback: observed_delta=%s commanded_delta=%s",
                    feedback["observed_delta"].tolist(),
                    feedback["commanded_delta"].tolist(),
                )
                if feedback["missing_feedback"]:
                    logger.warning(
                        "right-arm feedback did not change after the prior packet; "
                        "bridge sent follow2_pos but the robot did not report EE motion. "
                        "Check /follow_pos_cmd_2 subscribers, controller enable state, "
                        "/follow2_pos_back, and /joint_information2 on arm-pc."
                    )
                    if require_right_feedback:
                        logger.error(
                            "--require-right-feedback set; "
                            "stopping before another action packet"
                        )
                        return
                if stop_after_feedback:
                    logger.info("right-arm feedback check complete; stopping after one action packet")
                    return
            previous_right_plan = None
            # Turtle2 does not transmit velocity_decomposed.  The official
            # Turtle preprocessor accepts the field, so provide a neutral
            # initial value rather than failing on a missing state key.  This
            # is deliberately conservative until recurrent velocity semantics
            # are confirmed against a real checkpoint/robot run.
            state.setdefault("velocity_decomposed", [0.0, 0.0, 0.0])

            images = [await _recv_image(reader, i) for i in range(len(LEGACY_CAM_ORDER))]
            left_img, front_img, right_img = images
            logger.info(
                "recv state keys=%s images=%sx%s",
                list(state.keys()),
                front_img.shape[:2],
                right_img.shape[:2],
            )

            obs = {
                "state": dict(state),
                "views": {
                    SERVE_CAM_KEYS[0]: _img_to_b64(left_img),
                    SERVE_CAM_KEYS[1]: _img_to_b64(front_img),
                    SERVE_CAM_KEYS[2]: _img_to_b64(right_img),
                },
                "instruction": instruction,
                "infer_mode": "flow",
            }
            logger.info("forwarding to serving: %s", ws_url)
            await ws.send(msgpack.packb(obs))
            raw_resp = await ws.recv()
            if isinstance(raw_resp, str):
                # The official server sends a traceback as a text frame before
                # closing when inference fails. Preserve that diagnostic
                # instead of masking it with msgpack's bytes-only error.
                raise RuntimeError(
                    "serving returned a text error frame:\n" + raw_resp
                )
            resp = msgpack.unpackb(raw_resp)

            actions = prepare_robot_actions(
                resp,
                state_follow1_pos=state.get("follow1_pos"),
                state_follow2_pos=state.get("follow2_pos"),
                state_head_pos=state.get("head_pos"),
                state_lift=state.get("lift"),
                state_car_pose=state.get("car_pose"),
                action_horizon=action_horizon,
                action_end_ratio=action_end_ratio,
                action_interpolate_multiplier=action_interpolate_multiplier,
                allow_base_motion=allow_base_motion,
                hold_left_arm=hold_left_arm,
                fixed_car_pose=fixed_car_pose,
                fixed_lift=fixed_lift,
                allow_constant_fallbacks=allow_constant_fallbacks,
                clip_action_delta=clip_action_delta,
                max_position_delta=max_position_delta,
                max_rotation_delta=max_rotation_delta,
                gripper_min=gripper_min,
                gripper_max=gripper_max,
            )
            logger.info(
                "predicted follow1 T=%d last=%s", len(actions["follow1_pos"]), actions["follow1_pos"][-1]
            )
            logger.info(
                "predicted follow2 T=%d last=%s", len(actions["follow2_pos"]), actions["follow2_pos"][-1]
            )
            try:
                current_right = np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7)
                first_right = np.asarray(actions["follow2_pos"][0], dtype=np.float64)
                last_right = np.asarray(actions["follow2_pos"][-1], dtype=np.float64)
                logger.info(
                    "right-arm current=%s first=%s first_delta=%s last_delta=%s",
                    current_right.tolist(), first_right.tolist(),
                    (first_right - current_right).tolist(),
                    (last_right - current_right).tolist(),
                )
            except (KeyError, ValueError):
                logger.warning("could not compute right-arm current-to-first delta")

            if not allow_send:
                logger.warning("DRY-RUN: action NOT sent to robot")
                return
            if not ACTION_SEMANTICS_CONFIRMED:
                logger.error("ACTION_SEMANTICS_CONFIRMED=False; refusing to send to robot")
                return

            logger.info(
                "wire action T=%d left[first,last]=%s/%s right[first,last]=%s/%s "
                "lift=%s car_pose=%s",
                len(actions["follow1_pos"]),
                actions["follow1_pos"][0], actions["follow1_pos"][-1],
                actions["follow2_pos"][0], actions["follow2_pos"][-1],
                actions["lift"][0], actions["car_pose"][0],
            )
            payload = json.dumps(actions).encode("utf-8")
            writer.write(struct.pack("<L", len(payload)))
            writer.write(payload)
            await writer.drain()
            previous_right_plan = {
                "initial": np.asarray(state["follow2_pos"], dtype=np.float64).reshape(7),
                "target": np.asarray(actions["follow2_pos"][-1], dtype=np.float64).reshape(7),
            }
            logger.info("sent action back to robot T=%d", len(actions["follow1_pos"]))
            cycles += 1
            if max_action_cycles > 0 and cycles >= max_action_cycles:
                if require_right_feedback:
                    stop_after_feedback = True
                    logger.info(
                        "waiting for one robot feedback state before stopping after %d action cycle(s)",
                        cycles,
                    )
                    continue
                logger.warning("stopping after %d action cycle(s)", cycles)
                return


async def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tcp-host", default="0.0.0.0")
    parser.add_argument("--tcp-port", type=int, default=30123)
    parser.add_argument("--ws-url", default="ws://127.0.0.1:32195")
    parser.add_argument("--instruction", default="pick up the water bottel on the chair")
    parser.add_argument(
        "--action-horizon",
        type=int,
        default=32,
        help="Model action horizon (Wall-OSS-0.5 default: 32).",
    )
    parser.add_argument(
        "--allow-send",
        action="store_true",
        help="Request sending actions back to robot (still gated by ACTION_SEMANTICS_CONFIRMED).",
    )
    parser.add_argument(
        "--allow-constant-fallbacks",
        action="store_true",
        help=(
            "Use neutral head/lift/base trajectories when the serving response "
            "omits them. Intended only for protocol dry-runs."
        ),
    )
    parser.add_argument(
        "--action-end-ratio",
        type=float,
        default=0.2,
        help="Legacy action trim ratio (default: 0.2).",
    )
    parser.add_argument(
        "--action-interpolate-multiplier",
        type=int,
        default=32,
        help="Legacy linear interpolation multiplier (default: 32).",
    )
    parser.add_argument(
        "--allow-base-motion",
        action="store_true",
        help="Forward model-predicted car_pose instead of holding current pose.",
    )
    parser.add_argument(
        "--allow-left-arm-motion",
        action="store_true",
        help="Forward serializer-provided left-arm commands. Disabled by default for this right-arm-only checkpoint.",
    )
    parser.add_argument(
        "--fixed-car-pose",
        type=float,
        nargs=3,
        metavar=("X", "Y", "YAW"),
        default=None,
        help=(
            "Override every output car_pose with [X, Y, YAW]. "
            "Use only after confirming Turtle2 relative-target semantics."
        ),
    )
    parser.add_argument(
        "--fixed-lift",
        type=float,
        default=None,
        help="Force lift target for every action frame (meters, range 0.0-0.47).",
    )
    parser.add_argument(
        "--max-action-cycles", type=int, default=1,
        help="Maximum action packets per connection; 0 means unlimited.",
    )
    parser.add_argument(
        "--require-right-feedback",
        action="store_true",
        help="Stop before the next packet if the previous right-arm command produced no observed EE motion.",
    )
    parser.add_argument(
        "--clip-action-delta", action="store_true",
        help=(
            "Scale each right-arm action packet to safety limits instead of "
            "rejecting it; limits apply to packet excursion and inter-sample steps."
        ),
    )
    parser.add_argument(
        "--max-position-delta", type=float, default=0.10,
        help="Maximum right-arm position excursion per inference packet (m).",
    )
    parser.add_argument(
        "--max-rotation-delta", type=float, default=0.50,
        help="Maximum right-arm Euler rotation excursion per inference packet (rad).",
    )
    parser.add_argument(
        "--gripper-min", type=float, default=None,
        help="Optional gripper lower bound; disabled by default.",
    )
    parser.add_argument(
        "--gripper-max", type=float, default=None,
        help="Optional gripper upper bound; disabled by default.",
    )
    parser.add_argument("--log-level", default="INFO")
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s %(levelname)s %(name)s %(message)s",
    )

    async def client_connected(reader, writer):
        try:
            await _handle_client(
                reader,
                writer,
                args.ws_url,
                args.instruction,
                args.allow_send,
                args.allow_constant_fallbacks,
                args.action_horizon,
                args.action_end_ratio,
                args.action_interpolate_multiplier,
                args.allow_base_motion,
                not args.allow_left_arm_motion,
                args.fixed_car_pose,
                args.fixed_lift,
                args.max_action_cycles,
                args.require_right_feedback,
                args.clip_action_delta,
                args.max_position_delta,
                args.max_rotation_delta,
                args.gripper_min,
                args.gripper_max,
            )
        except Exception as exc:
            logger.error("handler error: %s: %s", type(exc).__name__, exc)
        finally:
            try:
                writer.close()
                await writer.wait_closed()
            except Exception:
                pass
            logger.info("robot session closed")

    server = await asyncio.start_server(client_connected, args.tcp_host, args.tcp_port)
    logger.info(
        "bridge listening on %s:%s -> %s (allow_send=%s)",
        args.tcp_host, args.tcp_port, args.ws_url, args.allow_send,
    )
    async with server:
        await server.serve_forever()


if __name__ == "__main__":
    try:
        raise SystemExit(asyncio.run(main()))
    except KeyboardInterrupt:
        sys.exit(0)