import gradio as gr
import numpy as np
from PIL import Image
import os
import json

# Try to import InsightFace
INSIGHTFACE_AVAILABLE = False
try:
    from insightface.app.face_analysis import FaceAnalysis
    INSIGHTFACE_AVAILABLE = True
    print("✓ InsightFace available")
except:
    print("InsightFace not available, using demo mode")

# Global variables
face_app = None
face_database = {}

def setup_models():
    global face_app
    if INSIGHTFACE_AVAILABLE:
        try:
            print("Loading InsightFace models...")
            face_app = FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
            face_app.prepare(ctx_id=0, det_thresh=0.5, det_size=(640, 640))
            print("✓ InsightFace models loaded")
        except Exception as e:
            print(f"Failed to load InsightFace: {e}")

def load_database():
    global face_database
    try:
        if os.path.exists('faces.json'):
            with open('faces.json', 'r') as f:
                face_database = json.load(f)
            print(f"Loaded {len(face_database)} faces from database")
    except:
        face_database = {}

def save_database():
    try:
        with open('faces.json', 'w') as f:
            json.dump(face_database, f)
    except:
        pass

def get_embedding(image):
    global face_app
    if not face_app or not image:
        if image is not None:
            seed = int(np.array(image).mean() * 1000) % 1000
            np.random.seed(seed)
            emb = np.random.rand(512)
            return emb / np.linalg.norm(emb), "Demo embedding"
        return None, "No image"
    
    try:
        img_array = np.array(image.convert('RGB'))
        faces = face_app.get(img_array)
        if not faces:
            return None, "No face detected"
        face = faces[0]
        return face.embedding, f"Face detected (confidence: {face.det_score:.2f})"
    except Exception as e:
        return None, f"Error: {str(e)}"

def add_face(image, name):
    if not name or not name.strip():
        return "Please enter a name"
    
    name = name.strip()
    embedding, msg = get_embedding(image)
    
    if embedding is None:
        return f"Failed: {msg}"
    
    face_database[name] = embedding.tolist()
    save_database()
    
    return f"✓ Added {name} ({msg}). Database now has {len(face_database)} faces."

def match_face(image):
    if not face_database:
        return "Database is empty. Please add faces first."
    
    if not image:
        return "Please upload an image"
    
    embedding, msg = get_embedding(image)
    if embedding is None:
        return f"Failed: {msg}"
    
    best_match = None
    best_score = -1
    
    for name, stored_emb in face_database.items():
        stored_emb = np.array(stored_emb)
        score = np.dot(embedding, stored_emb) / (np.linalg.norm(embedding) * np.linalg.norm(stored_emb))
        if score > best_score:
            best_score = score
            best_match = name
    
    if best_score > 0.6:
        return f"✓ Match Found: {best_match} (confidence: {best_score*100:.1f}%)"
    else:
        return f"❌ No match found. Best score: {best_score*100:.1f}% (threshold: 60%)"

def get_status():
    status = "✓ InsightFace loaded" if face_app else "Demo mode"
    db_count = len(face_database)
    return f"System: {status} | Database: {db_count} faces"

def clear_database():
    global face_database
    face_database = {}
    save_database()
    return "Database cleared successfully"

# Initialize
print("Starting FaceMatch system...")
setup_models()
load_database()

css = """
#col-left {
    margin: 0 auto;
    max-width: 450px;
}
#col-right {
    margin: 0 auto;
    max-width: 450px;
}
"""

with gr.Blocks(css=css) as demo:
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>🎯 FaceMatch Pro</h1>
        <p>Professional Face Recognition System</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(elem_id="col-left"):
            gr.HTML("""
            <div style="text-align: center; font-size: 18px; margin-bottom: 10px;">
                <b>📥 Add Face to Database</b>
            </div>
            """)
            add_image = gr.Image(label="Upload Photo", sources='upload', type="pil")
            add_name = gr.Textbox(label="Person Name", placeholder="Enter name...")
            add_button = gr.Button("Add to Database", variant="primary")
            add_result = gr.Textbox(label="Result")
            
        with gr.Column(elem_id="col-right"):
            gr.HTML("""
            <div style="text-align: center; font-size: 18px; margin-bottom: 10px;">
                <b>🔍 Find Face Match</b>
            </div>
            """)
            match_image = gr.Image(label="Upload Photo to Match", sources='upload', type="pil")
            match_button = gr.Button("Find Match", variant="primary")
            match_result = gr.Textbox(label="Match Result")
    
    gr.HTML("<hr>")
    
    with gr.Row():
        status_button = gr.Button("Check Status")
        clear_button = gr.Button("Clear Database", variant="stop")
        status_output = gr.Textbox(label="System Status")
    
    # Event handlers
    add_button.click(fn=add_face, inputs=[add_image, add_name], outputs=add_result)
    match_button.click(fn=match_face, inputs=match_image, outputs=match_result)
    status_button.click(fn=get_status, outputs=status_output)
    clear_button.click(fn=clear_database, outputs=status_output)

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