#!/usr/bin/env python3
"""
EZAI-Go2 OpenCV方法测试 - 使用原来的稳定方法
"""

import os
import sys
import time
import ezai.go2 as go2

os.environ["DISPLAY"] = ":0"

def opencv_video_demo():
    """使用OpenCV方法的视频演示（原来的稳定方法）"""
    print("=== EZAI-Go2 OpenCV方法演示 ===")
    
    # 1. 创建Go2实例
    robot = go2.Go2()
    robot.init()
    print(f"使用接口: {robot.interface}")
    
    # 2. 直接使用原来的OpenCV方法
    print("尝试使用OpenCV方法...")
    cap = robot.open_video()  # 使用原来的open_video方法
    
    if cap is None:
        print("OpenCV方法失败，尝试管道方法...")
        read_frame, show_frame, cleanup = robot.get_video_reader(method='pipe')
        if read_frame is None:
            print("所有方法都失败了")
            return
        use_pipe_method = True
    else:
        print("OpenCV方法成功！")
        use_pipe_method = False
    
    print("视频流启动成功，按 'q' 退出")
    
    try:
        frame_count = 0
        
        while True:
            if use_pipe_method:
                # 使用管道方法
                frame = read_frame()
                key = show_frame(frame)
            else:
                # 使用原来的OpenCV方法
                ret, frame = cap.read()
                if ret:
                    cv2.imshow('Go2视频流 (OpenCV方法)', frame)
                    frame_count += 1
                    
                    if frame_count % 30 == 0:
                        print(f"✅ 显示第 {frame_count} 帧")
                
                key = cv2.waitKey(1) & 0xFF
            
            # 检查退出
            if key == ord('q'):
                break
                
    except KeyboardInterrupt:
        print("\n用户中断")
    finally:
        # 清理资源
        if use_pipe_method:
            cleanup()
        else:
            cap.release()
            cv2.destroyAllWindows()
        print("程序结束")

def compare_methods():
    """对比不同方法的稳定性"""
    print("=== 方法对比测试 ===")
    
    robot = go2.Go2()
    robot.init()
    print(f"使用接口: {robot.interface}")
    
    methods = [
        ("OpenCV方法", lambda: robot.open_video()),
        ("管道方法", lambda: robot.get_video_reader(method='pipe'))
    ]
    
    for method_name, method_func in methods:
        print(f"\n--- 测试 {method_name} ---")
        
        try:
            if method_name == "OpenCV方法":
                cap = method_func()
                if cap is None:
                    print(f"{method_name} 初始化失败")
                    continue
                
                print(f"{method_name} 初始化成功")
                
                frame_count = 0
                start_time = time.time()
                
                # 测试30秒
                while time.time() - start_time < 30:
                    ret, frame = cap.read()
                    if ret:
                        cv2.imshow(f'Go2视频流 ({method_name})', frame)
                        frame_count += 1
                    
                    key = cv2.waitKey(1) & 0xFF
                    if key == ord('q'):
                        break
                
                cap.release()
                
            else:  # 管道方法
                read_frame, show_frame, cleanup = method_func()
                if read_frame is None:
                    print(f"{method_name} 初始化失败")
                    continue
                
                print(f"{method_name} 初始化成功")
                
                frame_count = 0
                start_time = time.time()
                
                # 测试30秒
                while time.time() - start_time < 30:
                    frame = read_frame()
                    key = show_frame(frame)
                    
                    if frame is not None:
                        frame_count += 1
                    
                    if key == ord('q'):
                        break
                
                cleanup()
            
            # 统计结果
            elapsed = time.time() - start_time
            fps = frame_count / elapsed if elapsed > 0 else 0
            print(f"{method_name} 结果: {frame_count} 帧, 平均FPS: {fps:.1f}")
            
        except Exception as e:
            print(f"{method_name} 测试出错: {e}")
        
        cv2.destroyAllWindows()
        
        # 方法间间隔
        if method_name != methods[-1][0]:
            print("等待3秒后测试下一个方法...")
            time.sleep(3)

if __name__ == "__main__":
    print("选择测试模式:")
    print("1. OpenCV方法演示")
    print("2. 方法对比测试")
    
    try:
        choice = input("请选择 (1/2): ").strip()
        
        if choice == "1":
            opencv_video_demo()
        elif choice == "2":
            compare_methods()
        else:
            print("默认运行OpenCV方法演示...")
            opencv_video_demo()
            
    except KeyboardInterrupt:
        print("\n程序被用户中断")