#!/usr/bin/env python3
"""CRAIC mechanical arm UDP controller.

Mechanical structure used here:
1. J2/J3/J4 rotate around the Z axis in the XY plane.
2. The gripper TCP has a fixed offset relative to the J4 frame: (x4, 0, z4).
3. ``phi`` is the TCP yaw in the XY plane and is always ``J2 + J3 + J4``.
4. All linear pose and geometry values are millimeters.

Supports two control modes:
1. Direct joint command mode.
2. Cartesian TCP pose mode using the inverse kinematics from ``docs/craic.md``.
"""

from __future__ import annotations

import argparse
import json
import math
import socket
import sys
import time
from dataclasses import dataclass
from pathlib import Path


DEFAULT_UDP_IP = "192.168.4.1"
DEFAULT_UDP_PORT = 8888

DEFAULT_HEIGHT_MIN = -290    # 内部 d1 坐标：底部，单位 mm（z 轴朝上）
DEFAULT_HEIGHT_MAX = 0       # 内部 d1 坐标：顶部，单位 mm（z 轴朝上）
DEFAULT_JOINT_MIN = -180
DEFAULT_JOINT_MAX = 180
DEFAULT_J2_MIN = -110
DEFAULT_J2_MAX = 115
DEFAULT_J3_MIN = -120
DEFAULT_J3_MAX = 145
DEFAULT_J4_MIN = -90
DEFAULT_J4_MAX = 130
J5_OPEN = 81
J5_CLOSED = -100
Z4_OPEN = 55
Z4_CLOSED = -100
DEFAULT_FIXED_J5 = J5_OPEN
# ==== 抓取/释放（由 J6 控制）====
# 请填写实际角度值：
GRIP_ANGLE = -5      # TODO: 填写抓取时 J6 的角度
RELEASE_ANGLE = 80   # TODO: 填写释放时 J6 的角度
DEFAULT_FIXED_J6 = RELEASE_ANGLE
DEFAULT_ZERO_J2 = 3
DEFAULT_ZERO_J3 = 7
DEFAULT_ZERO_J4 = 25
DEFAULT_L1 = 125.0
DEFAULT_L2 = 125.0
DEFAULT_X4 = 110.0
DEFAULT_Z4 = 80.0
DEFAULT_INTERP_DURATION = 1.0
DEFAULT_INTERP_RATE = 20.0
STATE_FILE = Path(__file__).with_name(".udp_control_state.json")


class ArmControlError(ValueError):
    """Raised when the requested arm pose is invalid."""


@dataclass(frozen=True)
class ArmGeometry:
    l1: float
    l2: float
    x4: float
    z4: float


@dataclass(frozen=True)
class ArmLimits:
    height_min: int = DEFAULT_HEIGHT_MIN
    height_max: int = DEFAULT_HEIGHT_MAX
    joint_min: int = DEFAULT_JOINT_MIN
    joint_max: int = DEFAULT_JOINT_MAX
    j2_min: int = DEFAULT_J2_MIN
    j2_max: int = DEFAULT_J2_MAX
    j3_min: int = DEFAULT_J3_MIN
    j3_max: int = DEFAULT_J3_MAX
    j4_min: int = DEFAULT_J4_MIN
    j4_max: int = DEFAULT_J4_MAX


@dataclass(frozen=True)
class ArmZeroOffsets:
    j2: int = DEFAULT_ZERO_J2
    j3: int = DEFAULT_ZERO_J3
    j4: int = DEFAULT_ZERO_J4


@dataclass(frozen=True)
class ArmJointState:
    height: int
    j2: int
    j3: int
    j4: int
    j5: int = DEFAULT_FIXED_J5
    j6: int = DEFAULT_FIXED_J6

    def to_udp_message(self) -> bytes:
        return (
            f"JXB:{self.height}:{self.j2}:{self.j3}:{self.j4}:"
            f"{self.j5}:{self.j6}:0:0:EZHY\n"
        ).encode("utf-8")


@dataclass(frozen=True)
class ArmPose:
    x: float
    y: float
    z: float
    phi_deg: float


@dataclass(frozen=True)
class ArmMathState:
    d1: float
    theta2_deg: float
    theta3_deg: float
    theta4_deg: float


@dataclass(frozen=True)
class Joint4Center:
    x: float
    y: float
    z: float


def default_command_state() -> ArmJointState:
    return ArmJointState(
        height=-DEFAULT_HEIGHT_MAX,
        j2=DEFAULT_ZERO_J2,
        j3=DEFAULT_ZERO_J3,
        j4=DEFAULT_ZERO_J4,
        j5=DEFAULT_FIXED_J5,
        j6=DEFAULT_FIXED_J6,
    )


def clamp_int(value: float, lower: int, upper: int, name: str) -> int:
    rounded = int(round(value))
    if rounded < lower or rounded > upper:
        raise ArmControlError(
            f"{name}={rounded} 超出范围 [{lower}, {upper}]"
        )
    return rounded


def normalize_angle_deg(angle_deg: float) -> float:
    normalized = (angle_deg + 180.0) % 360.0 - 180.0
    if normalized == -180.0 and angle_deg > 0:
        return 180.0
    return normalized


def lerp(start: float, end: float, t: float) -> float:
    return start + (end - start) * t


def lerp_angle_deg(start_deg: float, end_deg: float, t: float) -> float:
    delta = normalize_angle_deg(end_deg - start_deg)
    return normalize_angle_deg(start_deg + delta * t)


def tcp_to_joint4_center(geometry: ArmGeometry, pose: ArmPose) -> Joint4Center:
    """Project the TCP target back to the J4 rotation center.

    ``x4`` and ``z4`` only affect position conversion. They do not affect ``phi``.
    """

    phi = math.radians(pose.phi_deg)
    return Joint4Center(
        x=pose.x - geometry.x4 * math.cos(phi),
        y=pose.y - geometry.x4 * math.sin(phi),
        z=pose.z + geometry.z4,
    )


def forward_kinematics(geometry: ArmGeometry, state: ArmMathState) -> ArmPose:
    theta2 = math.radians(state.theta2_deg)
    theta3 = math.radians(state.theta3_deg)
    theta4 = math.radians(state.theta4_deg)
    phi = theta2 + theta3 + theta4
    j4_center_x = (
        geometry.l1 * math.cos(theta2)
        + geometry.l2 * math.cos(theta2 + theta3)
    )
    j4_center_y = (
        geometry.l1 * math.sin(theta2)
        + geometry.l2 * math.sin(theta2 + theta3)
    )
    x = j4_center_x + geometry.x4 * math.cos(phi)
    y = j4_center_y + geometry.x4 * math.sin(phi)
    z = state.d1 - geometry.z4
    return ArmPose(x=x, y=y, z=z, phi_deg=math.degrees(phi))


def inverse_kinematics(
    geometry: ArmGeometry,
    pose: ArmPose,
    limits: ArmLimits,
    elbow_up: bool,
    j5: int,
    j6: int,
) -> ArmMathState:
    joint4_center = tcp_to_joint4_center(geometry, pose)
    d1 = joint4_center.z

    r2 = joint4_center.x * joint4_center.x + joint4_center.y * joint4_center.y
    if r2 < 1e-9:
        raise ArmControlError("目标点过于接近奇异点，无法稳定求解 J2。")

    denom = 2.0 * geometry.l1 * geometry.l2
    if abs(denom) < 1e-9:
        raise ArmControlError("机械臂几何参数无效：L1 和 L2 不能为 0。")

    c3 = (r2 - geometry.l1 * geometry.l1 - geometry.l2 * geometry.l2) / denom
    if c3 < -1.0 - 1e-9 or c3 > 1.0 + 1e-9:
        reach = math.sqrt(r2)
        raise ArmControlError(
            f"目标超出工作空间，关节 4 投影距离为 {reach:.3f}mm。"
        )
    c3 = max(-1.0, min(1.0, c3))

    s3_abs = math.sqrt(max(0.0, 1.0 - c3 * c3))
    s3 = -s3_abs if elbow_up else s3_abs

    theta3 = math.atan2(s3, c3)
    theta2 = math.atan2(joint4_center.y, joint4_center.x) - math.atan2(
        geometry.l2 * s3,
        geometry.l1 + geometry.l2 * c3,
    )
    phi = math.radians(pose.phi_deg)
    theta4 = phi - theta2 - theta3

    return ArmMathState(
        d1=d1,
        theta2_deg=normalize_angle_deg(math.degrees(theta2)),
        theta3_deg=normalize_angle_deg(math.degrees(theta3)),
        theta4_deg=normalize_angle_deg(math.degrees(theta4)),
    )


def command_to_math_state(
    command_state: ArmJointState,
    zero_offsets: ArmZeroOffsets,
) -> ArmMathState:
    return ArmMathState(
        d1=command_state.height,
        theta2_deg=command_state.j2 - zero_offsets.j2,
        theta3_deg=command_state.j3 - zero_offsets.j3,
        theta4_deg=command_state.j4 - zero_offsets.j4,
    )


def math_to_command_state(
    math_state: ArmMathState,
    zero_offsets: ArmZeroOffsets,
    limits: ArmLimits,
    j5: int,
    j6: int,
) -> ArmJointState:
    return ArmJointState(
        height=clamp_int(
            round(math_state.d1),
            limits.height_min,
            limits.height_max,
            "height(cmd)",
        ),
        j2=clamp_int(
            math_state.theta2_deg + zero_offsets.j2,
            limits.j2_min,
            limits.j2_max,
            "J2(cmd)",
        ),
        j3=clamp_int(
            math_state.theta3_deg + zero_offsets.j3,
            limits.j3_min,
            limits.j3_max,
            "J3(cmd)",
        ),
        j4=clamp_int(
            math_state.theta4_deg + zero_offsets.j4,
            limits.j4_min,
            limits.j4_max,
            "J4(cmd)",
        ),
        j5=clamp_int(j5, limits.joint_min, limits.joint_max, "J5"),
        j6=clamp_int(j6, limits.joint_min, limits.joint_max, "J6"),
    )


def load_cached_command_state(limits: ArmLimits) -> ArmJointState | None:
    if not STATE_FILE.exists():
        return None

    try:
        payload = json.loads(STATE_FILE.read_text(encoding="utf-8"))
        return ArmJointState(
            height=clamp_int(payload["height"], limits.height_min, limits.height_max, "height(cmd)"),
            j2=clamp_int(payload["j2"], limits.j2_min, limits.j2_max, "J2"),
            j3=clamp_int(payload["j3"], limits.j3_min, limits.j3_max, "J3"),
            j4=clamp_int(payload["j4"], limits.j4_min, limits.j4_max, "J4"),
            j5=clamp_int(payload.get("j5", DEFAULT_FIXED_J5), limits.joint_min, limits.joint_max, "J5"),
            j6=clamp_int(payload.get("j6", DEFAULT_FIXED_J6), limits.joint_min, limits.joint_max, "J6"),
        )
    except (OSError, json.JSONDecodeError, KeyError, TypeError, ArmControlError):
        return None


def save_cached_command_state(state: ArmJointState) -> None:
    STATE_FILE.write_text(
        json.dumps(
            {
                "height": state.height,
                "j2": state.j2,
                "j3": state.j3,
                "j4": state.j4,
                "j5": state.j5,
                "j6": state.j6,
            },
            ensure_ascii=True,
            indent=2,
        ),
        encoding="utf-8",
    )


def resolve_start_command_state(
    limits: ArmLimits,
    use_state_cache: bool,
) -> ArmJointState:
    if use_state_cache:
        cached_state = load_cached_command_state(limits)
        if cached_state is not None:
            return cached_state
    return default_command_state()


def interpolate_command_states(
    start: ArmJointState,
    end: ArmJointState,
    steps: int,
) -> list[ArmJointState]:
    if steps <= 1:
        return [end]

    states: list[ArmJointState] = []
    for step_index in range(1, steps + 1):
        t = step_index / steps
        states.append(
            ArmJointState(
                height=int(round(lerp(start.height, end.height, t))),
                j2=int(round(lerp(start.j2, end.j2, t))),
                j3=int(round(lerp(start.j3, end.j3, t))),
                j4=int(round(lerp(start.j4, end.j4, t))),
                j5=int(round(lerp(start.j5, end.j5, t))),
                j6=int(round(lerp(start.j6, end.j6, t))),
            )
        )
    return states


def interpolate_pose(
    start: ArmPose,
    end: ArmPose,
    t: float,
) -> ArmPose:
    return ArmPose(
        x=lerp(start.x, end.x, t),
        y=lerp(start.y, end.y, t),
        z=lerp(start.z, end.z, t),
        phi_deg=lerp_angle_deg(start.phi_deg, end.phi_deg, t),
    )


def build_pose_command_path(
    start_pose: ArmPose,
    target_pose: ArmPose,
    steps: int,
    geometry: ArmGeometry,
    limits: ArmLimits,
    zero_offsets: ArmZeroOffsets,
    elbow_up: bool,
    j5: int,
    j6: int,
) -> list[ArmJointState]:
    if steps <= 1:
        math_state = inverse_kinematics(
            geometry=geometry,
            pose=target_pose,
            limits=limits,
            elbow_up=elbow_up,
            j5=j5,
            j6=j6,
        )
        return [math_to_command_state(math_state, zero_offsets, limits, j5, j6)]

    # Interpolate pose first (movement)
    path: list[ArmJointState] = []
    start_command_j5 = None
    start_command_j6 = None

    for step_index in range(1, steps + 1):
        t = step_index / steps
        pose = interpolate_pose(start_pose, target_pose, t)
        math_state = inverse_kinematics(
            geometry=geometry,
            pose=pose,
            limits=limits,
            elbow_up=elbow_up,
            j5=j5,
            j6=j6,
        )
        command = math_to_command_state(math_state, zero_offsets, limits, j5, j6)

        # Store the first command's j5/j6 values
        if start_command_j5 is None:
            start_command_j5 = command.j5
            start_command_j6 = command.j6

        # For all movement steps, use the starting j5/j6 values
        path.append(ArmJointState(
            height=command.height,
            j2=command.j2,
            j3=command.j3,
            j4=command.j4,
            j5=start_command_j5,
            j6=start_command_j6,
        ))

    # Add final grip/release adjustment if j5 or j6 changed
    final = path[-1]
    if final.j5 != j5 or final.j6 != j6:
        path.append(ArmJointState(
            height=final.height,
            j2=final.j2,
            j3=final.j3,
            j4=final.j4,
            j5=j5,
            j6=j6,
        ))

    return path


def compute_interpolation_steps(duration: float, rate: float) -> int:
    if duration <= 0.0 or rate <= 0.0:
        return 1
    return max(1, int(math.ceil(duration * rate)))


def send_udp_commands(
    ip: str,
    port: int,
    states: list[ArmJointState],
    dry_run: bool,
    duration: float,
) -> None:
    if not states:
        return

    delay = duration / len(states) if len(states) > 1 and duration > 0.0 else 0.0

    if dry_run:
        for state in states:
            print(state.to_udp_message().decode("utf-8").strip())
        return

    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
        for index, state in enumerate(states):
            sock.sendto(state.to_udp_message(), (ip, port))
            if delay > 0.0 and index < len(states) - 1:
                time.sleep(delay)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="CRAIC 机械臂 UDP 控制程序"
    )
    parser.add_argument("--ip", default=DEFAULT_UDP_IP, help="目标 UDP IP")
    parser.add_argument("--port", type=int, default=DEFAULT_UDP_PORT, help="目标 UDP 端口")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="只打印指令，不实际发送 UDP",
    )
    parser.add_argument(
        "--show-fk",
        action="store_true",
        help="输出对应关节角的 TCP 正运动学结果",
    )
    parser.add_argument(
        "--duration",
        type=float,
        default=DEFAULT_INTERP_DURATION,
        help="插值总时长（秒），默认 1.0；设为 0 则直接发送",
    )
    parser.add_argument(
        "--rate",
        type=float,
        default=DEFAULT_INTERP_RATE,
        help="插值发送频率（Hz），默认 20",
    )
    parser.add_argument(
        "--no-state-cache",
        action="store_true",
        help="不读取或更新上次发送的关节命令缓存",
    )
    parser.add_argument(
        "--height-min",
        type=int,
        default=DEFAULT_HEIGHT_MIN,
        help=f"高度下限 (内部 d1 坐标, mm)，默认 {DEFAULT_HEIGHT_MIN}",
    )
    parser.add_argument(
        "--height-max",
        type=int,
        default=DEFAULT_HEIGHT_MAX,
        help=f"高度上限 (内部 d1 坐标, mm)，默认 {DEFAULT_HEIGHT_MAX}",
    )
    parser.add_argument(
        "--joint-min",
        type=int,
        default=DEFAULT_JOINT_MIN,
        help="关节角下限，默认 -180",
    )
    parser.add_argument(
        "--joint-max",
        type=int,
        default=DEFAULT_JOINT_MAX,
        help="关节角上限（J5/J6），默认 180",
    )
    parser.add_argument(
        "--j2-min",
        type=int,
        default=DEFAULT_J2_MIN,
        help=f"J2 下限，默认 {DEFAULT_J2_MIN}",
    )
    parser.add_argument(
        "--j2-max",
        type=int,
        default=DEFAULT_J2_MAX,
        help=f"J2 上限，默认 {DEFAULT_J2_MAX}",
    )
    parser.add_argument(
        "--j3-min",
        type=int,
        default=DEFAULT_J3_MIN,
        help=f"J3 下限，默认 {DEFAULT_J3_MIN}",
    )
    parser.add_argument(
        "--j3-max",
        type=int,
        default=DEFAULT_J3_MAX,
        help=f"J3 上限，默认 {DEFAULT_J3_MAX}",
    )
    parser.add_argument(
        "--j4-min",
        type=int,
        default=DEFAULT_J4_MIN,
        help=f"J4 下限，默认 {DEFAULT_J4_MIN}",
    )
    parser.add_argument(
        "--j4-max",
        type=int,
        default=DEFAULT_J4_MAX,
        help=f"J4 上限，默认 {DEFAULT_J4_MAX}",
    )

    subparsers = parser.add_subparsers(dest="mode", required=True)

    joints = subparsers.add_parser("joints", help="直接发送关节角")
    joints.add_argument(
        "--dry-run",
        action="store_true",
        help="只打印指令，不实际发送 UDP",
    )
    joints.add_argument(
        "--show-fk",
        action="store_true",
        help="输出对应关节角的 TCP 正运动学结果",
    )
    joints.add_argument(
        "--duration",
        type=float,
        default=DEFAULT_INTERP_DURATION,
        help="插值总时长（秒），默认 1.0；设为 0 则直接发送",
    )
    joints.add_argument(
        "--rate",
        type=float,
        default=DEFAULT_INTERP_RATE,
        help="插值发送频率（Hz），默认 20",
    )
    joints.add_argument(
        "--no-state-cache",
        action="store_true",
        help="不读取或更新上次发送的关节命令缓存",
    )
    joints.add_argument("--height", type=int, required=True, help="内部 d1 坐标 mm (-290=底部, 0=顶部，z 轴朝上)")
    joints.add_argument("--j2", type=int, required=True, help="UDP 指令里的 J2 命令值")
    joints.add_argument("--j3", type=int, required=True, help="UDP 指令里的 J3 命令值")
    joints.add_argument("--j4", type=int, required=True, help="UDP 指令里的 J4 命令值")
    joints_grip = joints.add_mutually_exclusive_group()
    joints_grip.add_argument("--up", action="store_true", dest="up", default=None, help="夹爪抬起 (J5=-100°)")
    joints_grip.add_argument("--down", action="store_false", dest="up", default=None, help="夹爪放下 (J5=81°)")
    grip_release = joints.add_mutually_exclusive_group()
    grip_release.add_argument("--grip", action="store_true", help=f"抓取（J6={GRIP_ANGLE}°，待填写）")
    grip_release.add_argument("--release", action="store_true", help=f"释放（J6={RELEASE_ANGLE}°，待填写）")
    joints.add_argument("--j6", type=int, default=DEFAULT_FIXED_J6, help="UDP 指令里的 J6 命令值，默认固定 0")
    joints.add_argument("--l1", type=float, default=DEFAULT_L1, help=f"J2 到 J3 的连杆长度 mm (默认 {DEFAULT_L1})")
    joints.add_argument("--l2", type=float, default=DEFAULT_L2, help=f"J3 到 J4 的连杆长度 mm (默认 {DEFAULT_L2})")
    joints.add_argument("--x4", type=float, default=DEFAULT_X4, help=f"J4 到 TCP 的 X 偏移 mm (默认 {DEFAULT_X4})")
    joints.add_argument("--z4", type=float, default=DEFAULT_Z4, help=f"J4 到 TCP 的 Z 偏移 mm (默认 {DEFAULT_Z4})")

    pose = subparsers.add_parser("pose", help="根据末端位姿逆解后发送")
    pose.add_argument(
        "--dry-run",
        action="store_true",
        help="只打印指令，不实际发送 UDP",
    )
    pose.add_argument(
        "--show-fk",
        action="store_true",
        help="输出对应关节角的 TCP 正运动学结果",
    )
    pose.add_argument(
        "--duration",
        type=float,
        default=DEFAULT_INTERP_DURATION,
        help="插值总时长（秒），默认 1.0；设为 0 则直接发送",
    )
    pose.add_argument(
        "--rate",
        type=float,
        default=DEFAULT_INTERP_RATE,
        help="插值发送频率（Hz），默认 20",
    )
    pose.add_argument(
        "--no-state-cache",
        action="store_true",
        help="不读取或更新上次发送的关节命令缓存",
    )
    pose.add_argument("--x", type=float, required=True, help="TCP X 坐标 mm")
    pose.add_argument("--y", type=float, required=True, help="TCP Y 坐标 mm")
    pose.add_argument("--z", type=float, required=True, help="TCP Z 坐标 mm (-290=底部, 0=顶部, z 轴朝上)")
    pose.add_argument("--phi", type=float, required=True, help="TCP 偏航角，单位度；等于 J2+J3+J4")
    pose.add_argument("--l1", type=float, default=DEFAULT_L1, help=f"J2 到 J3 的连杆长度 mm (默认 {DEFAULT_L1})")
    pose.add_argument("--l2", type=float, default=DEFAULT_L2, help=f"J3 到 J4 的连杆长度 mm (默认 {DEFAULT_L2})")
    pose.add_argument("--x4", type=float, default=DEFAULT_X4, help=f"J4 到 TCP 的 X 偏移 mm (默认 {DEFAULT_X4})")
    pose.add_argument("--z4", type=float, default=DEFAULT_Z4, help=f"J4 到 TCP 的 Z 偏移 mm (默认 {DEFAULT_Z4})")
    pose.add_argument(
        "--elbow-up",
        action="store_true",
        help="使用肘部向上分支，默认使用肘部向下分支",
    )
    pose_grip_release = pose.add_mutually_exclusive_group()
    pose_grip_release.add_argument("--grip", action="store_true", help=f"抓取（J6={GRIP_ANGLE}°，待填写）")
    pose_grip_release.add_argument("--release", action="store_true", help=f"释放（J6={RELEASE_ANGLE}°，待填写）")
    pose_grip = pose.add_mutually_exclusive_group()
    pose_grip.add_argument("--up", action="store_true", dest="up", default=None, help="夹爪抬起 (J5=-100°)，覆盖 z 自动判断")
    pose_grip.add_argument("--down", action="store_false", dest="up", default=None, help="夹爪放下 (J5=81°)，覆盖 z 自动判断")
    pose.add_argument("--j6", type=int, default=DEFAULT_FIXED_J6, help="附加发送的 J6 命令值，默认固定 0")

    return parser


def geometry_from_args(args: argparse.Namespace, z4: float | None = None) -> ArmGeometry:
    return ArmGeometry(
        l1=float(args.l1),
        l2=float(args.l2),
        x4=float(args.x4),
        z4=float(args.z4) if z4 is None else z4,
    )


def limits_from_args(args: argparse.Namespace) -> ArmLimits:
    if args.height_min > args.height_max:
        raise ArmControlError("height-min 不能大于 height-max。")
    if args.joint_min > args.joint_max:
        raise ArmControlError("joint-min 不能大于 joint-max。")
    if args.j2_min > args.j2_max:
        raise ArmControlError("j2-min 不能大于 j2-max。")
    if args.j3_min > args.j3_max:
        raise ArmControlError("j3-min 不能大于 j3-max。")
    if args.j4_min > args.j4_max:
        raise ArmControlError("j4-min 不能大于 j4-max。")
    return ArmLimits(
        height_min=args.height_min,
        height_max=args.height_max,
        joint_min=args.joint_min,
        joint_max=args.joint_max,
        j2_min=args.j2_min,
        j2_max=args.j2_max,
        j3_min=args.j3_min,
        j3_max=args.j3_max,
        j4_min=args.j4_min,
        j4_max=args.j4_max,
    )


def resolve_j5(up: bool) -> int:
    return J5_CLOSED if up else J5_OPEN


def resolve_z4(up: bool) -> float:
    return Z4_CLOSED if up else Z4_OPEN


def resolve_gripper_from_z(z: float) -> bool:
    """Auto-select gripper state based on TCP z coordinate (user mm, z-up).

    - z in [-345, -55] -> down  (gripper open,  J5=81,  Z4=55)
    - z in [-190, 110] -> up   (gripper closed, J5=-100, Z4=-100)
    - Overlap [-190, -55] -> down (prefer down in intersection)
    """
    return z > -55   # True = up, False = down


def resolve_j6(grip: bool, release: bool, fallback: int) -> int:
    if grip:
        return GRIP_ANGLE
    if release:
        return RELEASE_ANGLE
    return fallback


def state_from_joint_args(args: argparse.Namespace, limits: ArmLimits) -> ArmJointState:
    up = args.up if args.up is not None else False
    return ArmJointState(
        height=clamp_int(args.height, limits.height_min, limits.height_max, "height(cmd)"),
        j2=clamp_int(args.j2, limits.j2_min, limits.j2_max, "J2"),
        j3=clamp_int(args.j3, limits.j3_min, limits.j3_max, "J3"),
        j4=clamp_int(args.j4, limits.j4_min, limits.j4_max, "J4"),
        j5=clamp_int(resolve_j5(up), limits.joint_min, limits.joint_max, "J5"),
        j6=clamp_int(resolve_j6(args.grip, args.release, args.j6), limits.joint_min, limits.joint_max, "J6"),
    )


def print_joint_summary(state: ArmJointState) -> None:
    print(
        "UDP command joints:",
        f"height={state.height}mm",
        f"J2={state.j2}",
        f"J3={state.j3}",
        f"J4={state.j4}",
        f"J5={state.j5}",
        f"J6={state.j6}",
    )


def print_pose_summary(pose: ArmPose) -> None:
    print(
        "TCP pose:",
        f"x={pose.x:.3f}mm",
        f"y={pose.y:.3f}mm",
        f"z={pose.z:.3f}mm",
        f"phi={pose.phi_deg:.3f}deg",
    )


def print_joint4_center_summary(center: Joint4Center) -> None:
    print(
        "J4 center:",
        f"x={center.x:.3f}mm",
        f"y={center.y:.3f}mm",
        f"z={center.z:.3f}mm",
    )


def print_math_summary(state: ArmMathState) -> None:
    print(
        "Math joints:",
        f"d1={state.d1:.3f}mm",
        f"J2={state.theta2_deg:.3f}",
        f"J3={state.theta3_deg:.3f}",
        f"J4={state.theta4_deg:.3f}",
    )


def print_interpolation_summary(
    duration: float,
    rate: float,
    steps: int,
    use_state_cache: bool,
) -> None:
    cache_mode = "on" if use_state_cache else "off"
    print(
        "Interpolation:",
        f"duration={duration:.3f}s",
        f"rate={rate:.3f}Hz",
        f"steps={steps}",
        f"state_cache={cache_mode}",
    )


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    try:
        limits = limits_from_args(args)
        zero_offsets = ArmZeroOffsets()
        use_state_cache = not args.no_state_cache
        steps = compute_interpolation_steps(args.duration, args.rate)
        print_interpolation_summary(args.duration, args.rate, steps, use_state_cache)

        if args.mode == "joints":
            start_command_state = resolve_start_command_state(limits, use_state_cache)
            command_state = state_from_joint_args(args, limits)
            math_state = command_to_math_state(command_state, zero_offsets)
            start_math_state = command_to_math_state(start_command_state, zero_offsets)
            command_path = interpolate_command_states(start_command_state, command_state, steps)
            print("Start state source:", "cache/default")
            print_joint_summary(start_command_state)
            print_math_summary(start_math_state)
            print_joint_summary(command_state)
            print_math_summary(math_state)

            if args.show_fk:
                up = args.up if args.up is not None else False
                z4 = resolve_z4(up)
                start_pose = forward_kinematics(
                    geometry_from_args(args, z4=z4), start_math_state,
                )
                print("Start FK:")
                print_pose_summary(start_pose)
                pose = forward_kinematics(
                    geometry_from_args(args, z4=z4), math_state,
                )
                print_pose_summary(pose)

        elif args.mode == "pose":
            if args.up is not None:
                auto_up = args.up
            else:
                auto_up = resolve_gripper_from_z(args.z)
            z4 = resolve_z4(auto_up)
            geometry = geometry_from_args(args, z4=z4)
            j5 = resolve_j5(auto_up)
            j6 = resolve_j6(args.grip, args.release, args.j6)
            start_command_state = resolve_start_command_state(limits, use_state_cache)
            start_math_state = command_to_math_state(start_command_state, zero_offsets)
            start_pose = forward_kinematics(geometry, start_math_state)
            target_pose = ArmPose(
                x=args.x,
                y=args.y,
                z=args.z,
                phi_deg=args.phi,
            )
            joint4_center = tcp_to_joint4_center(geometry, target_pose)
            math_state = inverse_kinematics(
                geometry=geometry,
                pose=target_pose,
                limits=limits,
                elbow_up=args.elbow_up,
                j5=j5,
                j6=j6,
            )
            command_state = math_to_command_state(
                math_state,
                zero_offsets,
                limits,
                j5=j5,
                j6=j6,
            )
            command_path = interpolate_command_states(start_command_state, command_state, steps)
            print("Start state source:", "cache/default")
            print_joint_summary(start_command_state)
            print_math_summary(start_math_state)
            print("Start FK:")
            print_pose_summary(start_pose)
            print_pose_summary(target_pose)
            print_joint4_center_summary(joint4_center)
            print_math_summary(math_state)
            print_joint_summary(command_state)

            if args.show_fk:
                solved_pose = forward_kinematics(geometry, math_state)
                print("Solved FK check:")
                print_pose_summary(solved_pose)

        else:
            raise ArmControlError(f"未知模式: {args.mode}")

        final_payload = command_state.to_udp_message()
        print("Final UDP payload:", final_payload.decode("utf-8").strip())
        send_udp_commands(args.ip, args.port, command_path, args.dry_run, args.duration)
        if use_state_cache and not args.dry_run:
            save_cached_command_state(command_state)
        if not args.dry_run:
            print(f"Sent to {args.ip}:{args.port}")
        return 0

    except ArmControlError as exc:
        print(f"错误: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())
