import gradio as gr
import numpy as np
from PIL import Image
import cv2
import os
import pickle  # For loading the face embeddings database

# Import your face detection and recognition functions
from face_detection import detect_faces
from face_recognition import recognize_face  # Assuming you have this

# --- Configuration ---
DATABASE_PATH = "face_database.pkl" # Path to your stored face embeddings
THRESHOLD = 0.6  # Similarity threshold for recognition

# --- Load the face embeddings database (if it exists) ---
face_embeddings_db = {}
if os.path.exists(DATABASE_PATH):
    with open(DATABASE_PATH, 'rb') as f:
        face_embeddings_db = pickle.load(f)

# --- Function to process an input image ---
def recognize_faces_in_image(image):
    img_array = np.array(image)
    detected_faces = detect_faces(img_array) # Returns a list of (box, confidence)

    output_image = img_array.copy()
    results = []

    for box, _ in detected_faces:
        x1, y1, w, h = box
        x2, y2 = x1 + w, y1 + h

        face_roi = img_array[y1:y2, x1:x2]

        if face_roi.shape[0] > 0 and face_roi.shape[1] > 0:
            identity = recognize_face(face_roi, face_embeddings_db, THRESHOLD)

            if identity:
                color = (0, 255, 0)  # Green for recognized
                text = identity
            else:
                color = (255, 0, 0)  # Red for unknown
                text = "Unknown"

            cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
            cv2.putText(output_image, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
            results.append(f"Face detected at ({x1}, {y1}) - {text}")
        else:
            results.append(f"Small or invalid face detected.")

    return [Image.fromarray(output_image), "\n".join(results)]

# --- Gradio Interface ---
iface = gr.Interface(
    fn=recognize_faces_in_image,
    inputs=gr.Image(label="Upload an Image"),
    outputs=[gr.Image(label="Detected Faces"), gr.Textbox(label="Recognition Results")],
    title="Face Recognition App",
    description="Upload an image and see the detected and recognized faces."
)

iface.launch()