import math
import time
from typing import Dict, List, Literal

from clients._path_setup import ensure_paths

ensure_paths()

from google.protobuf import timestamp_pb2

from clients.base_client import RobotClientBase
from clients.hand_client import DexHandClient
from cmvr.api import common_pb2
from cmvr.api import humanoid_robot_command_pb2 as pb
from cmvr.api import humanoid_robot_service_pb2_grpc as rpc



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]

        # 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)


# init_joint_list = [
# {'joint_name': 'R_SHOULDER_P', 'rad': -0.003768504782},
# {'joint_name': 'R_SHOULDER_R', 'rad': -0.256264020657},
# {'joint_name': 'R_SHOULDER_Y', 'rad': 2.526736892626},
# {'joint_name': 'R_ELBOW_R', 'rad': -1.542952107456},
# {'joint_name': 'R_WRIST_P', 'rad': 0.117851681162},
# {'joint_name': 'R_WRIST_Y', 'rad': 0.097711538603},
# {'joint_name': 'R_WRIST_R', 'rad': 0.021115966972},
# ]

init_joint_list = [
{'joint_name': 'R_SHOULDER_P', 'rad': -0.062633119026},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.013667237243},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.563999526183},
{'joint_name': 'R_ELBOW_R', 'rad': -0.275263170082},
{'joint_name': 'R_WRIST_P', 'rad': 0.21726806475},
{'joint_name': 'R_WRIST_Y', 'rad': -0.044427538871},
{'joint_name': 'R_WRIST_R', 'rad': -0.065047050427},
]

start_joint_list = [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.212733518819},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.391965314899},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.564591855398},
{'joint_name': 'R_ELBOW_R', 'rad': -1.503521397093},
{'joint_name': 'R_WRIST_P', 'rad': 0.217252876821},
{'joint_name': 'R_WRIST_Y', 'rad': 0.164711187099},
{'joint_name': 'R_WRIST_R', 'rad': 0.227967960443},
]

air_in_joint_list= [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.467093352929},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.397918033661},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.332349442301},
{'joint_name': 'R_ELBOW_R', 'rad': -1.108740619931},
{'joint_name': 'R_WRIST_P', 'rad': 0.184619713748},
{'joint_name': 'R_WRIST_Y', 'rad': 0.165660432636},
{'joint_name': 'R_WRIST_R', 'rad': 0.013328010806},
]

air_open_joint_list= [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.562264710474},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.375382944612},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.390867581924},
{'joint_name': 'R_ELBOW_R', 'rad': -1.143580778878},
{'joint_name': 'R_WRIST_P', 'rad': 0.19485637762},
{'joint_name': 'R_WRIST_Y', 'rad': 0.159855796177},
{'joint_name': 'R_WRIST_R', 'rad': 0.26605484086},
]

air_close_joint_list= air_open_joint_list

go_home_joint_list = [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.451632041621},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.307911521083},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.13536865317},
{'joint_name': 'R_ELBOW_R', 'rad': -1.294201365225},
{'joint_name': 'R_WRIST_P', 'rad': 0.378274346518},
{'joint_name': 'R_WRIST_Y', 'rad': 0.037015483938},
{'joint_name': 'R_WRIST_R', 'rad': 0.148664191295},
]

empty_joint_list = [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.255359389661},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.283898456732},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.498425644483},
{'joint_name': 'R_ELBOW_R', 'rad': -1.261196097901},
{'joint_name': 'R_WRIST_P', 'rad': 0.159945974503},
{'joint_name': 'R_WRIST_Y', 'rad': 0.124704284694},
{'joint_name': 'R_WRIST_R', 'rad': 0.089323055792},
]


music_in_joint_list= [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.534755574811},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.370206708698},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.318657524675},
{'joint_name': 'R_ELBOW_R', 'rad': -1.133308992921},
{'joint_name': 'R_WRIST_P', 'rad': 0.217250029084},
{'joint_name': 'R_WRIST_Y', 'rad': 0.151281261241},
{'joint_name': 'R_WRIST_R', 'rad': 0.087528290167},
]

music_open_joint_list = [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.562264710474},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.375382944612},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.390867581924},
{'joint_name': 'R_ELBOW_R', 'rad': -1.143580778878},
{'joint_name': 'R_WRIST_P', 'rad': 0.19485637762},
{'joint_name': 'R_WRIST_Y', 'rad': 0.149855796177},
{'joint_name': 'R_WRIST_R', 'rad': 0.25605484086},
]
music_close_joint_list = music_open_joint_list

set_in_joint_list = [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.494828409031},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.305430193249},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.207381267348},
{'joint_name': 'R_ELBOW_R', 'rad': -1.228615143333},
{'joint_name': 'R_WRIST_P', 'rad': 0.262071504853},
{'joint_name': 'R_WRIST_Y', 'rad': 0.097687807464},
{'joint_name': 'R_WRIST_R', 'rad': 0.069247461928},
]

dirve_mode_joint_list = [
{'joint_name': 'R_SHOULDER_P', 'rad': 0.563493983445},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.162294407962},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.249594216381},
{'joint_name': 'R_ELBOW_R', 'rad': -1.289937354272},
{'joint_name': 'R_WRIST_P', 'rad': 0.142454226992},
{'joint_name': 'R_WRIST_Y', 'rad': 0.000561004112},
{'joint_name': 'R_WRIST_R', 'rad': 0.167736130893},
]

dirve_soft_joint_list = [

{'joint_name': 'R_SHOULDER_P', 'rad': 0.535399163285},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.308316848927},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.505504168453},
{'joint_name': 'R_ELBOW_R', 'rad': -1.088469481487},
{'joint_name': 'R_WRIST_P', 'rad': -0.006021064442},
{'joint_name': 'R_WRIST_Y', 'rad': 0.165673722074},
{'joint_name': 'R_WRIST_R', 'rad': 0.062983648315},
]

dirve_jie_joint_list = [

{'joint_name': 'R_SHOULDER_P', 'rad': 0.511487668206},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.330508311093},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.33301771116},
{'joint_name': 'R_ELBOW_R', 'rad': -1.139994529239},
{'joint_name': 'R_WRIST_P', 'rad': 0.228347658658},
{'joint_name': 'R_WRIST_Y', 'rad': 0.164826995055},
{'joint_name': 'R_WRIST_R', 'rad': 0.075367851372},
]

head_front_joint_list = [
    {'joint_name': 'HEAD_P', 'rad': -0.15311278316},
    {'joint_name': 'HEAD_Y', 'rad': -0.04022955268},
    {'joint_name': 'HEAD_R', 'rad': 0.050031017646},
]

head_screen_joint_list = [
    {'joint_name': 'HEAD_P', 'rad': -0.307521160933},
    {'joint_name': 'HEAD_Y', 'rad': -0.474696809681},
    {'joint_name': 'HEAD_R', 'rad': 0.074742644342},
]

singal_joint_list = [
{'joint_name': 'L_SHOULDER_R', 'rad': 0},
]

hand_touch = [
    {"id": 0, "value": 0},
    {"id": 1, "value": 0},
    {"id": 2, "value": 0},
    {"id": 3, "value": 1.0},
    {"id": 4, "value": 0.3},
    {"id": 5, "value": 0.3},
]

hand_open = [
        {"id": 0, "value": 0.8},
        {"id": 1, "value": 0.8},
        {"id": 2, "value": 0.8},
        {"id": 3, "value": 0.8},
        {"id": 4, "value": 0.8},
        {"id": 5, "value": 1.0},
    ]


hello_l_joint_list = [
    # {'joint_name': 'R_SHOULDER_P', 'rad': 0.399221347784},
    # {'joint_name': 'R_SHOULDER_R', 'rad': -0.287301501983},
    # {'joint_name': 'R_SHOULDER_Y', 'rad': 2.066446782461},
    # {'joint_name': 'R_ELBOW_R', 'rad': -2.153093915083},
    # {'joint_name': 'R_WRIST_P', 'rad': 0.005766666638},
    # {'joint_name': 'R_WRIST_Y', 'rad': 0.292756816084},
    # {'joint_name': 'R_WRIST_R', 'rad': -0.466447865964},
{'joint_name': 'R_SHOULDER_P', 'rad': 0.500229616137},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.252456596808},
{'joint_name': 'R_SHOULDER_Y', 'rad': 2.183436548675},
{'joint_name': 'R_ELBOW_R', 'rad': -2.107785476354},
{'joint_name': 'R_WRIST_P', 'rad': -0.115909524794},
{'joint_name': 'R_WRIST_Y', 'rad': 0.10194042747},
{'joint_name': 'R_WRIST_R', 'rad': -0.283499773607},
]

hello_r_joint_list = [
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.406399542535},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.299956843483},
    {'joint_name': 'R_SHOULDER_Y', 'rad': 2.461169655645},
    {'joint_name': 'R_ELBOW_R', 'rad': -2.049333783919},
    {'joint_name': 'R_WRIST_P', 'rad': 0.185113321427},
    {'joint_name': 'R_WRIST_Y', 'rad': -0.44407319941},
    {'joint_name': 'R_WRIST_R', 'rad': -0.491400683396},
]



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




#


def _get_hand_client(hand_client):
    if hand_client is None:
        return DexHandClient(), True
    return hand_client, False


def action_hello(client, hand_client=None):
    hand_client, created = _get_hand_client(hand_client)
    hand_client.set_angles(hand_open, device_id="hand2")
    client.send(init_joint_list, vel=1.0, acc=1.0)
    client.send(hello_l_joint_list, vel=2.0, acc=2.0)
    for _ in range(3):
        client.send(hello_l_joint_list, vel=3.0, acc=5.0)
        time.sleep(0.1)
        client.send(hello_r_joint_list, vel=3.0, acc=5.0)
        time.sleep(0.1)
    client.send(init_joint_list, vel=1.5, acc=1.5)
    hand_client.set_angles(hand_open, device_id="hand2")
    if created:
        hand_client.close()

def action_air(client, hand_client=None):
    hand_client, created = _get_hand_client(hand_client)
    hand_client.set_angles(hand_touch, device_id="hand2")
    client.send(init_joint_list, vel=2.0, acc=2.0)
    client.send(head_screen_joint_list, vel=1.0, acc=1.0)
    client.send(start_joint_list, vel=1.0, acc=1.0)
    client.send(air_in_joint_list, vel=3.0, acc=3.0)
    client.send(start_joint_list, vel=3.0, acc=3.0)
    # # # # #
    # # time.sleep(1)
    # # client.send(air_open_joint_list, vel=3.0, acc=3.0)
    # # client.send(start_joint_list, vel=3.0, acc=3.0)
    time.sleep(1)
    client.send(air_close_joint_list, vel=3.0, acc=2.0)
    client.send(start_joint_list, vel=3.0, acc=2.0)
    time.sleep(1)
    client.send(go_home_joint_list, vel=3.0, acc=2.0)
    client.send(start_joint_list, vel=3.0, acc=3.0)
    client.send(head_front_joint_list, vel=1.0, acc=1.0)
    client.send(init_joint_list, vel=2.0, acc=2.0)
    hand_client.set_angles(hand_open, device_id="hand2")
    if created:
        hand_client.close()

def action_music(client, hand_client=None):
    hand_client, created = _get_hand_client(hand_client)
    # init_joint_list
    hand_client.set_angles(hand_touch, device_id="hand2")
    client.send(init_joint_list, vel=2.0, acc=2.0)
    client.send(head_screen_joint_list, vel=1.0, acc=1.0)
    client.send(start_joint_list, vel=1.0, acc=1.0)
    client.send(music_in_joint_list, vel=3.0, acc=3.0)
    client.send(start_joint_list, vel=3.0, acc=3.0)
    # # #
    time.sleep(1)
    client.send(music_open_joint_list, vel=3.0, acc=3.0)
    client.send(start_joint_list, vel=3.0, acc=3.0)
    # time.sleep(2)
    # client.send(music_close_joint_list, vel=3.0, acc=3.0)
    # client.send(start_joint_list, vel=3.0, acc=3.0)
    time.sleep(1)
    client.send(go_home_joint_list, vel=3.0, acc=2.0)
    client.send(start_joint_list, vel=3.0, acc=2.0)
    client.send(head_front_joint_list, vel=1.0, acc=1.0)
    client.send(init_joint_list, vel=2.0, acc=2.0)
    hand_client.set_angles(hand_open, device_id="hand2")
    if created:
        hand_client.close()
def action_drive_economy(client):
    # client.send(head_screen_joint_list, vel=1.0, acc=1.0)
    # client.send(start_joint_list, vel=1.0, acc=1.0)
    # client.send(set_in_joint_list, vel=3.0, acc=2.0)
    # client.send(start_joint_list, vel=3.0, acc=2.0)
    # # #
    time.sleep(1)
    # client.send(dirve_mode_joint_list, vel=3.0, acc=2.0)
    # client.send(start_joint_list, vel=3.0, acc=2.0)
    #
    # time.sleep(1)
    # client.send(dirve_jie_joint_list, vel=3.0, acc=2.0)
    # client.send(start_joint_list, vel=3.0, acc=2.0)
    # time.sleep(2)
    # client.send(go_home_joint_list, vel=3.0, acc=2.0)
    # client.send(start_joint_list, vel=3.0, acc=2.0)
    # client.send(head_front_joint_list, vel=1.0, acc=1.0)
def action_drive_comfort(client):
    # client.send(head_screen_joint_list, vel=1.0, acc=1.0)
    # client.send(start_joint_list, vel=1.0, acc=1.0)
    # client.send(set_in_joint_list, vel=3.0, acc=2.0)
    # client.send(start_joint_list, vel=3.0, acc=2.0)
    # # #
    time.sleep(1)
    # client.send(dirve_mode_joint_list, vel=3.0, acc=2.0)
    # client.send(start_joint_list, vel=3.0, acc=2.0)
    #
    # time.sleep(1)
    # client.send(dirve_soft_joint_list, vel=3.0, acc=2.0)
    # client.send(start_joint_list, vel=3.0, acc=2.0)
    # time.sleep(2)
    # client.send(go_home_joint_list, vel=3.0, acc=2)
    # client.send(start_joint_list, vel=3.0, acc=2)
    # client.send(head_front_joint_list, vel=1.0, acc=1.0)

    # client.send(init_joint_list, vel=1.0, acc=1.0)
if __name__ == "__main__":
    # Initialize client
    client = MoveJClient()
    hand_client = DexHandClient()
    # action_air(client)
    # action_hello(client)
    client.send(test_joint_list, vel=1.0, acc=1.0)
    # client.send(go_home_joint_list, vel=3.0, acc=3.0)
    # action_music(client)

    # action_music(client)
    # action_air(client)
    # action_one(client)
    # action_drive_soft(client)
    # action_drive_economy(client)
    # action_drive_comfort(client)
    # client.send(singal_joint_list, vel=1.0, acc=1.0)
    # client.send(start_joint_list, vel=1.0, acc=1.0)
    # while True:
    #     client.send(hello_l_joint_list, vel=1.0, acc=1.0)
    #     time.sleep(0.1)
    #     client.send(hello_r_joint_list, vel=1.0, acc=1.0)
    #     time.sleep(0.1)
    #     # hello_l_joint_list



    # Close client
    hand_client.close()
    client.close()
