#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
interactive_servo_test.py — 基于 Python 的 ESP32 舵机控制交互脚本
依赖：pip install pyserial
用法：
  python3 interactive_servo_test.py /dev/ttyUSB0 [baudrate]
功能：
  1. 控制单个舵机
  2. 控制多个舵机
  3. 退出
  4. 自动增减模式：从 0 到 180° 每步 +6° → -6° 循环，50Hz 频率发送
  5. 设置自动模式舵机数量（通道 1~N）
"""
import sys
import time
import serial

FRAME_HEADER = 0xAA  # 帧头


def wakeup_esp32(port, baud=115200):
    """通过 RTS/DTR 线复位并唤醒 ESP32"""
    try:
        ser = serial.Serial(port, baud, timeout=1)
        ser.setDTR(False)
        ser.setRTS(False)
        time.sleep(0.01)
        ser.setDTR(True)
        ser.setRTS(True)
        ser.write(b"\r\n")
        time.sleep(1)
    except Exception as e:
        print(f"ESP32 唤醒失败: {e}")
    finally:
        try:
            ser.close()
        except:
            pass


def open_serial(port, baud):
    """打开并配置串口"""
    try:
        ser = serial.Serial(port, baud, timeout=1)
        print(f"已打开串口: {port} @ {baud}bps")
        return ser
    except Exception as e:
        print(f"打开串口失败: {e}")
        sys.exit(1)


def calc_checksum(data: bytearray) -> int:
    """计算 XOR 校验和"""
    cs = 0
    for b in data:
        cs ^= b
    return cs


def pack_commands(cmds):
    """
    cmds: list of (addr, channel, angle, duration_ms)
    返回：完整帧字节数组
    格式： [FRAME_HEADER][count] [addr, ch, ang, durL, durH]... [checksum]
    """
    pkt = bytearray([FRAME_HEADER, len(cmds)])
    for addr, ch, ang, dur in cmds:
        pkt.extend([addr, ch, ang, dur & 0xFF, (dur >> 8) & 0xFF])
    pkt.append(calc_checksum(pkt))
    return pkt


def print_packet(pkt):
    """打印十六进制数据包"""
    print("发送数据包:", ' '.join(f"0x{b:02X}" for b in pkt))


def input_hex(prompt):
    """输入 0x 开头或十进制整型"""
    val = input(prompt).strip()
    try:
        return int(val, 0)
    except ValueError:
        print(f"无效数字: {val}")
        return input_hex(prompt)


def input_cmd_single():
    """交互输入单条命令"""
    addr = input_hex("输入 addr (hex or dec): ")
    ch   = input_hex("输入 channel (hex or dec): ")
    ang  = input_hex("输入 angle (0-180): ")
    dur  = input_hex("输入 duration(ms): ")
    return [(addr, ch, ang, dur)]


def input_cmd_multiple():
    """交互输入多条命令"""
    cnt = input_hex("输入命令数量: ")
    cmds = []
    for i in range(cnt):
        print(f"第 {i+1} 条:")
        cmds.extend(input_cmd_single())
    return cmds


def auto_mode(ser, channels):
    """自动增减模式：0→180→0，步长6，50Hz，多舵机"""
    addr = input_hex("自动模式: 输入 addr (hex or dec): ")
    interval = 1.0 / 50.0
    dur_ms   = int(interval * 1000)
    angle    = 0
    direction = 1
    step      = 6

    print(f"启动自动模式 @50Hz, addr=0x{addr:02X}, channels={channels}, step={step}°")
    try:
        while True:
            cmds = [(addr, ch, angle, dur_ms) for ch in channels]
            pkt  = pack_commands(cmds)
            print_packet(pkt)
            ser.write(pkt)
            # 更新角度
            angle += direction * step
            if angle >= 180:
                angle = 180
                direction = -1
            elif angle <= 0:
                angle = 0
                direction = 1
            time.sleep(interval)
    except KeyboardInterrupt:
        print("\n自动模式已停止，返回菜单")


def set_channels():
    """自定义舵机数量，返回通道列表 1~N"""
    cnt = input_hex("输入自动模式舵机数量 N: ")
    if cnt < 1:
        print("数量必须 >= 1")
        return None
    channels = list(range(1, cnt+1))
    print(f"已设置通道列表: {channels}")
    return channels


def main():
    if len(sys.argv) < 2:
        print(f"用法: {sys.argv[0]} /dev/ttyUSB0 [baudrate]")
        sys.exit(1)
    port = sys.argv[1]
    baud = int(sys.argv[2]) if len(sys.argv) > 2 else 115200

    print("唤醒 ESP32...")
    wakeup_esp32(port, baud)
    ser = open_serial(port, baud)
    channels = list(range(1, 10))  # 默认通道 1-9

    try:
        while True:
            print("\n=== 操作菜单 ===")
            print("1. 控制单个舵机")
            print("2. 控制多个舵机")
            print("3. 退出")
            print("4. 自动增减模式(50Hz)")
            print("5. 设置自动模式舵机数量")
            opt = input("选择: ").strip()
            if opt == '1':
                cmds = input_cmd_single()
            elif opt == '2':
                cmds = input_cmd_multiple()
            elif opt == '4':
                auto_mode(ser, channels)
                continue
            elif opt == '5':
                new_ch = set_channels()
                if new_ch:
                    channels = new_ch
                continue
            elif opt == '3':
                break
            else:
                print("无效选项，重试")
                continue

            pkt = pack_commands(cmds)
            print_packet(pkt)
            ser.write(pkt)
            print("✅ 指令已发送")
    except KeyboardInterrupt:
        print("\n用户中断，退出")
    finally:
        ser.close()
        print("串口已关闭")

if __name__ == '__main__':
    main()