# -*- coding: utf-8 -*-
import gradio as gr
import numpy as np
import cv2
from deepface import DeepFace
import json
import time
import os

# Ensure OpenCV finds the Haar Cascade file
# This might be needed if the default path isn't automatically found in the HF environment
cv2_base_dir = os.path.dirname(os.path.abspath(cv2.__file__))
haar_cascade_path = os.path.join(cv2_base_dir, 'data', 'haarcascade_frontalface_default.xml')

# Check if the cascade file exists, raise error if not
if not os.path.isfile(haar_cascade_path):
    # Attempt backup location if primary fails (less likely needed but safe)
    backup_path = "/home/user/.local/lib/python3.x/site-packages/cv2/data/haarcascade_frontalface_default.xml" # Adjust python version if needed
    if os.path.isfile(backup_path):
         haar_cascade_path = backup_path
    else:
        raise RuntimeError(f"Could not find Haar Cascade file at expected locations: {haar_cascade_path} or backup paths.")


EMBEDDINGS_FILE = "stored_embeddings.json"

def load_embeddings():
    """Load stored face embeddings (expects list of lists)."""
    try:
        with open(EMBEDDINGS_FILE, "r") as f:
            embeddings_data = json.load(f)
        # --- CRITICAL BUG FIX AREA (See Analysis) ---
        # Assuming the file SHOULD contain the structure saved by register_face:
        # a list of dictionaries like [{"embedding": [...], "name": "...", "timestamp": "..."}]
        # We need to extract just the embedding lists for the current logic.
        embeddings = [np.array(item["embedding"]) for item in embeddings_data if "embedding" in item]
        return embeddings
    except (FileNotFoundError, json.JSONDecodeError):
        return []
    except Exception as e:
        print(f"Error loading embeddings: {e}. Assuming empty list.") # Add more logging
        return []


# def save_embeddings(embeddings): # This function is defined but never used
#       with open(EMBEDDINGS_FILE, "w") as f:
#         json.dump([embedding.tolist() for embedding in embeddings], f)

# Load embeddings when the script starts
stored_embeddings = load_embeddings()
print(f"Loaded {len(stored_embeddings)} embeddings on startup.") # Add logging

def extract_face_embedding_from_frame(frame):
    """Extract a facial embedding from a single frame (image)."""
    # Use the validated haar_cascade_path
    face_cascade = cv2.CascadeClassifier(haar_cascade_path)
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

    if len(faces) > 0:
        # Process only the first detected face
        x, y, w, h = faces[0]
        face = frame[y:y + h, x:x + w]

        # Create a copy for marking to avoid modifying the original frame if needed elsewhere
        marked_frame = frame.copy()
        cv2.rectangle(marked_frame, (x, y), (x+w, y+h), (0, 255, 0), 2)

        # Use Facenet model for representation
        # Ensure backend='auto' or specify one like 'opencv', 'ssd', 'mtcnn' if needed
        # Error handling for representation is important
        try:
            embedding_objs = DeepFace.represent(face, model_name="Facenet", enforce_detection=False) # Don't re-detect
            if not embedding_objs:
                 raise Exception("DeepFace could not generate embedding.")
            embedding = embedding_objs[0]["embedding"] # Access the first item's embedding
            return np.array(embedding), marked_frame, face
        except Exception as represent_error:
            raise Exception(f"Failed to generate embedding: {represent_error}")

    # If no faces were detected by Haar Cascade
    raise Exception("No face detected in the frame.")

def verify_face_from_webcam(video, save_embedding=False, name=""):
    """Verifies face from video OR attempts registration if save_embedding is True."""
    global stored_embeddings

    if video is None:
        # Return signature must match the Gradio outputs: image, string, image, number
        return None, "Please record a video first.", None, 0

    try:
        cap = cv2.VideoCapture(video)
        frames = []
        frame_count = 0
        while frame_count < 60: # Limit frames read to prevent memory issues
            ret, frame = cap.read()
            if not ret:
                break
            frames.append(frame)
            frame_count += 1
        cap.release()

        if not frames:
            # Return signature: image, string, image, number
            return None, "Error: Could not read video.", None, 0

        # Use the middle frame for analysis
        frame = frames[len(frames)//2]

        # Extract embedding and get marked frame + face crop
        embedding, marked_frame, face = extract_face_embedding_from_frame(frame)

        max_similarity = 0
        is_match = False

        # Compare against loaded embeddings
        for stored_embedding in stored_embeddings:
            # Cosine Similarity Calculation
            similarity = np.dot(embedding, stored_embedding) / (np.linalg.norm(embedding) * np.linalg.norm(stored_embedding))
            max_similarity = max(max_similarity, similarity)

            # Verification threshold
            if similarity > 0.7: # Hardcoded threshold
                is_match = True
                # If only verifying (save_embedding=False), return success
                if not save_embedding:
                    # Return signature: image, string, image, number
                    return marked_frame, f"✅ **AUTHENTICATED!** Similarity score: {similarity:.2f}", face, similarity
                # If attempting to save but already matched, return warning (prevents duplicates somewhat)
                else:
                     # Return signature: image, string, image, number (use 0 for confidence here maybe?)
                     return marked_frame, f"⚠️ This face seems to be already in the database (Similarity: {similarity:.2f}). Registration aborted.", face, similarity


        # --- Logic for the "save_embedding" flag from the Authentication tab ---
        # This part is problematic because the Auth tab doesn't provide a name and
        # the return signature expects 4 items, but registration logic only naturally provides 3.
        if save_embedding:
            # This code block within verify_face_from_webcam is likely UNREACHABLE
            # or will error due to mismatched return values for the Gradio component expecting 4 outputs.
            # It duplicates logic from `register_face` unnecessarily.
            # It's better to remove this `save_embedding` logic entirely from the verification function.
            # Keeping it here as per "do not change code" rule, but highlighting it as flawed.
            print("WARNING: 'save_embedding' path triggered in 'verify_face_from_webcam'. This is likely incorrect.")
            if not name.strip():
                # Trying to match 4 return values: image, string, image, number
                return marked_frame, "⚠️ Please enter a name to save this face. (This shouldn't happen from Auth tab)", face, max_similarity # Problematic return

            # (This saving logic is redundant with register_face)
            metadata = {
                "embedding": embedding.tolist(),
                "name": name.strip(),
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
            }
            # --- BUG AREA ---
            # Should load the FULL data, append, then save, not just append embedding to in-memory list
            try:
                with open(EMBEDDINGS_FILE, "r") as f:
                    all_data = json.load(f)
            except (FileNotFoundError, json.JSONDecodeError):
                all_data = []
            all_data.append(metadata)
            # Need to update the global variable too if we want consistency during session
            stored_embeddings.append(embedding) # Appending only embedding, not metadata
            with open(EMBEDDINGS_FILE, "w") as f:
                 json.dump(all_data, f)

            # Trying to match 4 return values: image, string, image, number
            return marked_frame, f"✅ **NEW FACE REGISTERED!** Welcome, {name}! (From verify_face_from_webcam - likely wrong)", face, 1.0 # Problematic return


        # If we are here, it means we were only verifying (save_embedding=False) and no match was found > 0.7
        # Return signature: image, string, image, number
        return marked_frame, f"❌ **ACCESS DENIED!** Highest similarity: {max_similarity:.2f}", face, max_similarity

    except Exception as e:
        print(f"Error in verify_face_from_webcam: {e}") # Log the error
        # Return signature: image, string, image, number
        return None, f"⚠️ Error processing video: {str(e)}", None, 0


def register_face(video, name=""):
    """Registers a new face from video."""
    global stored_embeddings # Reference the global list

    if video is None:
        # Return signature: image, string, image
        return None, "Please record a video first.", None

    if not name.strip():
        # Return signature: image, string, image
        # Need to return *something* for the images if no video processed
        return None, "⚠️ Please enter a name to save this face.", None

    try:
        cap = cv2.VideoCapture(video)
        frames = []
        frame_count = 0
        while frame_count < 60: # Limit frames
            ret, frame = cap.read()
            if not ret:
                break
            frames.append(frame)
            frame_count += 1
        cap.release()

        if not frames:
            # Return signature: image, string, image
            return None, "Error: Could not read video.", None

        # Use the middle frame
        frame = frames[len(frames)//2]

        # Extract embedding
        embedding, marked_frame, face = extract_face_embedding_from_frame(frame)

        # Check for existing faces (using a higher threshold for registration uniqueness)
        for i, stored_embedding in enumerate(stored_embeddings):
            similarity = np.dot(embedding, stored_embedding) / (np.linalg.norm(embedding) * np.linalg.norm(stored_embedding))
            # Using a stricter threshold during registration check
            if similarity > 0.9: # Hardcoded threshold
                # Return signature: image, string, image
                # Optionally retrieve name if metadata loading is fixed:
                # existing_name = get_name_for_embedding(i) # Requires modification
                return marked_frame, f"⚠️ This face seems highly similar to an existing one in the database (Similarity: {similarity:.2f}). Registration aborted.", face


        # If checks pass, prepare metadata and save
        metadata = {
            "embedding": embedding.tolist(),
            "name": name.strip(),
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
        }

        # --- Corrected Saving Logic ---
        all_data = []
        try:
            # Read the existing full data file
            with open(EMBEDDINGS_FILE, "r") as f:
                all_data = json.load(f)
                # Ensure it's a list
                if not isinstance(all_data, list):
                    print(f"Warning: Embeddings file was not a list. Re-initializing.")
                    all_data = []
        except (FileNotFoundError, json.JSONDecodeError):
            # If file doesn't exist or is invalid, start with an empty list
            all_data = []
        except Exception as e:
             print(f"Error reading embeddings file before save: {e}. Starting fresh.")
             all_data = []


        # Append the new metadata dictionary
        all_data.append(metadata)

        # Write the updated list of dictionaries back to the file
        try:
            with open(EMBEDDINGS_FILE, "w") as f:
                json.dump(all_data, f, indent=4) # Add indent for readability
            # Update the in-memory list ONLY IF save was successful
            stored_embeddings.append(embedding) # Add the numpy array to the in-memory list for immediate use
            print(f"Successfully registered {name}. Total embeddings in memory: {len(stored_embeddings)}")
        except Exception as e:
             print(f"Error writing embeddings file: {e}")
             # Return signature: image, string, image
             return marked_frame, f"⚠️ Error saving embedding to file: {e}", face


        # Return signature: image, string, image
        return marked_frame, f"✅ **NEW FACE REGISTERED!** Welcome, {name}!", face

    except Exception as e:
        print(f"Error in register_face: {e}") # Log the error
        # Return signature: image, string, image
        return None, f"⚠️ Error processing registration: {str(e)}", None


def get_registered_count():
    """Get the count of registered faces from the JSON file."""
    try:
        with open(EMBEDDINGS_FILE, "r") as f:
            # Load the full data which should be a list of dictionaries
            data = json.load(f)
            if isinstance(data, list):
                return f"Total registered faces: {len(data)}"
            else:
                return "Registered faces: Error (Invalid data format)"
    except (FileNotFoundError, json.JSONDecodeError):
        return "Total registered faces: 0"
    except Exception as e:
        print(f"Error reading count: {e}")
        return "Registered faces: Error reading file"

# Define CSS
css = """
.gradio-container {
    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
    color: white;
}
.title-container {
    text-align: center;
    margin-bottom: 2rem;
}
.title-container h1 {
    font-size: 2.5rem;
    background: -webkit-linear-gradient(#eee, #0e93e0);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 0.5rem;
}
.title-container p {
    font-size: 1.2rem;
    color: #ccc;
}
.status-authenticated { /* This class seems unused in the Markdown outputs */
    color: #4CAF50;
    font-weight: bold;
    font-size: 1.2rem;
}
.status-denied { /* This class seems unused in the Markdown outputs */
    color: #F44336;
    font-weight: bold;
    font-size: 1.2rem;
}
.output-image {
    border-radius: 10px;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.3);
}
.footer {
    text-align: center;
    margin-top: 2rem;
    font-size: 0.9rem;
    color: #888;
}
/* Removed .face-container as it wasn't applied via elem_classes */
/* You might want to apply custom classes directly to Markdown/Image components if needed */
"""

# Define the Gradio Interface
with gr.Blocks(css=css, theme=gr.themes.Soft()) as iface:
    gr.HTML("""
    <div class="title-container">
        <h1>🔐 Secure Face Authentication System</h1>
        <p>Record a video using your webcam to authenticate or register a new face</p>
    </div>
    """)

    with gr.Tabs():
        with gr.Tab("Authentication"):
            with gr.Row():
                with gr.Column(scale=3):
                    # Use 'webcam' source for direct recording if preferred over upload
                    # format="mp4" might require ffmpeg installed in the environment
                    video_input = gr.Video(label="Record/Upload a short video (1-2 seconds)", sources=["webcam", "upload"], format="mp4")
                    auth_button = gr.Button("Authenticate", variant="primary")

                with gr.Column(scale=2):
                    output_image = gr.Image(label="Detection Result", type="numpy") # Specify type
                    face_image = gr.Image(label="Extracted Face", type="numpy", elem_classes="output-image") # Specify type

            with gr.Row():
                 # Use Markdown for rich text status
                status_text = gr.Markdown("Ready for authentication...")
                confidence_meter = gr.Number(label="Confidence Score", value=0) # Removed min/max as similarity can vary

        with gr.Tab("Registration"):
            with gr.Row():
                with gr.Column(scale=3):
                    # Use 'webcam' source for direct recording
                    reg_video_input = gr.Video(label="Record/Upload a short video (1-2 seconds)", sources=["webcam", "upload"], format="mp4")
                    name_input = gr.Textbox(label="Enter your name", placeholder="John Doe")
                    register_button = gr.Button("Register New Face", variant="secondary")

                with gr.Column(scale=2):
                    reg_output_image = gr.Image(label="Detection Result", type="numpy") # Specify type
                    reg_face_image = gr.Image(label="Extracted Face", type="numpy", elem_classes="output-image") # Specify type

            with gr.Row():
                 # Use Markdown for rich text status
                reg_status_text = gr.Markdown("Ready for registration...")

    with gr.Row():
        # Initial state loaded by function call
        stats_text = gr.Markdown(get_registered_count())
        # Button to refresh the count
        refresh_button = gr.Button("Refresh Stats", variant="tertiary", size="sm")


    gr.HTML("""
    <div class="footer">
        <p>Powered by DeepFace + Gradio | © 2025 Face Authentication System</p>
    </div>
    """)

    # Define interactions
    # Authentication button click action
    auth_button.click(
        fn=verify_face_from_webcam,
        # Pass video input. The save_embedding flag is hardcoded False here.
        # The name input isn't relevant for pure authentication.
        inputs=[video_input, gr.Checkbox(value=False, visible=False)],
        # Map outputs to the correct components
        outputs=[output_image, status_text, face_image, confidence_meter]
    )

    # Registration button click action
    register_button.click(
        fn=register_face,
        # Pass registration video and name input
        inputs=[reg_video_input, name_input],
        # Map outputs to the correct components
        outputs=[reg_output_image, reg_status_text, reg_face_image]
    )

    # Refresh button click action
    refresh_button.click(fn=get_registered_count, inputs=None, outputs=[stats_text])


# Launch the interface
# The if __name__ == "__main__": block is good practice but not strictly required by Spaces
if __name__ == "__main__":
    iface.launch() # debug=True can be helpful locally but remove for production