import time
import serial
import pygame


class RobotDriver:
    def __init__(self, port="COM9", baud=115200):
        self.ser = None
        try:
            self.ser = serial.Serial(port, baud, timeout=0.2)
            print(f"[System] Robot connected on {port} @ {baud}")
            time.sleep(2.0)
        except Exception as e:
            print(f"[Error] Robot connection failed: {e}")

    def _send(self, cmd):
        if self.ser is None:
            return
        packet = f"{cmd}\r\n"
        self.ser.write(packet.encode("ascii"))
        self.ser.flush()

    def motor_open(self):
        self._send("M:OPEN")

    def motor_close(self):
        self._send("M:CLOSE")

    def motor_stop(self):
        self._send("M:STOP")

    def set_servo(self, servo_id, angle):
        self._send(f"S{servo_id}:{angle}")

    def set_config(self, mode):
        # 与原脚本保持一致的三种构型
        if mode == 0:
            self.set_servo(1, 90)
            self.set_servo(2, 90)
        elif mode == 1:
            self.set_servo(1, 30)
            self.set_servo(2, 150)
        elif mode == 2:
            self.set_servo(1, 120)
            self.set_servo(2, 60)

    def close(self):
        self.motor_stop()
        if self.ser is not None:
            self.ser.close()
            self.ser = None


class GamepadRemoteController:
    # 默认按键映射（Xbox 常见布局）
    # A/B/X 控制构型 0/1/2
    # LB 打开，RB 闭合
    # START 退出
    BTN_A = 0
    BTN_B = 1
    BTN_X = 2
    BTN_LB = 4
    BTN_RB = 5
    BTN_BACK = 6
    BTN_START = 7

    def __init__(self, robot):
        self.robot = robot
        self.joystick = None
        self.estop_active = False

        self.last_open_pressed = False
        self.last_close_pressed = False

        self.last_mode_buttons = {
            self.BTN_A: False,
            self.BTN_B: False,
            self.BTN_X: False,
            self.BTN_BACK: False,
        }

    def init_gamepad(self):
        pygame.init()
        pygame.joystick.init()

        if pygame.joystick.get_count() < 1:
            raise RuntimeError("No gamepad detected. Please connect controller via Bluetooth first.")

        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()
        print(f"[System] Gamepad connected: {self.joystick.get_name()}")

    def _edge_pressed(self, btn_id, cache_dict):
        now_pressed = bool(self.joystick.get_button(btn_id))
        prev_pressed = cache_dict.get(btn_id, False)
        cache_dict[btn_id] = now_pressed
        return now_pressed and (not prev_pressed)

    def _handle_mode_buttons(self):
        if self._edge_pressed(self.BTN_A, self.last_mode_buttons):
            self.robot.set_config(0)
            print("[Action] Config 0")
        if self._edge_pressed(self.BTN_B, self.last_mode_buttons):
            self.robot.set_config(1)
            print("[Action] Config 1")
        if self._edge_pressed(self.BTN_X, self.last_mode_buttons):
            self.robot.set_config(2)
            print("[Action] Config 2")

    def _handle_estop_toggle(self):
        if not self._edge_pressed(self.BTN_BACK, self.last_mode_buttons):
            return

        self.estop_active = not self.estop_active
        self.robot.motor_stop()

        if self.estop_active:
            print("\n" + "!" * 40)
            print("!!! EMERGENCY STOP ACTIVE !!!")
            print("!!! Motion locked. Press BACK again to resume.")
            print("!" * 40)
        else:
            print("\n" + "=" * 40)
            print(">>> EMERGENCY STOP RELEASED")
            print("= " * 20)

    def _handle_open_close_hold(self):
        # 按住 LB 持续张开，按住 RB 持续闭合；松开则停止
        open_pressed = bool(self.joystick.get_button(self.BTN_LB))
        close_pressed = bool(self.joystick.get_button(self.BTN_RB))

        if open_pressed and not close_pressed:
            if not self.last_open_pressed or self.last_close_pressed:
                self.robot.motor_open()
                print("[Action] Palm OPEN")
        elif close_pressed and not open_pressed:
            if not self.last_close_pressed or self.last_open_pressed:
                self.robot.motor_close()
                print("[Action] Palm CLOSE")
        else:
            if self.last_open_pressed or self.last_close_pressed:
                self.robot.motor_stop()
                print("[Action] Palm STOP")

        self.last_open_pressed = open_pressed
        self.last_close_pressed = close_pressed

    def loop(self):
        print("\n=== Gamepad Remote Started ===")
        print("A -> Config 0")
        print("B -> Config 1")
        print("X -> Config 2")
        print("LB hold -> Palm OPEN")
        print("RB hold -> Palm CLOSE")
        print("Release LB/RB -> STOP")
        print("BACK -> Global E-STOP toggle")
        print("START -> Exit")

        running = True
        while running:
            pygame.event.pump()

            self._handle_estop_toggle()

            if self.joystick.get_button(self.BTN_START):
                print("[System] Exit requested by START")
                running = False
                continue

            if self.estop_active:
                # 急停锁定期间屏蔽所有动作命令
                self.last_open_pressed = bool(self.joystick.get_button(self.BTN_LB))
                self.last_close_pressed = bool(self.joystick.get_button(self.BTN_RB))
                time.sleep(0.02)
                continue

            self._handle_mode_buttons()
            self._handle_open_close_hold()

            time.sleep(0.02)

        self.robot.motor_stop()

    def shutdown(self):
        try:
            if self.joystick is not None:
                self.joystick.quit()
        finally:
            pygame.joystick.quit()
            pygame.quit()


def main():
    # 根据实际串口修改
    port = "COM9"
    baud = 115200

    robot = RobotDriver(port=port, baud=baud)
    if robot.ser is None:
        return

    controller = GamepadRemoteController(robot)

    try:
        controller.init_gamepad()
        controller.loop()
    except KeyboardInterrupt:
        print("\n[System] Keyboard interrupt")
    except Exception as e:
        print(f"[Error] {e}")
    finally:
        controller.shutdown()
        robot.close()
        print("[System] Shutdown complete.")


if __name__ == "__main__":
    main()