import sys
import time

from google.protobuf import timestamp_pb2

sys.path.append("../generated")

from generated.cmvr.api import humanoid_robot_command_pb2 as pb
from generated.cmvr.api import common_pb2
from generated.cmvr.api import humanoid_robot_service_pb2_grpc as rpc
from clients.base_client import RobotClientBase
import math
from typing import List, Dict, Literal


class MoveJClient(RobotClientBase):
    """Client to send MoveJ commands to the robot"""

    def send(self, joint_list, vel=1.0, acc=0.5, device_id="hc01"):
        """
        joint_list: list of dicts, e.g.,
            [
                {"joint_name": "L_SHOULDER_P", "rad": 0.0},
                {"joint_name": "L_SHOULDER_R", "rad": -1.31873},
                ...
            ]
        vel: velocity
        acc: acceleration
        device_id: device ID
        """
        # Construct JointCmd list
        # cmds = [pb.JointCmd(joint_name=j["joint_name"], rad=j["rad"], vel=vel) for j in joint_list]

        offset = {
            "R_WRIST_Y": -0.0,
        }

        cmds = [
            pb.JointCmd(
                joint_name=j["joint_name"],
                rad=j["rad"] + offset.get(j["joint_name"], 0),
                vel=vel
            )
            for j in joint_list
        ]

        # Construct request header
        header = common_pb2.CommandHeader.Request()
        header.device_id = device_id
        ts = timestamp_pb2.Timestamp()
        ts.GetCurrentTime()
        header.timestamp.CopyFrom(ts)

        # Construct MoveJ request
        req = pb.MoveJ.Request(
            header=header,
            vel=vel,
            acc=acc,
            cmds=cmds
        )

        # Call RPC
        try:
            resp = self.stub.moveJ(req, timeout=10)
            success = getattr(resp.header, "success", None)
            error_msg = getattr(resp.header, "error_message", "")
            timestamp_sec = getattr(resp.header.timestamp, "seconds", 0)
            print("MoveJ RPC call succeeded")
            print(f"Success: {success}")
            print(f"Error message: {error_msg}")
            print(f"Timestamp: {timestamp_sec}")
        except Exception as e:
            print("MoveJ RPC call failed:", e)


zero_joint_list = [

    {'joint_name': 'R_WRIST_R', 'rad': 0},
    {'joint_name': 'R_WRIST_Y', 'rad': 0},
    {'joint_name': 'R_WRIST_P', 'rad': 0},
    {'joint_name': 'R_ELBOW_R', 'rad': 0},
    {'joint_name': 'R_SHOULDER_Y', 'rad': 0},
    {'joint_name': 'R_SHOULDER_R', 'rad': 0},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0},
]

start_joint_list = [

    {'joint_name': 'R_WRIST_R', 'rad': -0.182982265197},
    {'joint_name': 'R_WRIST_Y', 'rad': 0.057211977764},
    {'joint_name': 'R_WRIST_P', 'rad': -2.450283707826},
    {'joint_name': 'R_ELBOW_R', 'rad': 1.691315387916},
    {'joint_name': 'R_SHOULDER_Y', 'rad': 1.741483963795},
    {'joint_name': 'R_SHOULDER_R', 'rad': 1.000908225412},
    {'joint_name': 'R_SHOULDER_P', 'rad': -0.275362840863},
]

end_joint_list = [

    {'joint_name': 'R_WRIST_R', 'rad': 0.590444918686},
    {'joint_name': 'R_WRIST_Y', 'rad': 0.133274073403},
    {'joint_name': 'R_WRIST_P', 'rad': -2.832957807407},
    {'joint_name': 'R_ELBOW_R', 'rad': 0.722278081411},
    {'joint_name': 'R_SHOULDER_Y', 'rad': 1.454527987188},
    {'joint_name': 'R_SHOULDER_R', 'rad': 1.319798720377},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.886766195809},
]



if __name__ == "__main__":
    # Initialize client
    client = MoveJClient()



    R_WRIST_R_joint = [
        {'joint_name': 'R_WRIST_R', 'rad': 0.0},
    ]

    # client.send(start_joint_list, 2.0, 1.0)
    # client.send(start_joint_list,2.0,2.0)
    # client.send(end_joint_list, 1.5, 1.5)
    while True:
        client.send(start_joint_list,1.5,1.5)
        client.send(end_joint_list,1.5,1.5)

    client.close()