import os
import time
import cv2
import numpy as np
import torch
from PIL import Image as PILImage
import gc
from datetime import datetime, timedelta

from ultralytics import YOLO
from keras_facenet import FaceNet
from transformers import pipeline
import gradio as gr

# -----------------------------
# Device Setup
# -----------------------------
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Disable GPU
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")

# -----------------------------
# Load YOLOv8 Face Model
# -----------------------------
MODEL_PATH = "yolov8n-face.pt"  # make sure this is in your repo
if not os.path.exists(MODEL_PATH):
    raise FileNotFoundError(f"Model file not found at {MODEL_PATH}")

print("Loading YOLOv8 face detector...")
face_model = YOLO(MODEL_PATH).to(DEVICE)
print("YOLOv8 loaded")

# -----------------------------
# Load FaceNet Embedder
# -----------------------------
print("Loading FaceNet model...")
embedder = FaceNet()
print("FaceNet loaded")

# -----------------------------
# Load HuggingFace Age & Gender Models
# -----------------------------
print("Loading HuggingFace Age & Gender models...")
age_model = pipeline(
    "image-classification",
    model="prithivMLmods/Age-Classification-SigLIP2",
    device=-1
)
gender_model = pipeline(
    "image-classification",
    model="dima806/fairface_gender_image_detection",
    device=-1
)
print("Age & Gender models loaded")

# -----------------------------
# Face DB
# -----------------------------
FACE_DB = []
NEXT_ID = 1

# -----------------------------
# Utilities
# -----------------------------
def clean_gpu():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# -----------------------------
# Core Function
# -----------------------------
def process_image(image: PILImage):
    global NEXT_ID, FACE_DB
    start_time = time.time()
    rgb_img = np.array(image)

    # Detect faces
    results = face_model(rgb_img, verbose=False)
    boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)

    now = datetime.now()
    # Remove old entries (>1 hour)
    FACE_DB = [f for f in FACE_DB if now - f["time"] <= timedelta(hours=1)]

    faces = []
    for (x1, y1, x2, y2) in boxes:
        face_crop = rgb_img[y1:y2, x1:x2]
        if face_crop.size == 0:
            continue

        face_embedding = embedder.embeddings([face_crop])[0]
        assigned_id = None
        age_pred, gender_pred = "unknown", "unknown"

        # Compare with DB
        if FACE_DB:
            similarities = [cosine_similarity(face_embedding, entry["embedding"]) for entry in FACE_DB]
            best_idx = int(np.argmax(similarities))
            best_score = similarities[best_idx]
            if best_score > 0.6:
                assigned_id = FACE_DB[best_idx]["id"]
                FACE_DB[best_idx]["time"] = now
                FACE_DB[best_idx]["seen_count"] += 1
                age_pred = FACE_DB[best_idx]["age"]
                gender_pred = FACE_DB[best_idx]["gender"]

        # New face
        if assigned_id is None:
            assigned_id = NEXT_ID
            face_pil = PILImage.fromarray(face_crop)
            try:
                age_pred = age_model(face_pil)[0]["label"]
                gender_pred = gender_model(face_pil)[0]["label"]
            except Exception:
                age_pred, gender_pred = "unknown", "unknown"

            FACE_DB.append({
                "id": assigned_id,
                "embedding": face_embedding,
                "time": now,
                "seen_count": 1,
                "age": age_pred,
                "gender": gender_pred
            })
            NEXT_ID += 1

        faces.append({
            "id": assigned_id,
            "age": age_pred,
            "gender": gender_pred,
            "box": [int(x1), int(y1), int(x2), int(y2)]
        })

    total_time = round(time.time() - start_time, 3)
    clean_gpu()

    summary = [
        {
            "id": entry["id"],
            "seen_count": entry["seen_count"],
            "age": entry["age"],
            "gender": entry["gender"]
        } for entry in FACE_DB
    ]

    return {
        "status": "ok",
        "faces": faces,
        "face_count": len(faces),
        "processing_time_sec": total_time,
        "active_faces_last_hour": len(FACE_DB),
        "seen_summary_last_hour": summary
    }

# -----------------------------
# Gradio Interface
# -----------------------------
demo = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs="json",
    title="Face Recognition + Age & Gender",
    description="YOLOv8 + FaceNet + HuggingFace Age/Gender"
)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        inbrowser=False,
        share=True
    )