"""YAM robot — direct raiden controller wrapper (on-russet use).

Drives both YAMs via raiden's ``RobotController`` over CAN. Runs on the
rig host (russet), no chiral / websocket involved.

The public API is :meth:`Robot.move_to` — a blocking smooth-move to a
14-DoF target at a given peak per-joint velocity. Returns when both
arms finish. Designed to be chained:

    with Robot() as robot:
        robot.move_to(pose_a, vel_rad_s=0.3)
        robot.move_to(pose_b, vel_rad_s=0.5)
        robot.move_to(pose_c, vel_rad_s=0.5)

For remote/chiral use (puget → ``rd serve``), see :mod:`policy_client`.
"""
from __future__ import annotations

from typing import Optional

import numpy as np

from raiden.robot.controller import RobotController

from .smooth_move import MoveResult, smooth_move


DOF = 7  # 6 arm joints + 1 gripper per arm


class Robot:
    """Direct raiden controller for one or both YAM arms.

    Parameters
    ----------
    use_left, use_right
        Which followers to drive. The other side's joints are ignored
        in any 14-DoF target you pass.
    gravity_comp
        Passed through to ``RobotController.initialize_robots``. Default
        False (position-control mode).
    """

    def __init__(
        self,
        use_left: bool = True,
        use_right: bool = True,
        gravity_comp: bool = False,
    ):
        if not (use_left or use_right):
            raise ValueError("must enable at least one arm")
        self.use_left = use_left
        self.use_right = use_right
        self.rc = RobotController(
            use_right_leader=False,
            use_left_leader=False,
            use_right_follower=use_right,
            use_left_follower=use_left,
        )
        self.rc.check_can_interfaces()
        self.rc.initialize_robots(gravity_comp_mode=gravity_comp)

    @property
    def follower_l(self):
        return self.rc.follower_l if self.use_left else None

    @property
    def follower_r(self):
        return self.rc.follower_r if self.use_right else None

    @property
    def joints14(self) -> np.ndarray:
        """Current 14-DoF state ``[L7, R7]``. Disabled arms read as zeros."""
        ql = self.follower_l.get_joint_pos() if self.follower_l else np.zeros(DOF)
        qr = self.follower_r.get_joint_pos() if self.follower_r else np.zeros(DOF)
        return np.concatenate([ql, qr]).astype(np.float32)

    def move_to(
        self,
        target14: np.ndarray,
        vel_rad_s: float,
        min_move_s: float = 0.2,
        steps: int = 100,
    ) -> MoveResult:
        """Smooth-move both arms to ``target14`` at the velocity budget.

        Blocks until both arms finish (parallel threads + thread.join).
        Returns a :class:`MoveResult` with planned vs actual timing.

        ``target14`` is ``[L7, R7]``. The slice for any disabled arm is
        ignored (you can pass zeros there).
        """
        return smooth_move(
            self,
            np.asarray(target14, dtype=np.float32),
            vel_rad_s=vel_rad_s,
            min_move_s=min_move_s,
            steps=steps,
        )

    def return_to_home(self) -> None:
        """Smooth-move both arms to their home pose. Blocks."""
        self.rc.return_to_home()

    def shutdown(self) -> None:
        """Release torque + close CAN. Idempotent."""
        self.rc.shutdown()

    def __enter__(self) -> "Robot":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.return_to_home()
        self.shutdown()