#!/usr/bin/env python3
"""
Go2视频流测试程序 - 使用tkinter显示
适用于灵芯派系统测试
"""

import os
import sys
import time
import cv2
import numpy as np
import tkinter as tk
from tkinter import messagebox
from PIL import Image, ImageTk
from threading import Thread
import queue

import ezai.go2 as go2

class VideoApp:
    def __init__(self, go2_robot):
        self.go2_robot = go2_robot
        self.root = tk.Tk()
        self.root.title("Go2视频流测试")
        self.root.geometry("800x600")
        
        # 创建视频显示标签
        self.video_label = tk.Label(self.root)
        self.video_label.pack(pady=10)
        
        # 状态标签
        self.status_label = tk.Label(self.root, text="状态: 初始化中...", font=("Arial", 12))
        self.status_label.pack(pady=5)
        
        # 退出按钮
        self.quit_button = tk.Button(self.root, text="退出", command=self.quit_app, font=("Arial", 12))
        self.quit_button.pack(pady=10)
        
        # 图像队列
        self.image_queue = queue.Queue()
        
        # 运行标志
        self.running = True
        
        # 启动视频线程
        self.video_thread = Thread(target=self.video_capture_thread)
        self.video_thread.daemon = True
        self.video_thread.start()
        
        # 启动UI更新
        self.update_video()
        
    def video_capture_thread(self):
        """视频捕获线程"""
        while self.running:
            try:
                # 读取图像
                image = self.go2_robot.read_image()
                
                if image is None:
                    self.status_label.config(text="状态: 无法读取图像")
                    time.sleep(0.1)
                    continue
                
                # 将图像放入队列
                self.image_queue.put(image)
                self.status_label.config(text=f"状态: 正常 ({image.shape})")
                
            except Exception as e:
                self.status_label.config(text=f"状态: 错误 - {str(e)}")
                time.sleep(0.1)
    
    def update_video(self):
        """更新视频显示"""
        try:
            # 从队列获取图像
            if not self.image_queue.empty():
                image = self.image_queue.get_nowait()
                
                # 调整图像大小以适应窗口
                height, width = image.shape[:2]
                max_width = 750
                max_height = 450
                
                if width > max_width or height > max_height:
                    ratio = min(max_width/width, max_height/height)
                    new_width = int(width * ratio)
                    new_height = int(height * ratio)
                    image = cv2.resize(image, (new_width, new_height))
                
                # 转换颜色空间 BGR -> RGB
                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                
                # 转换为PIL图像
                image_pil = Image.fromarray(image_rgb)
                image_tk = ImageTk.PhotoImage(image_pil)
                
                # 更新标签
                self.video_label.config(image=image_tk)
                self.video_label.image = image_tk  # 保持引用
            
        except queue.Empty:
            pass
        except Exception as e:
            print(f"更新视频错误: {e}")
        
        # 继续更新
        if self.running:
            self.root.after(30, self.update_video)
    
    def quit_app(self):
        """退出应用"""
        self.running = False
        self.root.quit()
        self.root.destroy()
    
    def run(self):
        """运行应用"""
        try:
            self.root.mainloop()
        except KeyboardInterrupt:
            self.quit_app()

def check_opencv_gstreamer():
    """检查OpenCV GStreamer支持"""
    print("检查OpenCV GStreamer支持...")
    try:
        build_info = cv2.getCvBuildInformation()
        gstreamer_support = 'GStreamer' in build_info
        print(f"GStreamer支持: {'是' if gstreamer_support else '否'}")
        if not gstreamer_support:
            print("警告: OpenCV未编译GStreamer支持，视频流可能无法工作")
        return gstreamer_support
    except Exception as e:
        print(f"检查GStreamer支持失败: {e}")
        return False

def test_direct_gstreamer(interface):
    """直接测试GStreamer管道"""
    print(f"\n直接测试GStreamer管道 (接口: {interface})...")
    
    gstreamer_str = (
        f"udpsrc address=230.1.1.1 port=1720 multicast-iface={interface} "
        "! application/x-rtp, media=video, encoding-name=H264 "
        "! rtph264depay ! h264parse "
        "! avdec_h264 ! videoconvert ! video/x-raw,width=640,height=480,format=BGR "
        "! appsink drop=1"
    )
    
    print(f"GStreamer管道: {gstreamer_str}")
    
    try:
        cap = cv2.VideoCapture(gstreamer_str, cv2.CAP_GSTREAMER)
        if cap.isOpened():
            print("GStreamer管道打开成功")
            
            # 尝试读取几帧
            for i in range(5):
                ret, frame = cap.read()
                if ret:
                    print(f"成功读取第{i+1}帧: {frame.shape}")
                else:
                    print(f"读取第{i+1}帧失败")
                    break
            
            cap.release()
            return True
        else:
            print("GStreamer管道打开失败")
            return False
            
    except Exception as e:
        print(f"GStreamer测试失败: {e}")
        return False

def main():
    # 设置显示环境（灵芯派需要）
    os.environ["DISPLAY"] = ":0"
    
    print("=== Go2视频流测试 (tkinter版本) ===")
    
    # 检查OpenCV GStreamer支持
    gstreamer_ok = check_opencv_gstreamer()
    
    print("\n初始化Go2机器人...")
    
    # 创建Go2实例
    go2_robot = go2.Go2()
    print(f"网络接口: {go2_robot.interface}")
    
    # 直接测试GStreamer
    if not test_direct_gstreamer(go2_robot.interface):
        print("\n错误: GStreamer管道测试失败")
        print("可能的原因:")
        print("1. 机器狗未开机或未连接网络")
        print("2. 网络接口名称不正确")
        print("3. OpenCV未编译GStreamer支持")
        print("4. 防火墙阻止UDP端口1720")
        return
    
    # 初始化连接
    try:
        go2_robot.init()
        print("Go2连接成功")
    except Exception as e:
        print(f"Go2连接失败: {e}")
        return
    
    # 检查视频流
    try:
        test_image = go2_robot.read_image()
        if test_image is None:
            print("错误: 无法读取视频流")
            return
        print(f"视频流测试成功: {test_image.shape}")
    except Exception as e:
        print(f"视频流测试失败: {e}")
        return
    
    # 启动tkinter应用
    try:
        app = VideoApp(go2_robot)
        print("\n视频流已启动，点击退出按钮关闭")
        app.run()
    except Exception as e:
        print(f"tkinter应用启动失败: {e}")
        import traceback
        traceback.print_exc()
    finally:
        print("程序结束")

if __name__ == "__main__":
    main()