#!/usr/bin/env python3
import argparse
import random
import sys
import time
from typing import List

import rclpy
from rclpy.action import ActionClient

from geometry_msgs.msg import Pose
from interfaces.action import DualArm
from interfaces.msg import ArmMotionParams


def make_movej(arm_id: int, joints: List[float]) -> ArmMotionParams:
    msg = ArmMotionParams()
    msg.arm_id = arm_id
    msg.motion_type = ArmMotionParams.MOVEJ
    msg.target_joints = joints
    msg.blend_radius = 0
    return msg


def _format_pose(pose: Pose) -> str:
    return (
        f"pos=({pose.position.x:.3f}, {pose.position.y:.3f}, {pose.position.z:.3f}) "
        f"quat=({pose.orientation.x:.3f}, {pose.orientation.y:.3f}, "
        f"{pose.orientation.z:.3f}, {pose.orientation.w:.3f})"
    )


def send_goal(action_client: ActionClient, goal_msg: DualArm.Goal, label: str) -> bool:
    print(f"Sending goal: {label}")
    for params in goal_msg.arm_motion_params:
        if params.motion_type == ArmMotionParams.MOVEJ:
            print(f"[{label}] MOVEJ arm={params.arm_id} target_joints={list(params.target_joints)}")
        else:
            print(f"[{label}] MOVEL arm={params.arm_id} target_pose={_format_pose(params.target_pose)}")
    if not action_client.wait_for_server(timeout_sec=5.0):
        print("Action server 'DualArm' not available", file=sys.stderr)
        return False

    feedback_log = {"last": 0.0}

    def feedback_cb(feedback_msg):
        progress = feedback_msg.feedback.progress
        if progress != feedback_log["last"]:
            feedback_log["last"] = progress
            print(f"[{label}] progress={progress:.2f} status={feedback_msg.feedback.status}")

    send_future = action_client.send_goal_async(goal_msg, feedback_callback=feedback_cb)
    rclpy.spin_until_future_complete(action_client._node, send_future)
    goal_handle = send_future.result()
    if not goal_handle.accepted:
        print(f"[{label}] goal rejected", file=sys.stderr)
        return False

    result_future = goal_handle.get_result_async()
    while not result_future.done():
        rclpy.spin_once(action_client._node, timeout_sec=0.1)
    result = result_future.result().result
    print(f"[{label}] success={result.success} progress={result.final_progress:.2f} msg={result.message}")
    return result.success


def main() -> int:
    parser = argparse.ArgumentParser(description="DualArm action test client")
    parser.add_argument("--velocity-scaling", type=int, default=20)
    parser.add_argument("--dual", action="store_true", help="Send left+right in a single goal")
    parser.add_argument(
        "--mode",
        choices=["mixed", "dual", "single"],
        default="dual",
        help="Test mode: mixed=dual+single, dual=both arms, single=left/right only",
    )
    parser.add_argument("--cases", type=int, default=3)
    parser.add_argument("--joints", type=int, default=6, help="Joint count per arm for MOVEJ")
    parser.add_argument("--joint-limit", type=float, default=0.5)
    parser.add_argument("--between-wait", type=float, default=1.0)
    parser.add_argument("--loops", type=int, default=3)
    args = parser.parse_args()

    rclpy.init()
    node = rclpy.create_node("dualarm_action_test")
    action_client = ActionClient(node, DualArm, "DualArm")

    def random_joints() -> List[float]:
        return [random.uniform(-args.joint_limit, args.joint_limit) for _ in range(args.joints)]

    def generate_cases(label: str) -> List[tuple[str, DualArm.Goal]]:
        cases: List[tuple[str, DualArm.Goal]] = []
        for case_idx in range(1, max(1, args.cases) + 1):
            goal = DualArm.Goal()
            goal.velocity_scaling = args.velocity_scaling
            if label == "dual":
                goal.arm_motion_params = [
                    make_movej(ArmMotionParams.ARM_LEFT, random_joints()),
                    make_movej(ArmMotionParams.ARM_RIGHT, random_joints()),
                ]
            elif label == "left":
                goal.arm_motion_params = [
                    make_movej(ArmMotionParams.ARM_LEFT, random_joints()),
                ]
            else:
                goal.arm_motion_params = [
                    make_movej(ArmMotionParams.ARM_RIGHT, random_joints()),
                ]
            cases.append((f"{label}_case{case_idx}", goal))
        return cases

    success = True
    for loop_idx in range(1, max(1, args.loops) + 1):
        loop_cases: List[tuple[str, DualArm.Goal]] = []
        mode = "dual" if args.dual else args.mode
        if mode in ("mixed", "dual"):
            loop_cases.extend(generate_cases("dual"))
        if mode in ("mixed", "single"):
            loop_cases.extend(generate_cases("left"))
            loop_cases.extend(generate_cases("right"))

        for label, goal in loop_cases:
            ok = send_goal(action_client, goal, f"{label}_loop{loop_idx}")
            success = success and ok
            time.sleep(max(0.0, args.between_wait))

    node.destroy_node()
    rclpy.shutdown()
    return 0 if success else 1


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

