"""Execute and measure a fixed-payload transfer in an isolated Gazebo ROS domain."""

import argparse
import json
import math
import os
import time

import rclpy
from action_msgs.msg import GoalStatus
from moveit_msgs.action import MoveGroup
from moveit_msgs.msg import Constraints, JointConstraint, MoveItErrorCodes
from rclpy.action import ActionClient
from rclpy.time import Time
from sensor_msgs.msg import JointState
from tf2_ros import Buffer, TransformException, TransformListener


JOINTS = (
    "shoulder_pan_joint",
    "shoulder_lift_joint",
    "elbow_joint",
    "wrist_1_joint",
    "wrist_2_joint",
    "wrist_3_joint",
)
SOURCE = (-0.30, -1.50, 1.60, -1.67, -1.57, 0.0)
DESTINATION = (0.42, -1.50, 1.60, -1.67, -1.57, 0.0)
TOLERANCE_RAD = 0.02
MIN_DISPLACEMENT_M = 0.20


def result_of(node, future, timeout):
    rclpy.spin_until_future_complete(node, future, timeout_sec=timeout)
    if not future.done():
        raise RuntimeError("ROS action timed out")
    return future.result()


def move(node, client, positions):
    goal = MoveGroup.Goal()
    goal.request.group_name = "ur_manipulator"
    goal.request.num_planning_attempts = 3
    goal.request.allowed_planning_time = 5.0
    goal.request.max_velocity_scaling_factor = 0.2
    goal.request.max_acceleration_scaling_factor = 0.2
    goal.request.start_state.is_diff = True
    goal.request.goal_constraints = [
        Constraints(
            joint_constraints=[
                JointConstraint(
                    joint_name=name,
                    position=position,
                    tolerance_above=0.01,
                    tolerance_below=0.01,
                    weight=1.0,
                )
                for name, position in zip(JOINTS, positions)
            ]
        )
    ]
    handle = result_of(node, client.send_goal_async(goal), 10)
    if not handle.accepted:
        raise RuntimeError("MoveIt rejected the goal")
    response = result_of(node, handle.get_result_async(), 35)
    if response.status != GoalStatus.STATUS_SUCCEEDED:
        raise RuntimeError(f"MoveIt action status {response.status}, code {response.result.error_code.val}")
    if response.result.error_code.val != MoveItErrorCodes.SUCCESS:
        raise RuntimeError(f"MoveIt error code {response.result.error_code.val}")


def measure(node, buffer, joints, goal):
    deadline = time.monotonic() + 5
    while time.monotonic() < deadline:
        rclpy.spin_once(node, timeout_sec=0.1)
        if not all(name in joints for name in JOINTS):
            continue
        error = max(abs(joints[name] - expected) for name, expected in zip(JOINTS, goal))
        if error > TOLERANCE_RAD:
            continue
        try:
            pose = buffer.lookup_transform("base_link", "maloni_payload", Time()).transform.translation
            return (pose.x, pose.y, pose.z), error
        except TransformException:
            continue
    raise RuntimeError("Joint tolerance or payload transform unavailable")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--runs", type=int, default=10)
    args = parser.parse_args()
    if not 1 <= args.runs <= 50:
        parser.error("runs must be between 1 and 50")
    if os.environ.get("ROS_DOMAIN_ID") != "85" or os.environ.get("ROS_LOCALHOST_ONLY") != "1":
        parser.error("launch only in the isolated simulation domain from README.md")

    rclpy.init()
    node = rclpy.create_node("maloni_gazebo_transfer")
    client = ActionClient(node, MoveGroup, "/move_action")
    buffer = Buffer()
    listener = TransformListener(buffer, node)
    joints = {}
    subscription = node.create_subscription(
        JointState, "/joint_states", lambda msg: joints.update(zip(msg.name, msg.position)), 10
    )
    try:
        if not client.wait_for_server(timeout_sec=25):
            raise RuntimeError("MoveIt action server unavailable")
        for attempt in range(1, args.runs + 1):
            stage = "source"
            started = time.monotonic()
            report = {"run": attempt, "success": False}
            try:
                move(node, client, SOURCE)
                start, start_error = measure(node, buffer, joints, SOURCE)
                stage = "destination"
                started = time.monotonic()
                move(node, client, DESTINATION)
                end, end_error = measure(node, buffer, joints, DESTINATION)
                displacement = math.dist(start, end)
                report.update(
                    source_xyz_m=start, destination_xyz_m=end, displacement_m=displacement,
                    max_joint_error_rad=max(start_error, end_error),
                )
                if displacement < MIN_DISPLACEMENT_M:
                    raise RuntimeError("Payload displacement below the test threshold")
                report["success"] = True
            except RuntimeError as exc:
                report.update(stage=stage, reason=str(exc))
            print(json.dumps({**report, "elapsed_wall_s": time.monotonic() - started}), flush=True)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()