from __future__ import annotations

from typing import Iterable, Mapping

from clients._path_setup import ensure_paths

ensure_paths()

from google.protobuf import timestamp_pb2

from clients.base_client import RobotClientBase
from cmvr.api import arm_command_pb2 as pb
from cmvr.api import common_pb2


FRAME_BY_NAME = {
    "base": pb.ARM_FRAME_BASE,
    "tool": pb.ARM_FRAME_TOOL,
    "world": pb.ARM_FRAME_WORLD,
    "user": pb.ARM_FRAME_USER,
}


class ArmMotionClient(RobotClientBase):
    def _header(self, device_id: str) -> common_pb2.CommandHeader.Request:
        header = common_pb2.CommandHeader.Request()
        header.device_id = device_id
        ts = timestamp_pb2.Timestamp()
        ts.GetCurrentTime()
        header.timestamp.CopyFrom(ts)
        return header

    def _frame(self, frame: str) -> int:
        return FRAME_BY_NAME.get(frame.strip().lower(), pb.ARM_FRAME_BASE)

    def _print_feedback(self, action: str, response) -> None:
        header = response.header
        print(f"{action} RPC call succeeded")
        print(f"Success: {getattr(header, 'success', None)}")
        print(f"Error message: {getattr(header, 'error_message', '')}")
        print(f"Timestamp: {getattr(header.timestamp, 'seconds', 0)}")

    def movej(
        self,
        joint_list: Iterable[Mapping[str, float]],
        vel: float = 1.0,
        acc: float = 0.5,
        device_id: str = "right_arm",
    ) -> None:
        target = pb.JointPositionCommand(position=[float(j["rad"]) for j in joint_list])
        req = pb.MoveJ.Request(
            header=self._header(device_id),
            target=target,
            options=pb.MotionOptions(velocity=vel, acceleration=acc),
        )
        self._print_feedback("MoveJ", self.stub.moveJ(req, timeout=10))

    def movel(
        self,
        pose: Mapping[str, float],
        vel: float = 0.2,
        acc: float = 0.2,
        blend_radius: float = 0.0,
        frame: str = "base",
        device_id: str = "right_arm",
    ) -> None:
        req = pb.MoveL.Request(
            header=self._header(device_id),
            target=pb.CartesianPose(
                x=float(pose["x"]),
                y=float(pose["y"]),
                z=float(pose["z"]),
                rx=float(pose["rx"]),
                ry=float(pose["ry"]),
                rz=float(pose["rz"]),
            ),
            options=pb.MotionOptions(
                velocity=vel,
                acceleration=acc,
                blend_radius=blend_radius,
            ),
            frame=self._frame(frame),
        )
        self._print_feedback("MoveL", self.stub.moveL(req, timeout=10))

    def speedj(
        self,
        velocity_list: Iterable[Mapping[str, float]],
        acc: float = 0.5,
        duration: float = 0.2,
        device_id: str = "right_arm",
    ) -> None:
        velocity = pb.JointVelocityCommand(
            velocity=[float(j["velocity"]) for j in velocity_list]
        )
        req = pb.SpeedJ.Request(
            header=self._header(device_id),
            velocity=velocity,
            acceleration=acc,
            duration=duration,
        )
        self._print_feedback("SpeedJ", self.stub.speedJ(req, timeout=10))

    def speedl(
        self,
        velocity: Mapping[str, float],
        acc: float = 0.2,
        duration: float = 0.2,
        frame: str = "base",
        device_id: str = "right_arm",
    ) -> None:
        req = pb.SpeedL.Request(
            header=self._header(device_id),
            velocity=pb.CartesianVelocity(
                vx=float(velocity["vx"]),
                vy=float(velocity["vy"]),
                vz=float(velocity["vz"]),
                wx=float(velocity["wx"]),
                wy=float(velocity["wy"]),
                wz=float(velocity["wz"]),
            ),
            acceleration=acc,
            duration=duration,
            frame=self._frame(frame),
        )
        self._print_feedback("SpeedL", self.stub.speedL(req, timeout=10))

    def stop_motion(self, device_id: str = "right_arm") -> None:
        req = self._header(device_id)
        resp = self.stub.stopMotion(req, timeout=10)
        print("StopMotion RPC call succeeded")
        print(f"Success: {getattr(resp, 'success', None)}")
        print(f"Error message: {getattr(resp, 'error_message', '')}")
        print(f"Timestamp: {getattr(resp.timestamp, 'seconds', 0)}")