import os
import cv2
import numpy as np
from keras_facenet import FaceNet
import gradio as gr

# Initialize FaceNet model
embedder = FaceNet()

def img_to_encoding(image, embedder):
    try:
        if image is None:
            print("Debug: Image is None")
            return None
        if not isinstance(image, np.ndarray):
            print(f"Debug: Image is not a numpy array, type(image)={type(image)}")
            return None

        print(f"Debug: Image shape = {image.shape}, dtype = {image.dtype}")

        # Convert BGR to RGB and resize to 160x160
        img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (160, 160))

        # Get the embedding
        embedding = embedder.embeddings([img])[0]
        return embedding
    except Exception as e:
        print(f"Error in img_to_encoding: {e}")
        return None

database = {}

def register_user(image, username):
    try:
        if not os.path.exists('images'):
            os.makedirs('images')
        
        if image is None or image.size == 0:
            print("Debug: Invalid image input")
            return "Error: Uploaded image is empty or invalid."

        print(f"Debug: Image received for registration, username={username}")

        # Save the image to disk
        img_path = f'images/{username}.jpg'
        success = cv2.imwrite(img_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
        
        if not success:
            print("Debug: Failed to save image to disk")
            return "Error: Failed to save the image."

        encoding = img_to_encoding(image, embedder)
        if encoding is None:
            return "Error: Failed to process the image."

        database[username] = encoding
        return f"User '{username}' registered successfully!"
    except Exception as e:
        print(f"Error during registration: {e}")
        return f"Error during registration: {e}"

def verify_user(image):
    try:
        if image is None or image.size == 0:
            print("Debug: Invalid image input for verification")
            return "Error: Uploaded image is empty or invalid."

        print("Debug: Image received for verification")

        encoding = img_to_encoding(image, embedder)
        if encoding is None:
            return "Error: Failed to process the image."

        min_dist = 1000
        identity = None

        for (name, db_enc) in database.items():
            dist = np.linalg.norm(encoding - db_enc)
            if dist < min_dist:
                min_dist = dist
                identity = name

        if min_dist > 5:
            return "Identity not recognized."
        else:
            return f"It's {identity}, the distance is {min_dist}."
    except Exception as e:
        print(f"Error during verification: {e}")
        return f"Error during verification: {e}"

# Gradio Interface for Registration
register_interface = gr.Interface(
    fn=register_user,
    inputs=[gr.Image(type="numpy", label="Take a Photo or Upload"), gr.Textbox(label="Username")],
    outputs="text",
    title="Register User",
    description="Take a photo or upload an image and enter a username to register."
)

# Gradio Interface for Verification
verify_interface = gr.Interface(
    fn=verify_user,
    inputs=gr.Image(type="numpy", label="Take a Photo or Upload"),
    outputs="text",
    title="Verify User",
    description="Take a photo or upload an image to verify the user's identity."
)

# Combine the interfaces into a single application
app = gr.Blocks()

with app:
    gr.Markdown("# Face Identification System")
    with gr.Tab("Register"):
        register_interface.render()
    with gr.Tab("Verify"):
        verify_interface.render()

# Run the Gradio application
if __name__ == "__main__":
    app.launch()