#!/usr/bin/env python3
"""
完全按照pipstream.py模式的Go2视频流测试
确保与灵芯派上工作的pipstream.py完全一致
"""

import os
import time
import cv2
import numpy as np

# 设置显示环境（与pipstream.py一致）
os.environ["DISPLAY"] = ":0"

def test_exact_pipstream():
    """完全按照pipstream.py的实现测试"""
    print("=== 完全按照pipstream.py模式测试 ===")
    
    pipe_path = "/tmp/go2_video_pipe"
    interface = "enx00e0986113a6"  # 与pipstream.py一致
    width, height = 1280, 720
    
    # 清理旧管道（与pipstream.py一致）
    if os.path.exists(pipe_path):
        os.remove(pipe_path)
    
    # 创建管道（与pipstream.py一致）
    os.mkfifo(pipe_path)
    print("命名管道已创建")
    
    # 启动GStreamer（完全复制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 &"
    )
    
    print("启动GStreamer...")
    os.system(gst_cmd)
    
    # 等待启动（与pipstream.py一致）
    time.sleep(3)
    
    try:
        print("开始读取视频...")
        frame_size = width * height * 3
        
        with open(pipe_path, 'rb') as pipe:
            frame_count = 0
            success_count = 0
            
            while True:
                # 读取一帧（与pipstream.py一致）
                frame_data = pipe.read(frame_size)
                
                if len(frame_data) == frame_size:
                    # 转换为图像并显示（与pipstream.py一致）
                    frame = np.frombuffer(frame_data, dtype=np.uint8).reshape((height, width, 3))
                    cv2.imshow('Go2视频流 (pipstream模式)', frame)
                    frame_count += 1
                    success_count += 1
                    
                    if frame_count % 30 == 0:
                        print(f"✅ 显示第 {success_count} 帧")
                
                # 检查退出
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
                    
    except Exception as e:
        print(f"错误: {e}")
    finally:
        # 清理（与pipstream.py一致）
        os.system("pkill -f gst-launch")
        if os.path.exists(pipe_path):
            os.remove(pipe_path)
        cv2.destroyAllWindows()
        print("程序结束")

def test_ezai_wrapper():
    """测试EZAI封装版本"""
    print("\n=== 测试EZAI封装版本 ===")
    
    try:
        # 导入ezai
        import sys
        sys.path.insert(0, '..')
        import ezai.go2 as go2
        
        print("初始化Go2机器人...")
        robot = go2.Go2()
        robot.init()
        print(f"网络接口: {robot.interface}")
        
        # 强制使用管道方法
        print("使用管道方法获取视频读取器...")
        read_frame, cleanup = robot.get_video_reader(method='pipe')
        
        if read_frame is None:
            print("❌ 无法创建视频读取器")
            return
        
        print("✅ 视频读取器创建成功")
        print("开始读取视频，按q键退出...")
        
        frame_count = 0
        success_count = 0
        
        while True:
            frame = read_frame()
            
            if frame is not None:
                cv2.imshow('Go2视频流 (EZAI封装)', frame)
                frame_count += 1
                success_count += 1
                
                if frame_count % 30 == 0:
                    print(f"✅ EZAI显示第 {success_count} 帧")
                
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            else:
                frame_count += 1
                if frame_count % 30 == 0:
                    print(f"⏳ EZAI等待数据... ({success_count}/{frame_count})")
                time.sleep(0.01)
                
    except Exception as e:
        print(f"❌ EZAI测试失败: {e}")
    finally:
        if cleanup:
            cleanup()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    print("选择测试模式:")
    print("1. pipstream.py原始模式")
    print("2. EZAI封装模式")
    print("3. 先运行1，再运行2（对比）")
    
    try:
        choice = input("请选择 (1/2/3): ").strip()
        
        if choice == "1":
            test_exact_pipstream()
        elif choice == "2":
            test_ezai_wrapper()
        elif choice == "3":
            print("先运行pipstream模式...")
            test_exact_pipstream()
            input("按Enter继续测试EZAI封装...")
            test_ezai_wrapper()
        else:
            print("默认运行pipstream模式...")
            test_exact_pipstream()
            
    except KeyboardInterrupt:
        print("\n测试被用户中断")
    except Exception as e:
        print(f"测试异常: {e}")