import csv
import math

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 arm_service_pb2_grpc as rpc


class GetJointStateClient(RobotClientBase):
    """客户端：获取机械臂关节数据并保存到CSV"""

    def fetch_and_save_joint_states(self, device_id="right_arm"):
        """获取关节数据一次并保存到CSV文件"""
        try:
            # 构建请求
            req = pb.JointRequest()
            req.header.device_id = device_id
            timestamp = timestamp_pb2.Timestamp()
            timestamp.GetCurrentTime()
            req.header.timestamp.CopyFrom(timestamp)

            # 调用RPC获取数据
            try:
                resp = self.stub.getJointState(req, timeout=10)
            except Exception as e:
                print(f"RPC调用失败: {e}")
                return

            # 定义需要保存的关节名称
            joint_order = [
                "R_SHOULDER_P", "R_SHOULDER_R", "R_SHOULDER_Y",
                "R_ELBOW_R", "R_WRIST_P", "R_WRIST_Y", "R_WRIST_R"
            ]

            # 将响应数据映射到字典
            joint_dict = {}
            state = resp.state
            for name, pos in zip(state.name, state.position):
                joint_dict[name] = round(pos, 6)

            # 将角度从弧度转换为度
            joint_list = [{"joint_name": name, "deg": round(joint_dict.get(name, 0.0) * 180 / math.pi, 10)}
                          for name in joint_order]

            # 准备保存到CSV的数据行
            data_row = [joint['deg'] for joint in joint_list]

            # 以追加模式打开CSV文件并写入数据
            with open("joint_states_10_23.csv", mode="a", newline="") as file:
                writer = csv.writer(file)
                # 如果文件为空，写入表头
                if file.tell() == 0:
                    writer.writerow(joint_order)  # 表头
                writer.writerow(data_row)  # 写入数据行

            print(f"数据已保存: {data_row}")

        except KeyboardInterrupt:
            print("\n停止获取关节数据。")


if __name__ == "__main__":
    client = GetJointStateClient()
    try:
        client.fetch_and_save_joint_states()
    finally:
        client.close()