#!/usr/bin/env python3
"""
EZAI-Go2 稳定版本测试 - 完全基于pipstream.py
"""

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

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

def stable_video_demo():
    """稳定的视频流演示（基于pipstream.py）"""
    print("=== EZAI-Go2 稳定视频演示 ===")
    
    # 1. 创建Go2实例
    robot = go2.Go2()
    robot.init()
    print(f"使用接口: {robot.interface}")
    
    # 2. 获取稳定的视频读取器（基于pipstream.py）
    read_frame, show_frame, cleanup = robot.get_video_reader(method='pipe')
    
    if read_frame is None:
        print("视频流初始化失败")
        return
    
    print("稳定视频流启动成功（基于pipstream.py方案）")
    print("按 'q' 退出")
    
    try:
        frame_count = 0
        
        while True:
            # 读取一帧（与pipstream.py一致）
            frame = read_frame()
            
            # 显示帧（与pipstream.py一致）
            key = show_frame(frame)
            
            # 计数（与pipstream.py一致）
            if frame is not None:
                frame_count += 1
                
                if frame_count % 30 == 0:
                    print(f"显示第 {frame_count} 帧")
            
            # 检查退出（与pipstream.py一致）
            if key == ord('q'):
                break
                
    except KeyboardInterrupt:
        print("\n用户中断")
    finally:
        # 清理资源
        cleanup()
        print("程序结束")

def compare_with_pipstream():
    """与pipstream.py对比测试"""
    print("=== 与pipstream.py对比测试 ===")
    
    print("1. 测试封装版本（动态接口）...")
    
    # 测试封装版本
    robot = go2.Go2()
    robot.init()
    print(f"   使用接口: {robot.interface}")
    
    read_frame, show_frame, cleanup = robot.get_video_reader(method='pipe')
    
    if read_frame is not None:
        try:
            frame_count = 0
            start_time = time.time()
            
            # 测试60秒
            while time.time() - start_time < 60:
                frame = read_frame()
                key = show_frame(frame)
                
                if frame is not None:
                    frame_count += 1
                
                if key == ord('q'):
                    break
            
            elapsed = time.time() - start_time
            fps = frame_count / elapsed if elapsed > 0 else 0
            print(f"   封装版本结果: {frame_count} 帧, 平均FPS: {fps:.1f}")
            
        finally:
            cleanup()
    
    print("\n等待3秒后测试原始pipstream.py...")
    time.sleep(3)
    
    # 测试原始pipstream.py
    print("2. 测试原始pipstream.py（固定接口）...")
    
    pipe_path = "/tmp/compare_pipstream"
    width, height = 1280, 720
    
    # 清理旧管道
    if os.path.exists(pipe_path):
        os.remove(pipe_path)
    
    # 创建管道
    os.mkfifo(pipe_path)
    
    # 使用固定接口（原始pipstream.py）
    gst_cmd = (
        "exec gst-launch-1.0 udpsrc address=230.1.1.1 port=1720 multicast-iface=enx00e0986113a6 "
        "! application/x-rtp, media=video, encoding-name=H264 ! rtph264depay ! h264parse "
        "! avdec_h264 ! videoconvert ! video/x-raw,format=BGR,width=1280,height=720 "
        "! filesink location=" + pipe_path + " > /dev/null 2>&1 &"
    )
    
    os.system(gst_cmd)
    time.sleep(3)
    
    try:
        frame_count = 0
        start_time = time.time()
        
        with open(pipe_path, 'rb') as pipe:
            frame_size = width * height * 3
            
            # 测试60秒
            while time.time() - start_time < 60:
                frame_data = pipe.read(frame_size)
                
                if len(frame_data) == frame_size:
                    frame = np.frombuffer(frame_data, dtype=np.uint8).reshape((height, width, 3))
                    cv2.imshow('原始pipstream.py', frame)
                    frame_count += 1
                
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    break
        
        elapsed = time.time() - start_time
        fps = frame_count / elapsed if elapsed > 0 else 0
        print(f"   原始版本结果: {frame_count} 帧, 平均FPS: {fps:.1f}")
        
    except Exception as e:
        print(f"   原始版本错误: {e}")
    finally:
        os.system("pkill -f gst-launch")
        if os.path.exists(pipe_path):
            os.remove(pipe_path)
        cv2.destroyAllWindows()

if __name__ == "__main__":
    print("选择测试模式:")
    print("1. 稳定视频演示")
    print("2. 与pipstream.py对比测试")
    
    try:
        choice = input("请选择 (1/2): ").strip()
        
        if choice == "1":
            stable_video_demo()
        elif choice == "2":
            compare_with_pipstream()
        else:
            print("默认运行稳定视频演示...")
            stable_video_demo()
            
    except KeyboardInterrupt:
        print("\n程序被用户中断")