#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试控制命令协议
发送测试控制命令到STM32以验证协议正确性
"""

import serial
import struct
import time


def calculate_checksum(data: bytes) -> int:
    """计算16位累加和校验"""
    return sum(data) & 0xFFFF


def send_control_command(ser: serial.Serial, steering_pwm: int, throttle_pwm: int):
    """
    发送控制命令
    
    Args:
        ser: 串口对象
        steering_pwm: 转向PWM值 (1000-2000us)
        throttle_pwm: 油门PWM值 (1000-2000us)
    """
    # 协议常量
    FRAME_HEADER1 = 0xAA
    FRAME_HEADER2 = 0x55
    TYPE_CONTROL = 0x10
    FRAME_FOOTER1 = 0x0D
    FRAME_FOOTER2 = 0x0A
    
    # 限制PWM范围
    steering_pwm = max(1000, min(2000, steering_pwm))
    throttle_pwm = max(1000, min(2000, throttle_pwm))
    
    # 构建数据包
    data_len = 4  # 2字节转向 + 2字节油门
    payload = struct.pack('<HH', steering_pwm, throttle_pwm)
    
    # 构建完整帧用于计算校验和
    checksum_data = (
        bytes([TYPE_CONTROL]) +
        struct.pack('<H', data_len) +
        payload
    )
    
    # 计算校验和
    checksum = calculate_checksum(checksum_data)
    
    # 完整数据包
    packet = (
        bytes([FRAME_HEADER1, FRAME_HEADER2]) +
        checksum_data +
        struct.pack('<H', checksum) +
        bytes([FRAME_FOOTER1, FRAME_FOOTER2])
    )
    
    # 发送
    ser.write(packet)
    ser.flush()
    
    print(f"发送控制命令: 转向={steering_pwm}us, 油门={throttle_pwm}us, "
          f"校验和=0x{checksum:04X}, 包长度={len(packet)}字节")
    
    # 打印十六进制
    hex_str = ' '.join(f'{b:02X}' for b in packet)
    print(f"  HEX: {hex_str}")


def test_neutral():
    """测试中立位置 (1500us)"""
    print("\n=== 测试1: 中立位置 ===")
    PORT = "COM17"  # 根据实际情况修改
    BAUDRATE = 921600
    
    try:
        ser = serial.Serial(PORT, BAUDRATE, timeout=1)
        print(f"串口已连接: {PORT} @ {BAUDRATE} bps\n")
        
        # 发送中立位置命令 (多次确保接收)
        for i in range(5):
            send_control_command(ser, 1500, 1500)
            time.sleep(0.1)
        
        print("\n中立位置命令已发送 (应该保持静止)")
        
    except Exception as e:
        print(f"错误: {e}")
    finally:
        if 'ser' in locals() and ser.is_open:
            ser.close()
            print("\n串口已关闭")


def test_steering():
    """测试转向控制"""
    print("\n=== 测试2: 转向控制 ===")
    PORT = "COM17"
    BAUDRATE = 921600
    
    try:
        ser = serial.Serial(PORT, BAUDRATE, timeout=1)
        print(f"串口已连接: {PORT} @ {BAUDRATE} bps\n")
        
        # 中立位置
        print("1. 中立位置 (1500us)")
        for _ in range(10):
            send_control_command(ser, 1500, 1500)
            time.sleep(0.1)
        
        time.sleep(1)
        
        # 左转
        print("\n2. 左转 (1000us)")
        for _ in range(10):
            send_control_command(ser, 1000, 1500)
            time.sleep(0.1)
        
        time.sleep(1)
        
        # 中立
        print("\n3. 回中 (1500us)")
        for _ in range(10):
            send_control_command(ser, 1500, 1500)
            time.sleep(0.1)
        
        time.sleep(1)
        
        # 右转
        print("\n4. 右转 (2000us)")
        for _ in range(10):
            send_control_command(ser, 2000, 1500)
            time.sleep(0.1)
        
        time.sleep(1)
        
        # 回中
        print("\n5. 回中 (1500us)")
        for _ in range(10):
            send_control_command(ser, 1500, 1500)
            time.sleep(0.1)
        
        print("\n转向测试完成")
        
    except Exception as e:
        print(f"错误: {e}")
    finally:
        if 'ser' in locals() and ser.is_open:
            ser.close()
            print("\n串口已关闭")


def test_throttle():
    """测试油门控制"""
    print("\n=== 测试3: 油门控制 ===")
    PORT = "COM17"
    BAUDRATE = 921600
    
    try:
        ser = serial.Serial(PORT, BAUDRATE, timeout=1)
        print(f"串口已连接: {PORT} @ {BAUDRATE} bps\n")
        
        # 中立位置
        print("1. 中立位置 (1500us)")
        for _ in range(10):
            send_control_command(ser, 1500, 1500)
            time.sleep(0.1)
        
        time.sleep(1)
        
        # 前进 (注意: 需要根据实际电机方向调整)
        print("\n2. 前进 (1700us)")
        for _ in range(20):
            send_control_command(ser, 1500, 1700)
            time.sleep(0.1)
        
        time.sleep(1)
        
        # 回中
        print("\n3. 停止 (1500us)")
        for _ in range(10):
            send_control_command(ser, 1500, 1500)
            time.sleep(0.1)
        
        print("\n油门测试完成 (请根据实际电机反应调整PWM值)")
        
    except Exception as e:
        print(f"错误: {e}")
    finally:
        if 'ser' in locals() and ser.is_open:
            ser.close()
            print("\n串口已关闭")


def test_combined():
    """测试组合控制 (转向+油门)"""
    print("\n=== 测试4: 组合控制 ===")
    PORT = "COM17"
    BAUDRATE = 921600
    
    try:
        ser = serial.Serial(PORT, BAUDRATE, timeout=1)
        print(f"串口已连接: {PORT} @ {BAUDRATE} bps\n")
        
        # 中立
        print("1. 中立")
        for _ in range(10):
            send_control_command(ser, 1500, 1500)
            time.sleep(0.1)
        
        time.sleep(1)
        
        # 前进 + 左转
        print("\n2. 前进 + 左转 (1000us)")
        for _ in range(20):
            send_control_command(ser, 1000, 1600)
            time.sleep(0.1)
        
        time.sleep(1)
        
        # 回中
        print("\n3. 停止")
        for _ in range(10):
            send_control_command(ser, 1500, 1500)
            time.sleep(0.1)
        
        print("\n组合控制测试完成")
        
    except Exception as e:
        print(f"错误: {e}")
    finally:
        if 'ser' in locals() and ser.is_open:
            ser.close()
            print("\n串口已关闭")


def main():
    """主菜单"""
    print("=" * 50)
    print("控制命令协议测试工具")
    print("=" * 50)
    print("\n请选择测试项目:")
    print("1. 测试中立位置 (1500us)")
    print("2. 测试转向控制 (左-中-右)")
    print("3. 测试油门控制 (前进-停止)")
    print("4. 测试组合控制 (转向+油门)")
    print("0. 退出")
    
    choice = input("\n请输入选择 (0-4): ").strip()
    
    if choice == '1':
        test_neutral()
    elif choice == '2':
        test_steering()
    elif choice == '3':
        test_throttle()
    elif choice == '4':
        test_combined()
    elif choice == '0':
        print("退出")
    else:
        print("无效选择")


if __name__ == "__main__":
    main()