import face_recognition
import numpy as np

class FaceLibrary:
    def __init__(self, face_db):
        """
        face_db: 字典 {"Name": {"image-path": "x.jpg", "age": 20}, ...}
        """
        self.known_encodings = []
        self.known_names = []
        self.known_infos = []
        self._load_database(face_db)

    def _load_database(self, face_db):
        print("正在加载人脸底库...")
        for name, info in face_db.items():
            try:
                image = face_recognition.load_image_file(info["image-path"])
                encodings = face_recognition.face_encodings(image)
                if encodings:
                    self.known_encodings.append(encodings[0])
                    self.known_names.append(name)
                    self.known_infos.append(info)
                    print(f"✅ 已加载: {name}")
                else:
                    print(f"⚠️ 无法提取特征: {name}")
            except Exception as e:
                print(f"❌ 加载失败 {name}: {e}")

    def identify(self, frame_rgb, face_location=None, tolerance=0.5):
        """
        frame_rgb: RGB 图片
        face_location: (top, right, bottom, left) 或者 None (全图搜索)
        """
        if not self.known_encodings:
            return None

        locations = [face_location] if face_location else None
        
        try: 
            encodings = face_recognition.face_encodings(frame_rgb, known_face_locations=locations)
            
            if not encodings:
                return None
            
            unknown_encoding = encodings[0]
            distances = face_recognition.face_distance(self.known_encodings, unknown_encoding)
            min_idx = np.argmin(distances)
            
            if distances[min_idx] <= tolerance:
                return {
                    "name": self.known_names[min_idx],
                    "info": self.known_infos[min_idx],
                    "distance": distances[min_idx]
                }
        except Exception as e:
            print(f"识别出错: {e}")
            
        return None