import serial
import time


class RobotDriver:
    def __init__(self, port='COM6', baud=115200):
        try:
            self.ser = serial.Serial(port, baud, timeout=1)
            print(f"[System] Connected to {port}")
            time.sleep(2)  # 等STM32复位
        except Exception as e:
            print(f"[Error] Connection failed: {e}")
            exit()

    def _send(self, cmd):
        """发送原始指令"""
        full_cmd = f"{cmd}\n"
        self.ser.write(full_cmd.encode())
        print(f"-> Sent: {cmd}")
        time.sleep(0.05)  # 给单片机一点反应时间

    # === 直线电机控制 API ===
    def motor_open(self):
        self._send("M:OPEN")

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

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

    # === 舵机控制 API ===
    def set_servo(self, servo_id, angle):
        self._send(f"S{servo_id}:{angle}")

    # === 构型切换 (组合动作) ===
    def set_config(self, mode):
        """
        mode 0: 初始 (S1=90, S2=90)
        mode 1: 错位 (S1=30, S2=150)
        mode 2: 对握 (S1=120, S2=60)
        """
        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()
        self.ser.close()


# === 主程序逻辑 ===
if __name__ == "__main__":
    # 请修改端口号
    robot = RobotDriver(port='COM9')

    print("\n=== 全能控制面板 ===")
    print(" [1] 变构型: Mode 1 (30, 150)")
    print(" [2] 变构型: Mode 2 (120, 60)")
    print(" [0] 变构型: Reset (90, 90)")
    print(" [o] 直线电机: 张开")
    print(" [c] 直线电机: 闭合")
    print(" [s] 直线电机: 停止")
    print(" [q] 退出")

    try:
        while True:
            cmd = input("指令 > ").strip().lower()

            if cmd == 'q':
                break

            # 直线电机
            elif cmd == 'o':
                robot.motor_open()
            elif cmd == 'c':
                robot.motor_close()
            elif cmd == 's':
                robot.motor_stop()

            # 舵机构型
            elif cmd == '1':
                robot.set_config(1)
            elif cmd == '2':
                robot.set_config(2)
            elif cmd == '0':
                robot.set_config(0)

    except KeyboardInterrupt:
        pass
    finally:
        robot.close()
        print("System Shutdown.")