"""
EZAI-Go2 简单测试示例
用于调试Go2连接和基本功能
"""

import time
import ezai.go2 as go2

def test_go2_connection():
    """测试Go2连接"""
    print("=== EZAI-Go2 连接测试 ===")
    
    try:
        # 创建Go2实例
        print("1. 创建Go2实例...")
        robot = go2.Go2()
        
        # 检查网络接口
        print("2. 检查网络接口...")
        if robot.interface:
            print(f"   找到网络接口: {robot.interface}")
        else:
            print("   未找到Go2网络接口")
            return False
        
        # 初始化连接
        print("3. 初始化Go2连接...")
        robot.init()
        print("   Go2初始化成功")
        
        # 测试基本运动
        print("4. 测试基本运动...")
        robot.Move(0.2, 0, 0)  # 缓慢前进
        time.sleep(1)
        robot.Move(0, 0, 0)    # 停止
        print("   运动测试完成")
        
        # 测试视频流
        print("5. 测试视频流...")
        if robot.cap is not None:
            for i in range(5):
                image = robot.read_image()
                if image is not None:
                    print(f"   第{i+1}帧: 成功获取图像 {image.shape}")
                else:
                    print(f"   第{i+1}帧: 获取失败")
                time.sleep(0.5)
        else:
            print("   视频流未初始化，跳过测试")
        
        # 断开连接
        print("6. 断开连接...")
        robot.disconnect()
        print("   连接已断开")
        
        print("=== 测试完成 ===")
        return True
        
    except Exception as e:
        print(f"测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_video_stream_only():
    """仅测试视频流"""
    print("=== 视频流专项测试 ===")
    
    try:
        robot = go2.Go2()
        robot.init()
        
        print("开始获取视频流（10帧）...")
        success_count = 0
        
        if robot.cap is not None:
            for i in range(10):
                image = robot.read_image()
                if image is not None:
                    success_count += 1
                    print(f"帧 {i+1}: 成功 {image.shape}")
                else:
                    print(f"帧 {i+1}: 失败")
                time.sleep(0.2)
        else:
            print("视频流未初始化，无法测试")
            return False
        
        print(f"视频流测试完成: {success_count}/10 成功")
        robot.disconnect()
        
        return success_count > 0
        
    except Exception as e:
        print(f"视频流测试失败: {e}")
        return False

if __name__ == "__main__":
    print("选择测试模式:")
    print("1. 完整连接测试")
    print("2. 仅视频流测试")
    
    try:
        choice = input("请选择 (1/2): ").strip()
        
        if choice == "1":
            test_go2_connection()
        elif choice == "2":
            test_video_stream_only()
        else:
            print("运行完整连接测试...")
            test_go2_connection()
            
    except KeyboardInterrupt:
        print("\n测试被用户中断")
    except Exception as e:
        print(f"测试异常: {e}")