#!/usr/bin/env python3
"""
Go2简单视频读取示例
最简单的获取Go2画面的方法
"""

import os
import time
import cv2
import ezai.go2 as go2

def simple_video_display():
    """最简单的视频显示示例"""
    # 设置显示环境
    os.environ["DISPLAY"] = ":0"
    
    print("初始化Go2机器人...")
    robot = go2.Go2()
    print(f"网络接口: {robot.interface}")
    
    # 初始化连接
    robot.init()
    print("Go2连接成功")
    
    # 获取视频读取器（自动选择最佳方法）
    read_frame, cleanup = robot.get_video_reader(method='auto')
    
    if read_frame is None:
        print("无法创建视频读取器")
        return
    
    print("开始读取视频，按q键退出...")
    
    frame_count = 0
    try:
        while True:
            # 读取一帧
            frame = read_frame()
            
            if frame is not None:
                # 显示图像
                cv2.imshow('Go2视频流', frame)
                frame_count += 1
                
                # 每30帧显示一次计数
                if frame_count % 30 == 0:
                    print(f"已显示 {frame_count} 帧")
                
                # 检查退出
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            else:
                print("读取帧失败，等待重试...")
                time.sleep(0.1)
                
    except KeyboardInterrupt:
        print("程序被中断")
    finally:
        # 清理资源
        if cleanup:
            cleanup()
        cv2.destroyAllWindows()
        print(f"程序结束，总共显示 {frame_count} 帧")

def video_processing_example():
    """视频处理示例 - 添加简单的图像处理"""
    os.environ["DISPLAY"] = ":0"
    
    print("初始化Go2机器人...")
    robot = go2.Go2()
    robot.init()
    
    # 获取视频读取器
    read_frame, cleanup = robot.get_video_reader(method='pipe')  # 强制使用管道方法
    
    if read_frame is None:
        print("无法创建视频读取器")
        return
    
    print("开始视频处理，按q键退出...")
    
    try:
        while True:
            frame = read_frame()
            
            if frame is not None:
                # 简单的图像处理示例
                # 1. 转换为灰度图
                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                
                # 2. 边缘检测
                edges = cv2.Canny(gray, 50, 150)
                
                # 3. 转换回BGR用于显示
                edges_bgr = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
                
                # 4. 并排显示原图和边缘图
                combined = cv2.hconcat([frame, edges_bgr])
                
                cv2.imshow('Go2视频处理', combined)
                
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
            else:
                time.sleep(0.1)
                
    except KeyboardInterrupt:
        print("程序被中断")
    finally:
        if cleanup:
            cleanup()
        cv2.destroyAllWindows()
        print("程序结束")

def save_video_frames():
    """保存视频帧到文件示例"""
    print("初始化Go2机器人...")
    robot = go2.Go2()
    robot.init()
    
    # 获取视频读取器
    read_frame, cleanup = robot.get_video_reader(method='pipe')
    
    if read_frame is None:
        print("无法创建视频读取器")
        return
    
    print("开始保存视频帧，保存10帧...")
    
    try:
        for i in range(10):
            frame = read_frame()
            
            if frame is not None:
                filename = f"go2_frame_{i:03d}.jpg"
                cv2.imwrite(filename, frame)
                print(f"保存 {filename}")
                time.sleep(0.5)  # 每0.5秒保存一帧
            else:
                print(f"第{i+1}帧读取失败")
                i -= 1  # 重试
                
    except KeyboardInterrupt:
        print("程序被中断")
    finally:
        if cleanup:
            cleanup()
        print("保存完成")

def main():
    print("Go2简单视频读取示例")
    print("1. 简单视频显示")
    print("2. 视频处理示例")
    print("3. 保存视频帧")
    
    choice = input("请选择示例 (1/2/3): ").strip()
    
    if choice == "1":
        simple_video_display()
    elif choice == "2":
        video_processing_example()
    elif choice == "3":
        save_video_frames()
    else:
        print("运行简单视频显示...")
        simple_video_display()

if __name__ == "__main__":
    main()