import gradio as gr
import cv2
import numpy as np
from face_recognition_system import FaceRecognitionSystem
import os
import tempfile

# Initialize the face recognition system
face_system = FaceRecognitionSystem()

def process_image(image, confidence_threshold=0.5, similarity_threshold=2.0):
    """Process a single image and return the annotated result"""
    # Convert from RGB (Gradio) to BGR (OpenCV)
    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Update thresholds
    face_system.confidence_threshold = confidence_threshold
    face_system.similarity_threshold = similarity_threshold
    
    # Process the frame
    processed_frame = face_system.process_frame(image_bgr)
    
    # Convert back to RGB for display
    return cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB)

def add_face(image, name):
    """Add a new face to the database"""
    if not name.strip():
        return "Error: Please enter a name"
    
    # Convert from RGB (Gradio) to BGR (OpenCV)
    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    if face_system.add_face_to_database(name, image_bgr):
        return f"Successfully added {name} to database"
    return "Failed to add face to database"

# Create the Gradio interface
with gr.Blocks(title="Face Recognition System") as demo:
    gr.Markdown("# Face Recognition System")
    gr.Markdown("Upload an image to detect and recognize faces, or add new faces to the database.")
    
    with gr.Tab("Recognize Faces"):
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Input Image", type="numpy")
                confidence_slider = gr.Slider(
                    minimum=0.1, 
                    maximum=1.0, 
                    value=0.5, 
                    step=0.1, 
                    label="Confidence Threshold"
                )
                similarity_slider = gr.Slider(
                    minimum=0.5, 
                    maximum=5.0, 
                    value=2.0, 
                    step=0.1, 
                    label="Similarity Threshold"
                )
                detect_btn = gr.Button("Detect Faces")
            
            with gr.Column():
                output_image = gr.Image(label="Output Image")
    
    with gr.Tab("Add New Face"):
        with gr.Row():
            with gr.Column():
                new_face_image = gr.Image(label="Face Image", type="numpy")
                name_input = gr.Textbox(label="Name")
                add_btn = gr.Button("Add Face to Database")
            
            with gr.Column():
                result_text = gr.Textbox(label="Result")
    
    # Set up event handlers
    detect_btn.click(
        fn=process_image,
        inputs=[input_image, confidence_slider, similarity_slider],
        outputs=output_image
    )
    
    add_btn.click(
        fn=add_face,
        inputs=[new_face_image, name_input],
        outputs=result_text
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()