import cv2
import numpy as np
import time
import os

# 导入我们的3D手部检测模块
from src.hand_detection_3d import load_mediapipe_model, process_frame_3d

# 尝试导入 Picamera2，如果不可用则使用标准OpenCV
try:
    from picamera2 import Picamera2
    PICAMERA_AVAILABLE = True
except ImportError:
    PICAMERA_AVAILABLE = False

def process_camera_3d(camera_id=0, output_path=None, use_picamera=True):
    """
    使用摄像头进行3D手部检测 - 简化版本
    
    参数:
    camera_id: 摄像头ID (默认为0)
    output_path: 输出视频文件路径 (可选)
    use_picamera: 是否优先使用Picamera2 (树莓派摄像头)
    """
    # 加载MediaPipe手部检测模型
    hands_model = load_mediapipe_model(max_num_hands=1)
    
    # 初始化摄像头
    cap = None
    picam2 = None
    
    if use_picamera and PICAMERA_AVAILABLE:
        # 使用 Picamera2 (树莓派摄像头)
        try:
            print("正在初始化 Picamera2...")
            picam2 = Picamera2()
            video_config = picam2.create_video_configuration(main={"size": (640, 480)})
            picam2.configure(video_config)
            picam2.start()
            print("✅ Picamera2 启动成功")
            frame_width, frame_height = 640, 480
        except Exception as e:
            print(f"❌ Picamera2 启动失败: {e}")
            print("⚠️  回退到 OpenCV VideoCapture")
            picam2 = None
            use_picamera = False
    
    if not use_picamera or not PICAMERA_AVAILABLE:
        # 使用标准 OpenCV VideoCapture
        cap = cv2.VideoCapture(camera_id)
        if not cap.isOpened():
            print(f"无法打开摄像头 ID: {camera_id}")
            return
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    fps = 30
    
    # 创建视频写入器（如果需要保存）
    video_writer = None
    if output_path:
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        video_writer = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
    
    # 前一帧的手部数据
    prev_hand_data = None
    
    # 初始化性能计数器
    frame_count = 0
    start_time = time.time()
    
    print("控制说明:")
    print("- 移动手部显示控制信号")
    print("- 拇指和食指捏合显示抓取状态") 
    print("- 按 Ctrl+C 退出")
    
    try:
        while True:
            # 读取一帧
            ret = False
            frame = None
            
            if picam2:
                try:
                    frame = picam2.capture_array()
                    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                    ret = True
                except Exception as e:
                    print(f"Picamera2 读取帧失败: {e}")
                    break
            else:
                ret, frame = cap.read()
            
            if not ret or frame is None:
                break
            
            frame_count += 1
            
            # 处理帧并获取三轴控制信号
            control_signal, prev_hand_data = process_frame_3d(frame, hands_model, prev_hand_data)
            
            # 显示控制信号
            if control_signal and frame_count % 30 == 0:  # 每秒显示一次
                print_control_signal(control_signal)
            
            # 写入输出视频
            if video_writer:
                video_writer.write(frame)
                
    except KeyboardInterrupt:
        print("\n收到中断信号，正在退出...")
    except Exception as e:
        print(f"程序运行时出错: {e}")
    finally:
        print("正在清理资源...")
        
        # 释放资源
        if picam2:
            try:
                picam2.stop()
                picam2.close()
            except:
                pass
        
        if cap:
            try:
                cap.release()
            except:
                pass
        
        if video_writer:
            try:
                video_writer.release()
                print(f"✅ 视频已保存: {output_path}")
            except:
                pass
        
        # 显示统计信息
        elapsed_time = time.time() - start_time
        if frame_count > 0:
            avg_fps = frame_count / elapsed_time
            print(f"📊 处理了 {frame_count} 帧，平均FPS: {avg_fps:.1f}")
        
        print("✅ 程序已退出")

def print_control_signal(control_signal):
    """显示手部控制信号信息"""
    print(f"手部控制信号:")
    print(f"  X轴角度: {control_signal['x_angle']:.1f}°")
    print(f"  Y轴角度: {control_signal['y_angle']:.1f}°") 
    print(f"  Z轴角度: {control_signal['z_angle']:.1f}°")
    print(f"  抓取状态: {'抓取' if control_signal['grip'] == 1 else '释放'}")
    print(f"  动作: {control_signal['action']}")
    print("-------------------------------")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='简化版3D手部检测')
    parser.add_argument('--camera-id', '-i', type=int, default=0, help='摄像头ID (默认为0)')
    parser.add_argument('--output', '-o', help='输出视频文件路径')
    
    args = parser.parse_args()
    
    process_camera_3d(
        camera_id=args.camera_id,
        output_path=args.output
    )