#!/usr/bin/env python3
"""独立双臂状态监控脚本（轮询版）。

优化点：
- 不再依赖 UDP 回调；
- 直接通过 rm_get_current_arm_state() 按句柄轮询左右臂状态；
- 避免同进程双回调潜在冲突导致的左臂显示异常。
"""

from __future__ import annotations

import argparse
import time
from dataclasses import dataclass
from numbers import Number
from pathlib import Path
from typing import Any

import numpy as np
import yaml


DEFAULT_HOST_IP = "0.0.0.0"
DEFAULT_LEFT_IP = "192.168.192.18"
DEFAULT_RIGHT_IP = "192.168.192.19"
DEFAULT_ROBOT_PORT = 8080


def default_config_path() -> Path:
    """同时兼容源码目录和 colcon install 后的配置文件路径。"""
    script_path = Path(__file__).resolve()
    candidates = [
        script_path.parents[1] / "config" / "dual_arm_rm75.yaml",
        script_path.parents[2] / "share" / "xr_rm_bringup" / "config" / "dual_arm_rm75.yaml",
    ]

    try:
        from ament_index_python.packages import get_package_share_directory

        candidates.append(Path(get_package_share_directory("xr_rm_bringup")) / "config" / "dual_arm_rm75.yaml")
    except Exception:
        pass

    for candidate in candidates:
        if candidate.is_file():
            return candidate
    return candidates[0]


DEFAULT_CONFIG_PATH = default_config_path()


@dataclass
class ArmSnapshot:
    """单臂状态快照。"""

    connected: bool
    ret_code: int
    last_update: float
    joint_deg: list[float]
    pose_euler: np.ndarray
    err_text: str


@dataclass
class MonitorDefaults:
    """从 YAML 配置中读取到的监控默认参数。"""

    config_path: Path
    left_ip: str = DEFAULT_LEFT_IP
    right_ip: str = DEFAULT_RIGHT_IP
    left_port: int = DEFAULT_ROBOT_PORT
    right_port: int = DEFAULT_ROBOT_PORT
    host_ip: str = DEFAULT_HOST_IP


def _load_yaml(path: Path) -> dict[str, Any]:
    """读取 ROS 参数 YAML，返回顶层 dict。"""
    if not path.is_file():
        raise FileNotFoundError(f"配置文件不存在: {path}")

    with path.open("r", encoding="utf-8") as f:
        data = yaml.safe_load(f) or {}

    if not isinstance(data, dict):
        raise ValueError(f"配置文件格式错误，顶层应为 dict: {path}")
    return data


def _node_parameters(config: dict[str, Any], node_name: str) -> dict[str, Any]:
    node_config = config.get(node_name, {})
    if not isinstance(node_config, dict):
        return {}
    params = node_config.get("ros__parameters", {})
    return params if isinstance(params, dict) else {}


def load_monitor_defaults(config_path: Path) -> MonitorDefaults:
    """从双臂配置读取左右臂连接参数。"""
    config = _load_yaml(config_path)
    left_params = _node_parameters(config, "left_arm_teleop")
    right_params = _node_parameters(config, "right_arm_teleop")

    return MonitorDefaults(
        config_path=config_path,
        left_ip=str(left_params.get("robot_ip", DEFAULT_LEFT_IP)),
        right_ip=str(right_params.get("robot_ip", DEFAULT_RIGHT_IP)),
        left_port=int(left_params.get("robot_port", DEFAULT_ROBOT_PORT)),
        right_port=int(right_params.get("robot_port", DEFAULT_ROBOT_PORT)),
    )


def load_realman_sdk() -> tuple[Any, Any]:
    """延迟导入睿尔曼 SDK，让 --help 和参数解析不依赖现场环境。"""
    try:
        from Robotic_Arm.rm_robot_interface import RoboticArm, rm_thread_mode_e
    except ImportError as exc:
        raise RuntimeError(
            "未安装睿尔曼 Python API2，无法连接真机。请先安装厂商 SDK，"
            "确认可以导入 Robotic_Arm.rm_robot_interface。"
        ) from exc
    return RoboticArm, rm_thread_mode_e


def resolve_thread_mode(rm_thread_mode_e: Any, mode: int) -> Any:
    """兼容不同 SDK 版本的线程模式枚举写法。"""
    try:
        return rm_thread_mode_e(mode)
    except Exception:
        pass

    mode_names = {
        0: "RM_SINGLE_MODE_E",
        1: "RM_DUAL_MODE_E",
        2: "RM_TRIPLE_MODE_E",
    }
    mode_name = mode_names.get(int(mode))
    if mode_name and hasattr(rm_thread_mode_e, mode_name):
        return getattr(rm_thread_mode_e, mode_name)
    return mode


def create_robot_arm(robot: Any, robot_ip: str, robot_port: int, level: int) -> Any:
    """兼容 rm_create_robot_arm(ip, port) 和 rm_create_robot_arm(ip, port, level)。"""
    try:
        return robot.rm_create_robot_arm(robot_ip, robot_port, level)
    except TypeError:
        return robot.rm_create_robot_arm(robot_ip, robot_port)


def robot_handle_id(handle: Any) -> int | None:
    """从 SDK handle 中取 id，兼容对象和整数返回值。"""
    if isinstance(handle, int):
        return handle
    return getattr(handle, "id", None)


def return_code(obj: Any) -> int:
    """从不同 SDK 返回结构中解析 ret/code。"""
    if isinstance(obj, tuple) and obj and isinstance(obj[0], int):
        return obj[0]
    if isinstance(obj, dict):
        for key in ("ret", "code", "err_code", "error_code"):
            value = obj.get(key)
            if isinstance(value, int):
                return value
    return 0


def payload(obj: Any) -> Any:
    """从 (ret, data) 风格返回值中取状态主体。"""
    if isinstance(obj, tuple) and len(obj) >= 2 and isinstance(obj[0], int):
        return obj[1]
    return obj


def _as_number_list(value: Any, expected_len: int) -> list[float] | None:
    if isinstance(value, np.ndarray):
        value = value.tolist()
    if isinstance(value, (list, tuple)) and len(value) >= expected_len:
        if all(isinstance(item, Number) for item in value[:expected_len]):
            return [float(item) for item in value[:expected_len]]
    return None


def _as_pose(value: Any) -> list[float] | None:
    pose = _as_number_list(value, 6)
    if pose is not None:
        return pose

    if isinstance(value, dict):
        position = value.get("position")
        euler = value.get("euler")
        if isinstance(position, dict) and isinstance(euler, dict):
            xyz_keys = ("x", "y", "z")
            rpy_keys = ("rx", "ry", "rz")
            if all(key in position for key in xyz_keys) and all(key in euler for key in rpy_keys):
                return [
                    float(position["x"]),
                    float(position["y"]),
                    float(position["z"]),
                    float(euler["rx"]),
                    float(euler["ry"]),
                    float(euler["rz"]),
                ]

    if all(hasattr(value, attr) for attr in ("position", "euler")):
        position = getattr(value, "position")
        euler = getattr(value, "euler")
        if all(hasattr(position, key) for key in ("x", "y", "z")) and all(
            hasattr(euler, key) for key in ("rx", "ry", "rz")
        ):
            return [
                float(position.x),
                float(position.y),
                float(position.z),
                float(euler.rx),
                float(euler.ry),
                float(euler.rz),
            ]
    return None


def find_pose(obj: Any) -> list[float] | None:
    """递归查找常见 TCP 位姿字段。"""
    pose = _as_pose(obj)
    if pose is not None:
        return pose

    if isinstance(obj, dict):
        for key in ("pose", "tool_pose", "tcp_pose", "current_pose"):
            pose = _as_pose(obj.get(key))
            if pose is not None:
                return pose
        for value in obj.values():
            pose = find_pose(value)
            if pose is not None:
                return pose
    elif isinstance(obj, (list, tuple)):
        for value in obj:
            pose = find_pose(value)
            if pose is not None:
                return pose
    elif hasattr(obj, "to_dictionary"):
        try:
            return find_pose(obj.to_dictionary(7))
        except TypeError:
            return find_pose(obj.to_dictionary())
    elif hasattr(obj, "to_dict"):
        return find_pose(obj.to_dict())
    else:
        for key in ("pose", "tool_pose", "tcp_pose", "current_pose"):
            if hasattr(obj, key):
                pose = _as_pose(getattr(obj, key))
                if pose is not None:
                    return pose
    return None


def find_joint(obj: Any) -> list[float] | None:
    """递归查找常见关节角字段。"""
    joint = _as_number_list(obj, 7)
    if joint is not None:
        return joint

    if isinstance(obj, dict):
        for key in ("joint", "joints", "joint_deg", "joint_position", "joint_positions"):
            joint = _as_number_list(obj.get(key), 7)
            if joint is not None:
                return joint
        for value in obj.values():
            joint = find_joint(value)
            if joint is not None:
                return joint
    elif isinstance(obj, (list, tuple)):
        for value in obj:
            joint = find_joint(value)
            if joint is not None:
                return joint
    elif hasattr(obj, "to_dictionary"):
        try:
            return find_joint(obj.to_dictionary(7))
        except TypeError:
            return find_joint(obj.to_dictionary())
    elif hasattr(obj, "to_dict"):
        return find_joint(obj.to_dict())
    else:
        for key in ("joint", "joints", "joint_deg", "joint_position", "joint_positions"):
            if hasattr(obj, key):
                joint = _as_number_list(getattr(obj, key), 7)
                if joint is not None:
                    return joint
    return None


def find_error_text(obj: Any) -> str:
    if isinstance(obj, dict):
        for key in ("err", "error", "error_msg", "err_msg"):
            if key in obj:
                return str(obj[key])
    return "ok"


class ArmStatePoller:
    """单臂状态轮询器。"""

    def __init__(
        self,
        name: str,
        robot_ip: str,
        robot_port: int,
        level: int = 3,
        mode: int = 2,
    ):
        self.name = name
        self.robot_ip = robot_ip
        self.robot_port = int(robot_port)
        self.level = int(level)
        self.mode = int(mode)

        self.robot: Any | None = None
        self.robot_handle = None
        self.connected = False

        self.last_snapshot = ArmSnapshot(
            connected=False,
            ret_code=-999,
            last_update=0.0,
            joint_deg=[0.0] * 7,
            pose_euler=np.zeros(6, dtype=np.float64),
            err_text="not connected",
        )

        self._connect()

    def _connect(self):
        """连接机械臂。"""
        try:
            RoboticArm, rm_thread_mode_e = load_realman_sdk()
            self.robot = RoboticArm(resolve_thread_mode(rm_thread_mode_e, self.mode))
            self.robot_handle = create_robot_arm(self.robot, self.robot_ip, self.robot_port, self.level)
            handle_id = robot_handle_id(self.robot_handle)
            if self.robot_handle is None or handle_id == -1:
                self.connected = False
                self.last_snapshot.connected = False
                self.last_snapshot.err_text = f"connect failed ({self.robot_ip}:{self.robot_port})"
                print(f"[{self.name}] 连接失败: {self.last_snapshot.err_text}")
                return

            self.connected = True
            self.last_snapshot.connected = True
            self.last_snapshot.err_text = "ok"
            print(f"[{self.name}] 已连接, handle_id={handle_id}, ip={self.robot_ip}, port={self.robot_port}")
        except Exception as exc:
            self.connected = False
            self.last_snapshot.connected = False
            self.last_snapshot.err_text = f"connect exception: {exc}"
            print(f"[{self.name}] 连接异常: {exc}")

    def poll(self) -> ArmSnapshot:
        """轮询一次当前状态。"""
        if (not self.connected) or (self.robot is None):
            return self.last_snapshot

        try:
            raw_state = self.robot.rm_get_current_arm_state()
            ret = return_code(raw_state)
            state = payload(raw_state)
            now_t = time.time()

            if ret != 0:
                self.last_snapshot.ret_code = ret
                self.last_snapshot.err_text = f"rm_get_current_arm_state ret={ret}"
                return self.last_snapshot

            joint = find_joint(state) or [0.0] * 7
            pose = find_pose(state) or [0.0] * 6
            err = find_error_text(state)

            if len(joint) < 7:
                joint.extend([0.0] * (7 - len(joint)))
            elif len(joint) > 7:
                joint = joint[:7]

            if len(pose) < 6:
                pose.extend([0.0] * (6 - len(pose)))
            elif len(pose) > 6:
                pose = pose[:6]

            self.last_snapshot = ArmSnapshot(
                connected=True,
                ret_code=ret,
                last_update=now_t,
                joint_deg=joint,
                pose_euler=np.array(pose, dtype=np.float64),
                err_text=str(err),
            )
            return self.last_snapshot
        except Exception as exc:
            self.last_snapshot.err_text = f"poll exception: {exc}"
            return self.last_snapshot

    def close(self):
        """断开连接。"""
        if self.robot is None:
            return
        try:
            self.robot.rm_delete_robot_arm()
        except Exception:
            pass


def format_arm_block(name: str, snap: ArmSnapshot, now_t: float) -> str:
    """格式化单臂输出。"""
    if not snap.connected:
        return f"[{name}] 未连接: {snap.err_text}"

    age = now_t - snap.last_update if snap.last_update > 0 else float("inf")
    stale_flag = "(stale)" if age > 2.0 else ""

    j = ", ".join(f"{v:7.2f}" for v in snap.joint_deg)
    p = snap.pose_euler
    return (
        f"[{name}] update_age={age:5.2f}s {stale_flag}, ret={snap.ret_code}\n"
        f"  Joint(deg): [{j}]\n"
        f"  Pose[x y z rx ry rz]: [{p[0]: .4f}, {p[1]: .4f}, {p[2]: .4f}, {p[3]: .4f}, {p[4]: .4f}, {p[5]: .4f}]\n"
        f"  Err: {snap.err_text}"
    )


def build_arg_parser() -> argparse.ArgumentParser:
    bootstrap = argparse.ArgumentParser(add_help=False)
    bootstrap.add_argument("--config", type=Path, default=DEFAULT_CONFIG_PATH)
    known_args, _ = bootstrap.parse_known_args()
    defaults = load_monitor_defaults(known_args.config)

    parser = argparse.ArgumentParser(description="Realman 双臂状态监控脚本（轮询版）")
    parser.add_argument("--config", type=Path, default=defaults.config_path, help="双臂 YAML 配置文件")
    parser.add_argument("--left-ip", default=defaults.left_ip, help="左臂 IP")
    parser.add_argument("--right-ip", default=defaults.right_ip, help="右臂 IP")
    parser.add_argument("--left-port", type=int, default=defaults.left_port, help="左臂 TCP 端口")
    parser.add_argument("--right-port", type=int, default=defaults.right_port, help="右臂 TCP 端口")
    parser.add_argument("--refresh", type=float, default=0.5, help="刷新周期（秒）")
    parser.add_argument("--level", type=int, default=3, help="连接级别")
    parser.add_argument("--mode", type=int, default=2, help="线程模式 0/1/2")
    parser.add_argument("--host-ip", default=defaults.host_ip, help="仅展示当前配置用，轮询版不使用")
    return parser


def main():
    args = build_arg_parser().parse_args()

    print(
        f"配置: config={args.config}, "
        f"left={args.left_ip}:{args.left_port}, "
        f"right={args.right_ip}:{args.right_port}, host_ip={args.host_ip}"
    )

    left = ArmStatePoller(
        name="LEFT",
        robot_ip=args.left_ip,
        robot_port=args.left_port,
        level=args.level,
        mode=args.mode,
    )
    right = ArmStatePoller(
        name="RIGHT",
        robot_ip=args.right_ip,
        robot_port=args.right_port,
        level=args.level,
        mode=args.mode,
    )

    print("\n开始监控双臂状态（轮询），按 Ctrl+C 退出。\n")

    try:
        while True:
            now_t = time.time()
            l_snap = left.poll()
            r_snap = right.poll()

            print("\033[2J\033[H", end="")
            print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(now_t)))
            print("=" * 88)
            print(format_arm_block("LEFT", l_snap, now_t))
            print("-" * 88)
            print(format_arm_block("RIGHT", r_snap, now_t))
            print("=" * 88)
            time.sleep(max(args.refresh, 0.05))
    except KeyboardInterrupt:
        print("\n退出监控...")
    finally:
        left.close()
        right.close()


if __name__ == "__main__":
    main()
