import argparse
import sys
import time

import serial
from robot_driver import RobotDriver


class GamepadRemoteController:
    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_buttons = {
            self.BTN_A: False,
            self.BTN_B: False,
            self.BTN_X: False,
            self.BTN_BACK: False,
        }

    def init_gamepad(self):
        global pygame
        import pygame

        pygame.init()
        pygame.joystick.init()

        if pygame.joystick.get_count() < 1:
            raise RuntimeError("No gamepad detected")

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

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

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

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

        self.estop_active = not self.estop_active
        self.robot.motor_stop()
        print("[Action] E-STOP ON" if self.estop_active else "[Action] E-STOP OFF")

    def _handle_open_close_hold(self):
        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] Linear motor 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] Linear motor CLOSE")
        else:
            if self.last_open_pressed or self.last_close_pressed:
                self.robot.motor_stop()
                print("[Action] Linear motor STOP")

        self.last_open_pressed = open_pressed
        self.last_close_pressed = close_pressed

    def loop(self):
        print("=== Gamepad Remote Started ===")
        print("A/B/X -> Config 0/1/2")
        print("LB/RB hold -> Linear motor open/close")
        print("BACK -> E-stop toggle, START -> Exit")

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

            if self.joystick.get_button(self.BTN_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):
        if self.joystick is not None:
            self.joystick.quit()
        pygame.joystick.quit()
        pygame.quit()


def run_self_test(robot):
    print("[Test] PING:", robot.ping())
    print("[Test] Servo 1: 90")
    robot.set_servo(1, 90)
    time.sleep(0.5)
    print("[Test] Servo 2: 90")
    robot.set_servo(2, 90)
    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.3)
    print("[Test] Linear motor close for 0.5 s")
    robot.motor_close()
    time.sleep(0.5)
    robot.motor_stop()
    print("[Test] Done")


def build_parser():
    parser = argparse.ArgumentParser(description="Gamepad 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")
    return parser


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

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

        if args.cmd:
            print(robot._send(args.cmd))
            time.sleep(3)   # 保持 3 秒，让电机持续运行
            return

        if args.test:
            run_self_test(robot)
            return

        controller = GamepadRemoteController(robot)
        controller.init_gamepad()
        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:
            controller.shutdown()
        if robot is not None:
            robot.close()
        print("[System] Shutdown complete")


if __name__ == "__main__":
    main()