"""Hardware-free LeRobot loop: observe -> compose action -> send, at 30 Hz.

Only the robot implementation and input source differ from the upstream example.
This is real-time control, NOT lockstep RL or dataset recording.
"""

import argparse
import json
import math
import os
import time
from importlib.metadata import version

from lerobot.robots.config import RobotConfig
from lerobot.robots.utils import make_robot_from_config
from lerobot.utils.import_utils import register_third_party_plugins


def make_robot(endpoint):
    register_third_party_plugins()
    config_type = RobotConfig.get_choice_class("lekiwi_sim")
    return make_robot_from_config(config_type(endpoint=endpoint, id="lekiwi-sim-demo"))


def run_demo(endpoint, duration=10.0):
    if not math.isfinite(duration) or not 1 <= duration <= 3600:
        raise ValueError("演示时长必须为1–3600秒")
    robot = make_robot(endpoint)
    robot.connect()
    try:
        start = time.monotonic()
        first = robot.get_sim_observation()
        latencies, gaps, max_motion = [], [], 0.0
        steps = math.ceil(duration * 30)
        previous_sim_time = first["simTime"]
        for i in range(steps):
            tick = time.monotonic()
            observation = robot.get_observation()
            t = i / 30
            action = {
                "arm_shoulder_pan.pos": 10 * math.sin(t),
                "arm_shoulder_lift.pos": 6 * math.sin(t + 0.2),
                "arm_elbow_flex.pos": 8 * math.sin(t + 0.4),
                "arm_wrist_flex.pos": 6 * math.sin(t + 0.6),
                "arm_wrist_roll.pos": 12 * math.sin(t + 0.8),
                "arm_gripper.pos": 50 + 25 * math.sin(t),
                "x.vel": 0.05 * math.cos(t),
                "y.vel": 0.05 * math.sin(t),
                "theta.vel": 10 * math.sin(t / 2),
            }
            accepted = robot.send_action(action)
            platform = robot.get_sim_observation()
            latencies.append((time.monotonic() - tick) * 1000)
            gaps.append(platform["simTime"] - previous_sim_time)
            previous_sim_time = platform["simTime"]
            max_motion = max(
                max_motion,
                math.hypot(
                    platform["values"]["base.x"] - first["values"]["base.x"],
                    platform["values"]["base.y"] - first["values"]["base.y"],
                ),
            )
            time.sleep(max(0, start + (i + 1) / 30 - time.monotonic()))
        ordered = sorted(latencies)
        return {
            "lerobot": version("lerobot"),
            "torch": version("torch"),
            "steps": steps,
            "elapsedSeconds": time.monotonic() - start,
            "simSeconds": platform["simTime"] - first["simTime"],
            "rttMeanMs": sum(latencies) / len(latencies),
            "rttP95Ms": ordered[int(0.95 * (len(ordered) - 1))],
            "rttMaxMs": max(latencies),
            "maxObservationSimGap": max(gaps),
            "maxTranslationM": max_motion,
            "observation": observation,
            "accepted": accepted,
            "sequence": platform["sequence"],
            "appliedActionSeq": platform["appliedActionSeq"],
            "timeouts": 0,
            "droppedRequests": 0,
        }
    finally:
        robot.disconnect()


def main():
    parser = argparse.ArgumentParser(description="LeKiwi 仿真：先播放并在浏览器显式授权")
    parser.add_argument(
        "--endpoint", default=os.environ.get("MUJOCO_CONTROL_ENDPOINT", "http://127.0.0.1:8766")
    )
    parser.add_argument("--duration", type=float, default=10)
    args = parser.parse_args()
    print(json.dumps(run_demo(args.endpoint, args.duration), ensure_ascii=False))


if __name__ == "__main__":
    main()