#!/usr/bin/env python3
"""
EZAI-Go2 简洁使用示例 - 最终版本
基于go2_window_fix.py的成功实现
"""

import os
import sys
import ezai.go2 as go2

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

def simple_video_demo():
    """最简洁的视频流使用示例"""
    print("=== EZAI-Go2 简洁视频演示 ===")
    
    # 1. 创建Go2实例（自动检测接口）
    robot = go2.Go2()
    robot.init()
    print(f"使用接口: {robot.interface}")
    
    # 2. 获取视频读取器（自动选择最佳方法）
    read_frame, show_frame, cleanup = robot.get_video_reader()
    
    if read_frame is None:
        print("视频流初始化失败")
        return
    
    print("视频流启动成功，按 'q' 退出")
    
    try:
        frame_count = 0
        while True:
            # 读取一帧
            frame = read_frame()
            
            # 显示帧
            key = show_frame(frame)
            
            # 计数
            if frame is not None:
                frame_count += 1
                if frame_count % 30 == 0:
                    print(f"✅ 已显示 {frame_count} 帧")
            
            # 检查退出
            if key == ord('q'):
                break
                
    except KeyboardInterrupt:
        print("\n用户中断")
    finally:
        # 清理资源
        cleanup()
        print("程序结束")

def video_processing_demo():
    """视频处理示例 - 获取帧数据进行处理"""
    print("=== EZAI-Go2 视频处理演示 ===")
    
    # 1. 初始化
    robot = go2.Go2()
    robot.init()
    
    # 2. 获取视频读取器
    read_frame, show_frame, cleanup = robot.get_video_reader()
    
    if read_frame is None:
        print("视频流初始化失败")
        return
    
    print("视频流启动成功，按 'q' 退出，按 's' 保存当前帧")
    
    try:
        frame_count = 0
        while True:
            # 读取一帧
            frame = read_frame()
            
            if frame is not None:
                # 在这里可以进行图像处理
                # 例如：转换为灰度图
                # import cv2
                # gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                
                # 显示原始帧
                key = show_frame(frame)
                
                frame_count += 1
                if frame_count % 30 == 0:
                    print(f"✅ 处理第 {frame_count} 帧")
                
                # 保存功能
                if key == ord('s'):
                    import cv2
                    filename = f"/tmp/go2_frame_{frame_count}.jpg"
                    cv2.imwrite(filename, frame)
                    print(f"保存帧到: {filename}")
            else:
                # 没有帧数据时也要等待
                key = show_frame(None)
            
            # 检查退出
            if key == ord('q'):
                break
                
    except KeyboardInterrupt:
        print("\n用户中断")
    finally:
        cleanup()
        print("程序结束")

if __name__ == "__main__":
    print("选择演示模式:")
    print("1. 简单视频演示")
    print("2. 视频处理演示")
    
    try:
        choice = input("请选择 (1/2): ").strip()
        
        if choice == "1":
            simple_video_demo()
        elif choice == "2":
            video_processing_demo()
        else:
            print("默认运行简单视频演示...")
            simple_video_demo()
            
    except KeyboardInterrupt:
        print("\n程序被用户中断")