import cv2
import time
import argparse
import numpy as np
from face_utils import FaceDetector, FaceRecognizer, draw_face_box
from face_db import FaceDatabase
from voice_prompt import VoicePrompt

def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='人脸识别应用')
    parser.add_argument('--camera', type=int, default=0, help='摄像头索引')
    parser.add_argument('--confidence', type=float, default=0.5, help='人脸检测置信度阈值')
    parser.add_argument('--similarity', type=float, default=0.6, help='人脸相似度阈值')
    parser.add_argument('--display_width', type=int, default=640, help='显示窗口宽度')
    parser.add_argument('--display_height', type=int, default=480, help='显示窗口高度')
    return parser.parse_args()

def main():
    """人脸识别应用主函数"""
    args = parse_args()
    
    # 初始化组件
    print("正在初始化人脸识别系统...")
    voice = VoicePrompt()
    detector = FaceDetector(conf_thres=args.confidence)
    recognizer = FaceRecognizer(detector=detector, similarity_threshold=args.similarity)
    db = FaceDatabase()
    
    # 加载已知人脸特征
    known_faces = db.get_all_features()
    if not known_faces:
        print("警告: 数据库中没有注册的人脸，请先运行注册程序")
        voice.speak("数据库中没有注册的人脸，请先运行注册程序", block=True)
    else:
        print(f"已加载 {len(known_faces)} 个人脸特征")
    
    # 打开摄像头
    cap = cv2.VideoCapture(args.camera)
    if not cap.isOpened():
        print("无法打开摄像头")
        return
    
    # 设置摄像头分辨率
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, args.display_width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, args.display_height)
    
    print("系统初始化完成，开始人脸识别...")
    voice.speak("人脸识别系统已启动", block=True)
    
    # 跟踪识别状态，避免重复语音提示
    recognized_persons = {}  # {person_id: last_recognition_time}
    recognition_cooldown = 5.0  # 同一个人识别成功后的冷却时间（秒）
    
    # 绘制UI相关
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.6
    line_thickness = 2
    fps_history = []
    
    try:
        while True:
            # 计时（用于计算FPS）
            start_time = time.time()
            
            # 读取摄像头帧
            ret, frame = cap.read()
            if not ret:
                print("无法获取摄像头画面")
                break
            
            # 镜像翻转以便更直观
            frame = cv2.flip(frame, 1)
            
            # 创建显示帧
            display_frame = frame.copy()
            
            # 检测人脸
            faces = detector.detect_faces(frame)
            
            # 处理每个检测到的人脸
            for face_box in faces:
                # 识别人脸
                result = recognizer.identify_face(frame, known_faces, face_box[:4])
                
                # 如果识别成功
                if result:
                    person_id, name, similarity = result
                    
                    # 绘制人脸框和标签（绿色表示识别成功）
                    draw_face_box(display_frame, face_box, name, similarity, color=(0, 255, 0))
                    
                    # 检查是否需要播放语音提示（避免频繁重复）
                    current_time = time.time()
                    if (person_id not in recognized_persons or 
                        current_time - recognized_persons[person_id] > recognition_cooldown):
                        
                        # 更新识别时间
                        recognized_persons[person_id] = current_time
                        
                        # 播放语音提示
                        voice.speak(f"您好，{name}", block=False)
                        print(f"识别成功: {name} (相似度: {similarity:.2f})")
                else:
                    # 未识别的人脸使用红色框
                    draw_face_box(display_frame, face_box, "未知", color=(0, 0, 255))
            
            # 计算并显示FPS
            end_time = time.time()
            fps = 1.0 / (end_time - start_time)
            fps_history.append(fps)
            if len(fps_history) > 30:
                fps_history.pop(0)
            avg_fps = sum(fps_history) / len(fps_history)
            
            cv2.putText(
                display_frame,
                f"FPS: {avg_fps:.1f}",
                (10, 30),
                font,
                font_scale,
                (0, 255, 255),
                line_thickness
            )
            
            # 显示数据库信息
            cv2.putText(
                display_frame,
                f"数据库: {len(known_faces)} 个人脸",
                (10, 60),
                font,
                font_scale,
                (0, 255, 255),
                line_thickness
            )
            
            # 显示操作指南
            cv2.putText(
                display_frame,
                "按ESC退出, 按R刷新数据库",
                (10, display_frame.shape[0] - 10),
                font,
                font_scale,
                (255, 255, 255),
                line_thickness
            )
            
            # 显示图像
            cv2.imshow("人脸识别系统", display_frame)
            
            # 按键处理
            key = cv2.waitKey(1) & 0xFF
            
            # 按ESC键退出
            if key == 27:
                break
            
            # 按R键刷新数据库
            elif key == ord('r'):
                print("正在刷新人脸数据库...")
                known_faces = db.get_all_features()
                print(f"已刷新: 加载了 {len(known_faces)} 个人脸特征")
                voice.speak("已刷新人脸数据库", block=False)
    
    except Exception as e:
        print(f"发生错误: {e}")
    
    finally:
        # 释放资源
        cap.release()
        cv2.destroyAllWindows()
        db.close()
        print("人脸识别系统已关闭")

if __name__ == "__main__":
    main()