import pygame
import time
import sys

class FlightStickDetector:
    def __init__(self):
        # 初始化pygame
        pygame.init()
        pygame.joystick.init()
        
        self.joysticks = []
        self.running = True
        
    def detect_joysticks(self):
        """检测连接的手柄"""
        joystick_count = pygame.joystick.get_count()
        print(f"检测到 {joystick_count} 个手柄设备")
        
        if joystick_count == 0:
            print("未检测到任何手柄设备！")
            return False
            
        # 初始化所有检测到的手柄
        self.joysticks = []
        for i in range(joystick_count):
            joystick = pygame.joystick.Joystick(i)
            joystick.init()
            self.joysticks.append(joystick)
            
            print(f"\n手柄 {i}:")
            print(f"  名称: {joystick.get_name()}")
            print(f"  轴数量: {joystick.get_numaxes()}")
            print(f"  按钮数量: {joystick.get_numbuttons()}")
            print(f"  帽子开关数量: {joystick.get_numhats()}")
            
        return True
    
    def get_joystick_data(self, joystick_id=0):
        """获取指定手柄的数据"""
        if joystick_id >= len(self.joysticks):
            return None
            
        joystick = self.joysticks[joystick_id]
        
        # 获取轴数据
        axes = []
        for i in range(joystick.get_numaxes()):
            axis_value = joystick.get_axis(i)
            axes.append(round(axis_value, 3))
        
        # 获取按钮数据
        buttons = []
        for i in range(joystick.get_numbuttons()):
            button_pressed = joystick.get_button(i)
            buttons.append(button_pressed)
        
        # 获取帽子开关数据
        hats = []
        for i in range(joystick.get_numhats()):
            hat_value = joystick.get_hat(i)
            hats.append(hat_value)
        
        return {
            'name': joystick.get_name(),
            'axes': axes,
            'buttons': buttons,
            'hats': hats
        }
    
    def monitor_joystick(self, joystick_id=0, update_interval=0.1):
        """实时监控手柄数据"""
        print(f"\n开始监控手柄 {joystick_id} (按Ctrl+C停止)")
        print("=" * 50)
        
        try:
            while self.running:
                pygame.event.pump()  # 更新事件队列
                
                data = self.get_joystick_data(joystick_id)
                if data is None:
                    print("无法获取手柄数据")
                    break
                
                # 清屏并显示数据
                print("\033[H\033[J", end="")  # ANSI清屏
                print(f"手柄: {data['name']}")
                print("-" * 40)
                
                # 显示轴数据（通常用于飞行控制）
                axis_names = ['X轴(副翼)', 'Y轴(升降舵)', 'Z轴(方向舵)', 
                             '油门', '轴4', '轴5', '轴6', '轴7']
                
                print("轴数据:")
                for i, value in enumerate(data['axes']):
                    name = axis_names[i] if i < len(axis_names) else f'轴{i}'
                    # 创建进度条显示
                    bar_length = 20
                    bar_pos = int((value + 1) * bar_length / 2)
                    bar = ['─'] * bar_length
                    if 0 <= bar_pos < bar_length:
                        bar[bar_pos] = '█'
                    bar_str = ''.join(bar)
                    print(f"  {name:12}: {value:6.3f} [{bar_str}]")
                
                # 显示按钮数据
                pressed_buttons = [i for i, pressed in enumerate(data['buttons']) if pressed]
                if pressed_buttons:
                    print(f"\n按下的按钮: {pressed_buttons}")
                
                # 显示帽子开关数据
                if data['hats']:
                    print(f"帽子开关: {data['hats']}")
                
                time.sleep(update_interval)
                
        except KeyboardInterrupt:
            print("\n\n监控已停止")
        
    def cleanup(self):
        """清理资源"""
        pygame.joystick.quit()
        pygame.quit()

# 使用示例
def main():
    detector = FlightStickDetector()
    
    try:
        # 检测手柄
        if detector.detect_joysticks():
            # 如果检测到手柄，开始监控第一个手柄
            detector.monitor_joystick(0)
        else:
            print("请连接飞行手柄后重试")
            
    finally:
        detector.cleanup()

if __name__ == "__main__":
    main()