import gradio as gr
import os
import shutil
import threading
import time
from PIL import Image
import numpy as np
from utils import verify_identity

registration_dir = "registered_users"
os.makedirs(registration_dir, exist_ok=True)

# Register user
def register_user(name, image):
    if image is None or not name.strip():
        return "❌ Please provide a name and webcam photo."
    
    image = Image.fromarray(image.astype("uint8"))
    user_path = os.path.join(registration_dir, name)
    os.makedirs(user_path, exist_ok=True)
    image.save(os.path.join(user_path, f"{name}_1.jpg"))
    return f"✅ User '{name}' registered successfully."

# Webcam verify
def process_image(frame):
    if frame is None:
        return None, "No frame captured."
    result = verify_identity(None, frame, registration_dir)
    return frame, result

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Real-Time Face Verification (FaceNet)")

    with gr.Tab("👤 Register"):
        name = gr.Textbox(label="Enter Your Name")
        image = gr.Image(label="Capture Image")
        register_btn = gr.Button("Register")
        register_out = gr.Textbox(label="Status")
        register_btn.click(fn=register_user, inputs=[name, image], outputs=register_out)

    with gr.Tab("🔍 Live Verify"):
        with gr.Row():
            input_img = gr.Image(label="📸 Webcam", streaming=True, sources=["webcam"])
            output_img = gr.Image(label="🧠 Output Frame")
        result_text = gr.Textbox(label="Prediction", interactive=False)
        input_img.stream(fn=process_image, inputs=input_img, outputs=[output_img, result_text])

def auto_cleanup(interval=600):
    while True:
        if os.path.exists(registration_dir):
            shutil.rmtree(registration_dir)
            os.makedirs(registration_dir)
            print("🧹 Cleaned up registered_users/")
        time.sleep(interval)

threading.Thread(target=auto_cleanup, daemon=True).start()

if __name__ == "__main__":
    demo.launch()