# Install required packages
import os
import subprocess
import sys
import importlib
import pkg_resources

def install_package(package, version=None):
    package_spec = f"{package}=={version}" if version else package
    print(f"Installing {package_spec}...")
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-cache-dir", package_spec])
    except subprocess.CalledProcessError as e:
        print(f"Failed to install {package_spec}: {e}")
        raise

def ensure_package(package, version=None):
    try:
        if version:
            pkg_resources.require(f"{package}=={version}")
        else:
            importlib.import_module(package)
        print(f"{package} is already installed with the correct version.")
    except (ImportError, pkg_resources.VersionConflict, pkg_resources.DistributionNotFound) as e:
        print(f"Package requirement failed: {e}")
        install_package(package, version)

# Check environment and install dependencies
if not os.path.exists("/.dockerenv") and not os.path.exists("/kaggle"):
    print("Setting up environment...")
    
    # Install core dependencies
    ensure_package("numpy", "1.23.5")
    ensure_package("protobuf", "3.20.3")
    ensure_package("tensorflow", "2.10.0")
    ensure_package("opencv-python-headless", "4.7.0.72")  # Fix for cv2 error
    ensure_package("deepface", "0.0.79")
    ensure_package("gradio", "3.50.2")
    
    # Install additional required packages
    for pkg in ["matplotlib", "pillow", "pandas"]:
        ensure_package(pkg)

# Now import the required modules
import gradio as gr
import json
import cv2
import numpy as np
from deepface import DeepFace
import matplotlib.pyplot as plt
from PIL import Image
import tempfile
import pandas as pd
import shutil

# Google Drive integration (for Colab users)
if 'google.colab' in sys.modules:
    from google.colab import drive
    drive.mount('/content/drive')

def verify_faces(img1, img2, threshold=0.70, model="VGG-Face"):
    temp_dir = tempfile.mkdtemp()
    img1_path = os.path.join(temp_dir, "image1.jpg")
    img2_path = os.path.join(temp_dir, "image2.jpg")
    
    if isinstance(img1, np.ndarray):
        Image.fromarray(img1).save(img1_path)
    else:
        img1.save(img1_path)
        
    if isinstance(img2, np.ndarray):
        Image.fromarray(img2).save(img2_path)
    else:
        img2.save(img2_path)
    
    try:
        result = DeepFace.verify(
            img1_path=img1_path, 
            img2_path=img2_path,
            model_name=model,
            distance_metric="cosine",
            threshold=threshold
        )
        
        fig, ax = plt.subplots(1, 2, figsize=(10, 5))
        
        img1_display = cv2.imread(img1_path)
        img1_display = cv2.cvtColor(img1_display, cv2.COLOR_BGR2RGB)
        img2_display = cv2.imread(img2_path)
        img2_display = cv2.cvtColor(img2_display, cv2.COLOR_BGR2RGB)
        
        ax[0].imshow(img1_display)
        ax[0].set_title("Image 1")
        ax[0].axis("off")
        
        ax[1].imshow(img2_display)
        ax[1].set_title("Image 2")
        ax[1].axis("off")
        
        verification_result = "✅ FACE MATCHED" if result["verified"] else "❌ FACE NOT MATCHED"
        confidence = round((1 - result["distance"]) * 100, 2)
        
        plt.suptitle(f"{verification_result}\nConfidence: {confidence}%\nDistance: {result['distance']:.4f}", 
                    fontsize=16, fontweight='bold', 
                    color='green' if result["verified"] else 'red')
        
        plt.tight_layout()
        
        os.remove(img1_path)
        os.remove(img2_path)
        os.rmdir(temp_dir)
        
        return fig, json.dumps(result, indent=2)
    
    except Exception as e:
        if os.path.exists(img1_path):
            os.remove(img1_path)
        if os.path.exists(img2_path):
            os.remove(img2_path)
        if os.path.exists(temp_dir):
            os.rmdir(temp_dir)
        
        error_msg = f"Error: {str(e)}"
        if "No face detected" in str(e):
            error_msg = "No face detected in one or both images. Please try different images."
        
        return None, error_msg

def find_faces(query_img, db_folder, threshold=0.70, model="VGG-Face"):
    temp_dir = tempfile.mkdtemp()
    query_path = os.path.join(temp_dir, "query.jpg")
    
    if isinstance(query_img, np.ndarray):
        Image.fromarray(query_img).save(query_path)
    else:
        query_img.save(query_path)
    
    # Handle cloud storage paths and uploaded files
    if isinstance(db_folder, str):
        # Check if it's a Google Drive path
        if db_folder.startswith("/content/drive"):
            db_path = db_folder
        else:
            # Handle regular path
            db_path = os.path.abspath(db_folder)
            if not os.path.exists(db_path):
                return None, "Invalid database path - directory does not exist"
    else:
        # Handle uploaded files
        db_path = os.path.join(temp_dir, "db")
        os.makedirs(db_path, exist_ok=True)
        
        for i, file in enumerate(db_folder):
            orig_filename = file.orig_name
            file_ext = os.path.splitext(orig_filename)[1]
            new_filename = f"image_{i}{file_ext}"
            shutil.copy(file.name, os.path.join(db_path, new_filename))
    
    try:
        dfs = DeepFace.find(
            img_path=query_path,
            db_path=db_path,
            model_name=model,
            distance_metric="cosine",
            threshold=threshold,
            silent=True
        )
        
        if isinstance(dfs, list):
            if len(dfs) == 0:
                return None, "No matching faces found in the database."
            df = dfs[0]
        else:
            df = dfs
            
        if df.empty:
            return None, "No matching faces found in the database."
        
        df = df.sort_values(by=["distance"])
        
        num_matches = min(4, len(df))
        fig, axes = plt.subplots(1, num_matches + 1, figsize=(15, 5))
        
        query_display = cv2.imread(query_path)
        query_display = cv2.cvtColor(query_display, cv2.COLOR_BGR2RGB)
        axes[0].imshow(query_display)
        axes[0].set_title("Query Image")
        axes[0].axis("off")
        
        valid_matches = 0
        for i in range(num_matches):
            if i >= len(df):
                break
                
            match_path = df.iloc[i]["identity"]
            if not os.path.exists(match_path):
                continue
                
            distance = df.iloc[i]["distance"]
            confidence = round((1 - distance) * 100, 2)
            
            try:
                match_img = cv2.imread(match_path)
                if match_img is None:
                    continue
                    
                match_img = cv2.cvtColor(match_img, cv2.COLOR_BGR2RGB)
                axes[valid_matches+1].imshow(match_img)
                axes[valid_matches+1].set_title(f"Match #{valid_matches+1}\nConfidence: {confidence}%")
                axes[valid_matches+1].axis("off")
                valid_matches += 1
            except:
                continue
        
        # Hide empty axes
        for j in range(valid_matches+1, num_matches+1):
            axes[j].axis("off")
        
        plt.suptitle(f"Found {len(df)} matching faces", fontsize=16, fontweight='bold')
        plt.tight_layout()
        
        results = df[["identity", "distance"]].copy()
        results["confidence"] = (1 - results["distance"]) * 100
        results["confidence"] = results["confidence"].round(2)
        results = results.rename(columns={"identity": "Image Path"})
        
        os.remove(query_path)
        if not isinstance(db_folder, str):
            shutil.rmtree(db_path)
        
        return fig, results.to_dict('records')
    
    except Exception as e:
        if os.path.exists(query_path):
            os.remove(query_path)
        
        error_msg = f"Error: {str(e)}"
        if "No face detected" in str(e):
            error_msg = "No face detected in the query image. Please try a different image."
        elif "No such file or directory" in str(e):
            error_msg = "Invalid database path or corrupted image files"
        
        return None, error_msg

def analyze_face(img, actions=['age', 'gender', 'race', 'emotion']):
    temp_dir = tempfile.mkdtemp()
    img_path = os.path.join(temp_dir, "analyze.jpg")
    
    if isinstance(img, np.ndarray):
        Image.fromarray(img).save(img_path)
    else:
        img.save(img_path)
    
    try:
        results = DeepFace.analyze(
            img_path=img_path,
            actions=actions,
            enforce_detection=True,
            detector_backend='opencv'
        )
        
        if isinstance(results, list):
            num_faces = len(results)
        else:
            num_faces = 1
            results = [results]
        
        fig = plt.figure(figsize=(14, 7))
        
        img_display = cv2.imread(img_path)
        img_display = cv2.cvtColor(img_display, cv2.COLOR_BGR2RGB)
        
        main_ax = plt.subplot2grid((2, 4), (0, 0), colspan=2, rowspan=2)
        main_ax.imshow(img_display)
        main_ax.set_title(f"Analyzed Image ({num_faces} face{'s' if num_faces > 1 else ''} detected)")
        main_ax.axis('off')
        
        for i, face_result in enumerate(results):
            if i >= 4:
                break
                
            age = face_result.get('age', 'N/A')
            gender = face_result.get('dominant_gender', 'N/A')
            race = face_result.get('dominant_race', 'N/A')
            emotion = face_result.get('dominant_emotion', 'N/A')
            
            gender_conf = 'N/A'
            if 'gender' in face_result and isinstance(face_result['gender'], dict):
                for g, conf in face_result['gender'].items():
                    if g.lower() == gender.lower():
                        gender_conf = f"{conf:.1f}%"
                        break
            
            race_conf = 'N/A'
            if 'race' in face_result and isinstance(face_result['race'], dict):
                for r, conf in face_result['race'].items():
                    if r.lower() == race.lower():
                        race_conf = f"{conf:.1f}%"
                        break
            
            emotion_conf = 'N/A'
            if 'emotion' in face_result and isinstance(face_result['emotion'], dict):
                for e, conf in face_result['emotion'].items():
                    if e.lower() == emotion.lower():
                        emotion_conf = f"{conf:.1f}%"
                        break
            
            ax = plt.subplot2grid((2, 4), (0 if i < 2 else 1, 2 + (i % 2)))
            
            text = (
                f"Face #{i+1}\n\n"
                f"Age: {age}\n\n"
                f"Gender: {gender} ({gender_conf})\n\n"
                f"Race: {race} ({race_conf})\n\n"
                f"Emotion: {emotion} ({emotion_conf})"
            )
            
            ax.text(0.5, 0.5, text, ha='center', va='center', fontsize=11)
            ax.axis('off')
        
        plt.tight_layout()
        
        os.remove(img_path)
        os.rmdir(temp_dir)
        
        formatted_results = []
        for i, res in enumerate(results[:8]):
            face_data = {
                "face_number": i+1,
                "age": res.get("age", "N/A"),
                "gender": {
                    "dominant": res.get("dominant_gender", "N/A"),
                    "confidence": res.get("gender", {})
                },
                "race": {
                    "dominant": res.get("dominant_race", "N/A"),
                    "confidence": res.get("race", {})
                },
                "emotion": {
                    "dominant": res.get("dominant_emotion", "N/A"),
                    "confidence": res.get("emotion", {})
                }
            }
            formatted_results.append(face_data)
        
        return fig, formatted_results
    
    except Exception as e:
        if os.path.exists(img_path):
            os.remove(img_path)
        if os.path.exists(temp_dir):
            os.rmdir(temp_dir)
        
        error_msg = f"Error: {str(e)}"
        if "No face detected" in str(e):
            error_msg = "No face detected in the image. Please try a different image."
        
        return None, error_msg

# Create Gradio interface
with gr.Blocks(title="Complete Face Recognition Tool", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🔍 Complete Face Recognition Tool
    
    This tool provides three face recognition features:
    - **Verify Faces**: Compare two specific images to check if they contain the same person
    - **Find Faces**: Search for matching faces in a database/folder (supports Google Drive paths in Colab)
    - **Analyze Face**: Determine age, gender, race, and emotion from a facial image
    """)
    
    with gr.Tabs():
        with gr.TabItem("Verify Faces"):
            with gr.Row():
                img1_input = gr.Image(label="First Image", type="pil")
                img2_input = gr.Image(label="Second Image", type="pil")
            
            with gr.Row():
                verify_threshold = gr.Slider(minimum=0.1, maximum=0.9, value=0.6, step=0.05, 
                                       label="Similarity Threshold (lower = stricter matching)")
                verify_model = gr.Dropdown(
                    choices=["VGG-Face", "Facenet", "OpenFace", "DeepFace", "ArcFace"], 
                    value="VGG-Face", 
                    label="Face Recognition Model"
                )
            
            verify_button = gr.Button("Verify Faces", variant="primary")
            
            verify_result_plot = gr.Plot(label="Verification Result")
            verify_json = gr.JSON(label="Technical Details")

        with gr.TabItem("Find Faces"):
            query_img = gr.Image(label="Query Image (Face to find)", type="pil")
            db_path_input = gr.Textbox(
                label="Database Path (folder path or Google Drive path in Colab)",
                placeholder="/content/drive/MyDrive/your_folder"
            )
            db_files_input = gr.File(label="Or upload images for database", file_count="multiple")
            
            with gr.Row():
                find_threshold = gr.Slider(minimum=0.1, maximum=0.9, value=0.6, step=0.05, 
                                     label="Similarity Threshold (lower = stricter matching)")
                find_model = gr.Dropdown(
                    choices=["VGG-Face", "Facenet", "OpenFace", "DeepFace", "ArcFace"], 
                    value="VGG-Face", 
                    label="Face Recognition Model"
                )
            
            find_button = gr.Button("Find Matching Faces", variant="primary")
            
            find_result_plot = gr.Plot(label="Search Results")
            find_results_table = gr.JSON(label="Detailed Results")

        with gr.TabItem("Analyze Face"):
            analyze_img = gr.Image(label="Upload Image for Analysis", type="pil")
            actions_checkboxes = gr.CheckboxGroup(
                choices=["age", "gender", "race", "emotion"],
                value=["age", "gender", "race", "emotion"],
                label="Select Attributes to Analyze"
            )
            
            analyze_button = gr.Button("Analyze Face", variant="primary")
            
            analyze_result_plot = gr.Plot(label="Analysis Results")
            analyze_json = gr.JSON(label="Detailed Analysis")

    # Setup all button clicks
    verify_button.click(
        verify_faces, 
        inputs=[img1_input, img2_input, verify_threshold, verify_model], 
        outputs=[verify_result_plot, verify_json]
    )
    
    find_button.click(
        find_faces,
        inputs=[query_img, db_path_input, find_threshold, find_model],
        outputs=[find_result_plot, find_results_table]
    )
    
    db_files_input.change(
        lambda x: "",
        inputs=db_files_input,
        outputs=db_path_input
    )
    
    analyze_button.click(
        analyze_face,
        inputs=[analyze_img, actions_checkboxes],
        outputs=[analyze_result_plot, analyze_json]
    )

# Launch the app
demo.launch()