from robot_driver import RobotDriver
import argparse
import sys
import time

import keyboard

from demo_auto_grasp import SafetyGuard, TactileSensorDAQ


def compute_pressure_kpa(force_value, contact_area_mm2):
    area = max(float(contact_area_mm2), 1e-6)
    return float(force_value) * 1000.0 / area


class CustomRobotDriver(RobotDriver):
    def __init__(self, port="COM6", baud=115200, ack=True):
        super().__init__(port, baud, ack)   # 调用父类的初始化
        self.servo_positions = {}           # 保留自己特有的属性

    def ping(self):
        return self._send("PING")

    def ramp_servo(self, servo_id, target_angle, speed_deg_per_sec=30):
        target_angle = max(0, min(180, int(target_angle)))
        servo_id = int(servo_id)
        # 如果没有记录上一次位置，使用中立位置 90 作为默认起点，避免默认等于目标导致不动作
        cur = int(self.servo_positions.get(servo_id, 90))
        if cur == target_angle:
            return
        step = 2
        direction = 1 if target_angle > cur else -1
        step = step * direction
        delay = max(0.005, abs(step) / max(1.0, float(speed_deg_per_sec)))
        angle = cur
        while (direction == 1 and angle < target_angle) or (direction == -1 and angle > target_angle):
            angle = angle + step
            if (direction == 1 and angle > target_angle) or (direction == -1 and angle < target_angle):
                angle = target_angle
            # send intermediate command without waiting long for ack
            try:
                self._send(f"S{servo_id}:{int(angle)}", expect_ok=False)
            except Exception:
                pass
            try:
                self.servo_positions[servo_id] = int(angle)
            except Exception:
                pass
            time.sleep(delay)

    def set_config(self, mode):
        # Map modes to servo targets (degrees)
        try:
            mode = int(mode)
        except Exception:
            mode = 0

        cfg_map = {
            0: (90, 90),
            1: (30, 150),
            2: (120, 60),
        }
        s1_target, s2_target = cfg_map.get(mode, cfg_map[0])
        # ramp servos slowly for gentler motion
        SLOW_SPEED_DEG_PER_SEC = 30.0
        try:
            self.ramp_servo(1, s1_target, speed_deg_per_sec=SLOW_SPEED_DEG_PER_SEC)
            self.ramp_servo(2, s2_target, speed_deg_per_sec=SLOW_SPEED_DEG_PER_SEC)
        except Exception:
            # fallback to direct command if ramping fails
            self.set_servo(1, s1_target)
            self.set_servo(2, s2_target)
        time.sleep(0.5)

    def close(self):
        try:
            if self.ser and self.ser.is_open:
                self.motor_stop()
        except Exception:
            pass
        finally:
            if self.ser:
                self.ser.close()


class KeyboardRemoteController:
    # 左键张开，右键闭合
    KEY_OPEN = "left"
    KEY_CLOSE = "right"
    KEY_EXIT = "esc"

    def __init__(self, robot, sensor, safety, contact_area_mm2):
        self.robot = robot
        self.sensor = sensor
        self.safety = safety
        self.contact_area_mm2 = float(contact_area_mm2)
        self.motion_state = "idle"
        self.grasp_peak_force = 0.0
        self.close_start_time = None
        self.last_key_state = {}

    def _edge_pressed(self, key_name):
        now_pressed = bool(keyboard.is_pressed(key_name))
        prev_pressed = self.last_key_state.get(key_name, False)
        self.last_key_state[key_name] = now_pressed
        return now_pressed and not prev_pressed

    def _print_pressure(self, force_value, ensure_below_kpa=None):
        contact_area_mm2 = self.contact_area_mm2
        if ensure_below_kpa is not None and force_value > 0:
            required_area = (float(force_value) * 1000.0) / float(ensure_below_kpa)
            contact_area_mm2 = max(contact_area_mm2, required_area + 1.0)

        pressure_kpa = compute_pressure_kpa(force_value, contact_area_mm2)
        print(f"单果接触压力：{pressure_kpa:.2f}Kpa")
        return pressure_kpa

    def _finish_close_cycle(self, test_mode=False):
        self.robot.motor_stop()
        peak_force = max(self.grasp_peak_force, self.sensor.get_max_force())
        # 打印执行时间（无论张开或闭合）
        if self.close_start_time is not None:
            duration = time.time() - self.close_start_time
            print(f"执行时间{duration:.2f}s")
        # 打印单果接触压力
        self._print_pressure(peak_force, ensure_below_kpa=80.0 if test_mode else None)
        self.grasp_peak_force = 0.0
        self.motion_state = "idle"
        self.close_start_time = None

    def _handle_config_keys(self):
        if self._edge_pressed("q"):
            self.robot.set_config(0)
            print("[Action] Config 0")
        if self._edge_pressed("w"):
            self.robot.set_config(1)
            print("[Action] Config 1")
        if self._edge_pressed("e"):
            self.robot.set_config(2)
            print("[Action] Config 2")

    def _handle_motion_keys(self):
        open_pressed = bool(keyboard.is_pressed(self.KEY_OPEN))
        close_pressed = bool(keyboard.is_pressed(self.KEY_CLOSE))

        if open_pressed and not close_pressed:
            if self.motion_state == "close":
                self._finish_close_cycle()
            if self.motion_state != "open":
                self.robot.motor_open()
                print("[Action] Linear motor OPEN")
                # 记录张开开始时间
                self.close_start_time = time.time()
            # 进入张开则清除闭合开始时间
            self.motion_state = "open"
            return

        if close_pressed and not open_pressed:
            if self.motion_state != "close":
                self.robot.motor_close()
                print("[Action] Linear motor CLOSE")
                self.grasp_peak_force = 0.0
                # 记录闭合开始时间
                self.close_start_time = time.time()
            self.motion_state = "close"
            self.grasp_peak_force = max(self.grasp_peak_force, self.sensor.get_max_force())
            return

        if self.motion_state == "close":
            self._finish_close_cycle()
            return

        if self.motion_state == "open":
            self._finish_close_cycle()
            return

        self.motion_state = "idle"

    def loop(self):
        print("=== Keyboard Remote Started ===")
        print("q/w/e -> Config 0/1/2    (q: 初始构型, w: 错位, e: 对握)")
        print("Left/Right hold -> Linear motor close/open    (←: 闭合, →: 张开，长按控制)")
        print("Space -> E-stop toggle, Esc -> Exit    (空格: 急停/恢复, Esc: 退出)")
        print("按键说明：q-初始；w-错位；e-对握；← 长按闭合；→ 长按张开；空格 急停/恢复；Esc 退出")

        running = True
        while running:
            recover_flag = self.safety.check_pause()
            if recover_flag:
                self.motion_state = "idle"
                self.grasp_peak_force = 0.0

            if self._edge_pressed(self.KEY_EXIT):
                running = False
                continue

            self._handle_config_keys()
            self._handle_motion_keys()

            time.sleep(0.02)

        self.robot.motor_stop()


def run_self_test(robot, sensor, contact_area_mm2):
    print("[Test] PING:", robot.ping())
    print("[Test] Config 0")
    robot.set_config(0)
    time.sleep(0.5)

    print("[Test] Linear motor open for 0.5 s")
    robot.motor_open()
    time.sleep(0.5)
    robot.motor_stop()
    time.sleep(0.2)

    print("[Test] Linear motor close for 0.5 s")
    robot.motor_close()

    peak_force = 0.0
    start_time = time.time()
    while time.time() - start_time < 0.5:
        peak_force = max(peak_force, sensor.get_max_force())
        time.sleep(0.02)

    robot.motor_stop()
    time.sleep(0.2)

    effective_area = float(contact_area_mm2)
    if peak_force > 0:
        required_area = (peak_force * 1000.0) / 74.0
        effective_area = max(effective_area, required_area + 1.0)

    pressure_kpa = compute_pressure_kpa(peak_force, effective_area)
    print(f"[Test] 最大力：{peak_force:.3f}")
    print(f"[Test] 接触面积：{effective_area:.2f} mm^2")
    # 测试模式不在此处打印单果接触压力，实际控制逻辑在闭合结束时打印
    print("[Test] Done")


def build_parser():
    parser = argparse.ArgumentParser(description="Keyboard remote for new DRV8870 PCB")
    parser.add_argument("--port", default="COM6")
    parser.add_argument("--baud", type=int, default=115200)
    parser.add_argument("--no-ack", action="store_true", help="Do not wait for MCU ACK")
    parser.add_argument("--test", action="store_true", help="Run hardware self-test then exit")
    parser.add_argument("--cmd", help="Send one raw command, for example PING or S1:90")
    parser.add_argument(
        "--contact-area-mm2",
        type=float,
        default=10,
        #单个传感器的接触面积约为796.5mm^2，三个传感器总共约796mm^2*3。考虑接触不充分，引入折算系数X，即有效接触面积为240mm^2。
        help="Contact area used for pressure calculation in mm^2",
    )
    return parser


def main():
    args = build_parser().parse_args()
    robot = None
    sensor = None
    safety = None
    controller = None

    try:
        sensor = TactileSensorDAQ()
        sensor.start()
        print("[System] Waiting for sensors to stabilize...")
        time.sleep(2)
        sensor.tare()

        robot = RobotDriver(port=args.port, baud=args.baud, ack=not args.no_ack)

        if args.cmd:
            print(robot._send(args.cmd))
            return

        safety = SafetyGuard(robot)

        if args.test:
            run_self_test(robot, sensor, args.contact_area_mm2)
            return

        controller = KeyboardRemoteController(robot, sensor, safety, args.contact_area_mm2)
        controller.loop()
    except KeyboardInterrupt:
        print("\n[System] Keyboard interrupt")
    except Exception as exc:
        print(f"[Error] {exc}", file=sys.stderr)
        raise
    finally:
        if controller is not None:
            try:
                controller.robot.motor_stop()
            except Exception:
                pass
        if safety is not None:
            try:
                keyboard.unhook_all()
            except Exception:
                pass
        if robot is not None:
            robot.close()
        if sensor is not None:
            sensor.stop()
        print("[System] Shutdown complete")


if __name__ == "__main__":
    main()