"""Remote policy client — chiral websocket wrapper.

Used by inference (puget → russet ``rd serve``). For russet-local
scripts that drive the CAN bus directly, use :mod:`robot` instead.

Owns the websocket lifecycle, the obs stream, and the latest-obs read.
Smooth-move logic lives in :mod:`smooth_move`.

Fail-loud policy: connection failures, obs decode errors, missing keys
all propagate. No identity fallbacks, no retry loops.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import numpy as np


@dataclass
class Obs:
    """A single observation snapshot from the rig."""
    # TODO: define fields. At minimum:
    #   scene_rgb: np.ndarray (H, W, 3) uint8
    #   wrist_rgb: np.ndarray (H, W, 3) uint8
    #   scene_K:   np.ndarray (3, 3) float64
    #   wrist_K:   np.ndarray (3, 3) float64
    #   T_wrist_in_world: np.ndarray (4, 4) float64
    #   joints14:  np.ndarray (14,) float32   # [left7, right7]
    #   timestamps: scene_t, wrist_t (or id-based dedupe if timestamps are 0)
    pass


class PolicyClient:
    """Wraps the chiral PolicyClient + a few high-level conveniences."""

    def __init__(self, server_uri: str):
        # TODO: instantiate raiden chiral PolicyClient, connect, store handle.
        # TODO: start obs_stream at the given hz (param? sensible default 30).
        raise NotImplementedError

    def latest_obs(self) -> Obs:
        """Return the most recent obs snapshot, parsed into :class:`Obs`.

        Raises if the obs stream has never produced a frame. Do **not**
        return None — caller must handle missing data via control flow,
        not a magic-value sentinel.
        """
        # TODO: read env.latest_obs from chiral; build Obs; return.
        raise NotImplementedError

    def wait_for_first_obs(self, timeout_s: float = 5.0) -> Obs:
        """Block until the obs stream produces a frame (one-time at startup)."""
        # TODO: loop with short sleep until latest_obs becomes available; raise
        # TimeoutError if it doesn't arrive within timeout_s.
        raise NotImplementedError

    @property
    def joints14(self) -> np.ndarray:
        """Current 14-DoF state (left7 + right7) read from latest obs."""
        # TODO: return self.latest_obs().joints14
        raise NotImplementedError

    def smooth_move_to(
        self, target14: np.ndarray, vel_rad_s: float, min_move_s: float = 0.1
    ) -> dict:
        """Block until both arms reach target14. See :mod:`smooth_move`."""
        # TODO: delegate to lib.smooth_move.smooth_move(self, target14, vel_rad_s).
        raise NotImplementedError

    def reset(self) -> None:
        """Home both arms via env.reset()."""
        # TODO: env.reset(). Will trigger rd serve estop+shutdown on next
        # disconnect — that's intentional (handled by yam-rd-serve respawn).
        raise NotImplementedError

    def close(self) -> None:
        """Disconnect the websocket cleanly."""
        # TODO: env.stop_obs_stream() then env.close().
        raise NotImplementedError