"""Linux terminal base + joint-space arm teleop; no hardware or pynput required."""

import argparse
import math
import os
import select
import sys
import termios
import time
import tty

CONTROL_HZ = 30
KEY_PULSE_SECONDS = 0.18
# Positive / negative in LeRobot units: degrees, or 0–100 gripper opening.
ARM_KEYS = {
    "arm_shoulder_pan.pos": ("u", "j"),
    "arm_shoulder_lift.pos": ("i", "k"),
    "arm_elbow_flex.pos": ("o", "l"),
    "arm_wrist_flex.pos": ("t", "g"),
    "arm_wrist_roll.pos": ("y", "h"),
    "arm_gripper.pos": ("v", "b"),
}


class KeyboardController:
    """Compose base velocities and incremental arm targets in one writer/lease."""

    def __init__(self, robot, arm_speed=20.0, gripper_speed=50.0):
        self.robot = robot
        self.arm_speed = arm_speed
        self.gripper_speed = gripper_speed
        self.targets = robot.get_observation()  # Never jump to a hard-coded zero pose.
        self.active = {}
        self.speed_keys = {robot.teleop_keys[k] for k in ("speed_up", "speed_down")}
        self.motion_keys = {
            robot.teleop_keys[k]
            for k in ("forward", "backward", "left", "right", "rotate_left", "rotate_right")
        } | {key for pair in ARM_KEYS.values() for key in pair}

    def step(self, keys, now):
        keys = set(keys)
        if " " in keys:
            self.active.clear()
            pressed = set()  # Space wins over every movement/speed key in the batch.
        else:
            self.active = {key: expiry for key, expiry in self.active.items() if expiry > now}
            self.active.update(dict.fromkeys(keys & self.motion_keys, now + KEY_PULSE_SECONDS))
            # Speed changes are input events, not a 180 ms pulse repeated every frame.
            pressed = set(self.active) | (keys & self.speed_keys)
        self.robot.get_observation()  # Require fresh feedback even while holding still.
        action = self.robot._from_keyboard_to_base_action(pressed)
        for channel, (positive, negative) in ARM_KEYS.items():
            direction = int(positive in pressed) - int(negative in pressed)
            if direction:
                speed = self.gripper_speed if channel == "arm_gripper.pos" else self.arm_speed
                # At most one control tick per update: no large catch-up jump after a stall.
                action[channel] = self.targets[channel] + direction * speed / CONTROL_HZ
        # The plugin holds omitted arm channels and clamps via the browser descriptor.
        # Accumulate from confirmed targets only, including any joint/gripper clipping.
        self.targets = self.robot.send_action(action)
        return self.targets


def read_keys(fd):
    keys = []
    # Bound input draining so queued terminal data cannot starve the control watchdog.
    for _ in range(64):
        if not select.select([fd], [], [], 0)[0]:
            break
        value = os.read(fd, 1)
        if not value:
            raise EOFError("终端输入已断开")
        keys.append(value.decode(errors="ignore"))
    return keys


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(
        "--arm-speed", type=float, default=20.0, help="臂关节点动速度，0–90 度/秒，默认20"
    )
    parser.add_argument(
        "--gripper-speed", type=float, default=50.0, help="夹爪点动速度，0–100 百分点/秒，默认50"
    )
    args = parser.parse_args()
    for name, speed, limit in (
        ("--arm-speed", args.arm_speed, 90),
        ("--gripper-speed", args.gripper_speed, 100),
    ):
        if not math.isfinite(speed) or not 0 < speed <= limit:
            parser.error(f"{name} 必须为大于0且不超过{limit}的有限数值")
    if not sys.stdin.isatty():
        raise ValueError("请在交互式终端运行；自动测试请用 demo_control.py")
    # Keep --help and argument validation usable without importing the LeRobot stack.
    from demo_control import make_robot

    fd = sys.stdin.fileno()
    previous = termios.tcgetattr(fd)
    robot = make_robot(args.endpoint)
    try:
        robot.connect()
        controller = KeyboardController(robot, args.arm_speed, args.gripper_speed)
        tty.setcbreak(fd)
        print(
            "底盘：w/s 前后，a/d 左右，z/x 旋转，r/f 底盘调速；空格停止点动，q 退出。\n"
            "机械臂（前键增大/后键减小）：u/j 肩转，i/k 肩俯仰，o/l 肘，t/g 腕俯仰，"
            "y/h 腕旋转；v/b 夹爪开/合。\n"
            f"臂 {args.arm_speed:g} 度/秒，夹爪 {args.gripper_speed:g} 百分点/秒。"
            "终端按键为180ms脉冲；脉冲结束后底盘停止、机械臂保持目标。",
            flush=True,
        )
        while True:
            start = time.monotonic()
            keys = read_keys(fd)
            if robot.teleop_keys["quit"] in keys:
                return
            controller.step(keys, start)
            time.sleep(max(0, 1 / CONTROL_HZ - (time.monotonic() - start)))
    except KeyboardInterrupt:
        pass
    finally:
        try:
            robot.disconnect()
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, previous)


if __name__ == "__main__":
    main()