from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
from time import perf_counter, sleep
from typing import Optional, Sequence


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Visualize first-stage IK in the normalized dual RM75 scene"
    )
    parser.add_argument("--arm", choices=("left", "right"), default="left")
    parser.add_argument(
        "--trajectory",
        choices=("joint", "line", "arc", "orientation", "combined"),
        default="combined",
    )
    parser.add_argument("--points", type=int, default=180)
    parser.add_argument("--dt", type=float, default=1.0 / 90.0)
    parser.add_argument(
        "--target-position",
        nargs=3,
        type=float,
        metavar=("X", "Y", "Z"),
        help="target position in the controlled-arm base frame, in meters",
    )
    parser.add_argument(
        "--target-rpy",
        nargs=3,
        type=float,
        metavar=("ROLL", "PITCH", "YAW"),
        help="target RPY in radians; defaults to the start orientation",
    )
    parser.add_argument(
        "--duration",
        type=float,
        default=8.0,
        help="target-point motion duration in seconds",
    )
    parser.add_argument("--wait-before", type=float, default=2.0)
    parser.add_argument("--hold-after", type=float, default=3.0)
    parser.add_argument(
        "--manual-drag",
        action="store_true",
        help="enable zero-gravity mouse perturbation of the controlled arm",
    )
    parser.add_argument(
        "--drag-damping",
        type=float,
        default=1.5,
        help="joint damping used by manual-drag mode",
    )
    parser.add_argument(
        "--headless",
        action="store_true",
        help="render the final frame without opening the interactive viewer",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("stage2_demo.png"),
        help="headless output image",
    )
    parser.add_argument(
        "--close-after-play",
        action="store_true",
        help="close the interactive viewer after one playback (smoke testing)",
    )
    return parser


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    if args.manual_drag and args.headless:
        raise SystemExit("--manual-drag requires the interactive viewer")
    if args.target_rpy is not None and args.target_position is None:
        raise SystemExit("--target-rpy requires --target-position")
    for name in ("dt", "duration", "wait_before", "hold_after"):
        if getattr(args, name) < 0.0 or (name in {"dt", "duration"} and getattr(args, name) == 0.0):
            raise SystemExit(f"--{name.replace('_', '-')} must be positive")
    if args.headless:
        os.environ.setdefault("MUJOCO_GL", "egl")
    else:
        os.environ.setdefault("MUJOCO_GL", "glfw")

    import mujoco.viewer
    import numpy as np
    import pinocchio as pin
    from PIL import Image

    from .dual_arm import DualArmAssembly
    from .mujoco_backend import CONTROLLED_HOME_Q_RAD, DualArmMuJoCo
    from .mujoco_trajectories import (
        cartesian_demo_targets,
        joint_demo_trajectory,
        se3_target_trajectory,
        solve_pose_trajectory,
    )
    from .kinematics import RM75Kinematics, pose_errors
    from .solver import RM75IkSolver
    from .types import teleop_joint_limits

    scene = DualArmMuJoCo(controlled_arm=args.arm)
    if args.manual_drag:
        scene.configure_manual_drag(args.drag_damping)
        print("Manual drag mode")
        print("  1. Double-click a link to select it.")
        print("  2. Hold Ctrl and drag with the right mouse button.")
        print("  3. Close the viewer or press Ctrl+C to finish.")
        viewer = mujoco.viewer.launch_passive(scene.model, scene.data)
        try:
            while viewer.is_running():
                step_started = perf_counter()
                scene.step_manual_drag()
                viewer.sync()
                remaining = scene.model.opt.timestep - (perf_counter() - step_started)
                if remaining > 0.0:
                    sleep(remaining)
        except KeyboardInterrupt:
            pass
        finally:
            viewer.close()
            sleep(0.2)
        print("Final controlled-arm joints (rad):")
        print(scene.get_arm_configuration(args.arm).tolist())
        return 0

    kinematics = RM75Kinematics(limits=teleop_joint_limits())
    solver = RM75IkSolver(kinematics)
    targets = None
    target_pose = None
    if args.target_position is not None:
        start_pose = kinematics.forward(CONTROLLED_HOME_Q_RAD)
        rotation = (
            start_pose.rotation
            if args.target_rpy is None
            else pin.rpy.rpyToMatrix(*args.target_rpy)
        )
        target_pose = pin.SE3(rotation, np.asarray(args.target_position, dtype=float))
        target_points = max(2, int(round(args.duration / args.dt)) + 1)
        try:
            targets = se3_target_trajectory(start_pose, target_pose, target_points)
            trajectory = solve_pose_trajectory(
                solver, targets, CONTROLLED_HOME_Q_RAD
            )
        except (RuntimeError, TypeError, ValueError) as exc:
            raise SystemExit(f"target trajectory rejected: {exc}") from exc
        print("Target mode")
        print("  position (m):", target_pose.translation.tolist())
        print("  rpy (rad):", pin.rpy.matrixToRpy(target_pose.rotation).tolist())
        print(f"  duration: {args.duration:.3f} s, points: {target_points}")
    elif args.trajectory == "joint":
        trajectory = joint_demo_trajectory(CONTROLLED_HOME_Q_RAD, args.points)
    else:
        targets = cartesian_demo_targets(
            args.trajectory,
            kinematics.forward(CONTROLLED_HOME_Q_RAD),
            args.points,
        )
        trajectory = solve_pose_trajectory(
            solver, targets, CONTROLLED_HOME_Q_RAD
        )

    assembly = DualArmAssembly.from_source_urdf()
    mount = (
        assembly.mounts.left_base
        if args.arm == "left"
        else assembly.mounts.right_base
    )
    if target_pose is not None:
        scene.set_target_marker(mount * target_pose)
    elif targets is not None:
        scene.set_target_marker(mount * targets[-1])
    else:
        scene.set_arm_configuration(args.arm, trajectory[-1])
        scene.set_target_marker(scene.get_flange_pose(args.arm))
        scene.set_arm_configuration(args.arm, trajectory[0])

    if args.headless:
        result = scene.play_trajectory(trajectory, dt=args.dt, realtime=False)
        args.output.parent.mkdir(parents=True, exist_ok=True)
        Image.fromarray(scene.render()).save(args.output)
        print(
            f"Played {result.samples} samples; max joint step "
            f"{result.max_joint_step_rad:.6f} rad"
        )
        if target_pose is not None:
            actual_local = mount.actInv(result.final_flange_pose)
            position_error, orientation_error = pose_errors(actual_local, target_pose)
            print(f"Final position error: {position_error:.9f} m")
            print(f"Final orientation error: {orientation_error:.9f} rad")
        print(args.output)
        return 0

    viewer = mujoco.viewer.launch_passive(scene.model, scene.data)
    try:
        wait_until = perf_counter() + args.wait_before
        while viewer.is_running() and perf_counter() < wait_until:
            viewer.sync()
            sleep(0.02)
        scene.play_trajectory(
            trajectory,
            dt=args.dt,
            realtime=True,
            viewer=viewer,
        )
        if target_pose is not None:
            actual_local = mount.actInv(scene.get_flange_pose(args.arm))
            position_error, orientation_error = pose_errors(actual_local, target_pose)
            print(f"Final position error: {position_error:.9f} m")
            print(f"Final orientation error: {orientation_error:.9f} rad")
        hold_until = perf_counter() + args.hold_after
        while viewer.is_running() and perf_counter() < hold_until:
            viewer.sync()
            sleep(0.02)
        if not args.close_after_play:
            while viewer.is_running():
                sleep(0.05)
    except KeyboardInterrupt:
        pass
    finally:
        viewer.close()
        # GLFW tears down asynchronously; allow its render thread to exit.
        sleep(0.2)
    return 0


if __name__ == "__main__":
    sys.exit(main())
