import csv

import math
import sys
from google.protobuf import timestamp_pb2

sys.path.append("../generated")
import math
from typing import List, Dict, Literal
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


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)

    def read_joint_from_csv(self, file_name, row_num):
        """Read joint data from a CSV file and return joint list for the specified row."""
        joint_order = [
            'R_SHOULDER_P', 'R_SHOULDER_R', 'R_SHOULDER_Y',
            'R_ELBOW_R', 'R_WRIST_P', 'R_WRIST_Y', 'R_WRIST_R'
        ]

        try:
            with open(file_name, newline='') as csvfile:
                reader = csv.reader(csvfile)
                # Skip header row
                next(reader)
                # Read specific row
                for i, row in enumerate(reader, start=1):
                    if i == row_num:
                        # Assign joint values to the list
                        joint_list = [
                            {'joint_name': joint_order[j], 'rad': round(float(row[j]) * math.pi / 180, 10)}
                            for j in range(len(joint_order))
                        ]
                        return joint_list
            print(f"Row {row_num} not found.")
            return []
        except Exception as e:
            print(f"Error reading CSV: {e}")
            return []

    def apply_joint_offsets(self,
            joint_list: List[Dict[str, float]],
            on_missing: Literal["ignore", "error", "warn"] = "ignore",
    ) -> List[Dict[str, float]]:
        """
        将零偏置 joint_zero_offset 叠加到 joint_list 上，按 joint_name 对齐。

        参数：
            joint_list: [{'joint_name': str, 'rad': float}, ...]  原始指令
            joint_zero_offset: [{'joint_name': str, 'rad': float}, ...]  零偏置
            on_missing: 当 joint_list 中出现 zero_offset 里没有的关节名时的处理：
                - "ignore": 当作偏置为 0（默认）
                - "warn":   打印警告，再当作 0
                - "error":  抛出 KeyError

        返回：
            与 joint_list 顺序一致、已叠加偏置的新列表
        """
        # 构造偏置表，并检查是否有重复关节名

        deg_to_rad = math.pi / 180.0

        joint_zero_offset = [
            {'joint_name': 'R_SHOULDER_P', 'rad': 0},
            {'joint_name': 'R_SHOULDER_R', 'rad': 0 * deg_to_rad},
            {'joint_name': 'R_SHOULDER_Y', 'rad': 0.2 * deg_to_rad},
            {'joint_name': 'R_ELBOW_R', 'rad': 0.2 * deg_to_rad},
            {'joint_name': 'R_WRIST_P', 'rad': 0.6 * deg_to_rad},
            {'joint_name': 'R_WRIST_Y', 'rad': 2.5 * deg_to_rad},
            {'joint_name': 'R_WRIST_R', 'rad': 0},
        ]

        offset_map: Dict[str, float] = {}
        for j in joint_zero_offset:
            name = j["joint_name"]
            if name in offset_map:
                raise ValueError(f"Duplicated joint in zero_offset: {name}")
            offset_map[name] = float(j["rad"])

        joint_cmd: List[Dict[str, float]] = []
        for j in joint_list:
            name = j["joint_name"]
            base = float(j["rad"])
            if name not in offset_map:
                if on_missing == "error":
                    raise KeyError(f"Missing zero offset for joint: {name}")
                elif on_missing == "warn":
                    print(f"[apply_joint_offsets] WARN: Missing zero offset for {name}, using 0.0")
                off = 0.0
            else:
                off = offset_map[name]

            joint_cmd.append({"joint_name": name, "rad": base + off})
        return joint_cmd


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

    # Specify CSV file and row number to read from
    file_name = 'joint_states.csv'
    row_num = 2  # Row number you want to read (e.g., 2 for the second row)

    # Read joint data from CSV
    joint_list = client.read_joint_from_csv(file_name, 8)
    joint_list = client.apply_joint_offsets(joint_list)
    print(joint_list)

    if joint_list:

        # Send MoveJ command with the read joint data
        client.send(joint_list, vel=0.8, acc=0.8)

    # Close client
    client.close()