import cv2
import time
import argparse
import os
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('--name', type=str, help='要注册的人员姓名')
    parser.add_argument('--camera', type=int, default=0, help='摄像头索引')
    parser.add_argument('--threshold', type=float, default=0.5, help='人脸检测置信度阈值')
    return parser.parse_args()

def register_face():
    """人脸注册主函数"""
    args = parse_args()
    
    # 初始化组件
    voice = VoicePrompt()
    detector = FaceDetector(conf_thres=args.threshold)
    recognizer = FaceRecognizer(detector=detector)
    db = FaceDatabase()
    
    # 打开摄像头
    cap = cv2.VideoCapture(args.camera)
    if not cap.isOpened():
        print("无法打开摄像头")
        return
    
    # 获取姓名（如果命令行没有提供，则交互输入）
    name = args.name
    if name is None:
        name = input("请输入姓名: ")
    
    # 创建新的人员记录
    person_id = db.add_person(name)
    print(f"已创建人员记录: {name} (ID: {person_id})")
    voice.speak(f"开始为{name}注册人脸信息，请正视摄像头", block=True)
    
    # 采集计数器和状态
    face_count = 0
    target_count = 5  # 需要采集的人脸图像数量
    last_capture_time = 0
    capture_interval = 1.0  # 每次采集间隔秒数
    
    print(f"请看向摄像头，将采集{target_count}张人脸图像...")
    
    while face_count < target_count:
        # 读取摄像头帧
        ret, frame = cap.read()
        if not ret:
            print("无法获取摄像头画面")
            break
        
        # 镜像翻转以便更直观
        frame = cv2.flip(frame, 1)
        
        # 显示帧
        display_frame = frame.copy()
        cv2.putText(
            display_frame,
            f"请直视摄像头 ({face_count}/{target_count})",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (0, 255, 0),
            2
        )
        
        # 检测人脸
        faces = detector.detect_faces(frame)
        
        # 如果检测到人脸
        current_time = time.time()
        if faces and (current_time - last_capture_time) >= capture_interval:
            # 只处理最大的一个人脸
            faces.sort(key=lambda x: (x[2]-x[0])*(x[3]-x[1]), reverse=True)
            face_box = faces[0]
            
            # 绘制人脸框
            draw_face_box(display_frame, face_box)
            
            # 提取人脸特征
            face_feature = recognizer.extract_face_features(frame, face_box[:4])
            
            if face_feature is not None:
                # 保存人脸特征到数据库
                feature_id = db.add_face_feature(person_id, face_feature)
                
                face_count += 1
                last_capture_time = current_time
                print(f"已采集第 {face_count}/{target_count} 张人脸")
                voice.speak(f"已完成第{face_count}次采集", block=False)
        
        # 显示图像
        cv2.imshow("人脸注册", display_frame)
        
        # 按ESC键退出
        key = cv2.waitKey(1) & 0xFF
        if key == 27:
            break
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
    db.close()
    
    if face_count >= target_count:
        print(f"成功完成{name}的人脸注册!")
        voice.speak(f"已完成{name}的人脸信息注册", block=True)
    else:
        print("人脸注册未完成")

if __name__ == "__main__":
    register_face()