from PIL import Image
from io import BytesIO
import numpy as np
import base64
import uuid
import pymongo
import certifi
from deepface import DeepFace
import gradio as gr
import json

# =======================
# MongoDB setup
# =======================
client = pymongo.MongoClient(
    "mongodb+srv://username:password@cluster0.n8pboqq.mongodb.net/?retryWrites=true&w=majority",
    tlsCAFile=certifi.where()
)
db = client["faceDB"]
face_collection = db["faceImg"]

# =======================
# Upload ảnh + MongoDB
# =======================
def upload_image(image, name):
    try:
        if isinstance(image, Image.Image):
            image = np.array(image)

        image_id = str(uuid.uuid4())
        pil_image = Image.fromarray(image)
        buffer = BytesIO()
        pil_image.save(buffer, format="PNG")
        img_bytes = buffer.getvalue()
        img_base64 = base64.b64encode(img_bytes).decode('utf-8')

        doc = {
            "_id": image_id,
            "name": name,
            "image": img_base64
        }
        face_collection.insert_one(doc)

        return f"✅ Upload thành công!\nID: {image_id}\nTên: {name}", image_id
    except Exception as e:
        return f"❌ Lỗi upload: {str(e)}", None

# =======================
# Nhận diện khuôn mặt
# =======================
def recognize_face(image, image_id):
    try:
        doc = face_collection.find_one({"_id": image_id})
        if not doc:
            return "❌ Ảnh ID không tồn tại"

        # Decode base64 → lưu tạm
        img_bytes = base64.b64decode(doc['image'])
        temp_path = f"temp_{image_id}.png"
        with open(temp_path, "wb") as f:
            f.write(img_bytes)

        # Lưu input ảnh tạm
        if isinstance(image, Image.Image):
            image = np.array(image)
        pil_input = Image.fromarray(image)
        input_path = f"input_{image_id}.png"
        pil_input.save(input_path)

        result = DeepFace.verify(img1_path=input_path, img2_path=temp_path)
        output = {
            "Verified": result["verified"],
            "Khoảng cách": round(result["distance"], 4),
            "Tên ảnh gốc": doc['name']
        }
        return json.dumps(output, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ Lỗi nhận diện: {str(e)}"

# =======================
# Phân tích khuôn mặt
# =======================
def analyze_face(image):
    try:
        if isinstance(image, Image.Image):
            image = np.array(image)
        pil_image = Image.fromarray(image)
        temp_path = f"analyze_{uuid.uuid4().hex}.png"
        pil_image.save(temp_path)

        result = DeepFace.analyze(
            img_path=temp_path,
            actions=["age", "gender", "emotion", "race"],
            enforce_detection=False
        )
        info = result[0]
        output = {
            "Tuổi ước tính": info["age"],
            "Giới tính": "Nam" if info["dominant_gender"] == "Man" else "Nữ",
            "Cảm xúc chính": info["dominant_emotion"],
            "Chủng tộc": info["dominant_race"]
        }
        return json.dumps(output, ensure_ascii=False, indent=2)
    except Exception as e:
        return f"❌ Lỗi phân tích: {str(e)}"

# =======================
# Gradio interface
# =======================
with gr.Blocks() as demo:
    gr.Markdown("## 📁 Upload ảnh vào MongoDB")
    with gr.Row():
        upload_input = gr.Image(label="Upload ảnh", type="numpy")
        name_input = gr.Textbox(label="Tên người trong ảnh")
        upload_btn = gr.Button("Upload")
        upload_output = gr.Textbox(label="Kết quả Upload")
        image_id_store = gr.State()

    upload_btn.click(upload_image, inputs=[upload_input, name_input], outputs=[upload_output, image_id_store])

    gr.Markdown("## 🆔 Nhận diện khuôn mặt")
    with gr.Row():
        recognize_input = gr.Image(label="Ảnh cần nhận diện", type="numpy")
        recognize_btn = gr.Button("Nhận diện")
        recognize_output = gr.Textbox(label="Kết quả nhận diện")

    recognize_btn.click(recognize_face, inputs=[recognize_input, image_id_store], outputs=recognize_output)

    gr.Markdown("## 🔍 Phân tích khuôn mặt")
    with gr.Row():
        analyze_input = gr.Image(label="Ảnh cần phân tích", type="numpy")
        analyze_btn = gr.Button("Phân tích")
        analyze_output = gr.Textbox(label="Kết quả phân tích")

    analyze_btn.click(analyze_face, inputs=analyze_input, outputs=analyze_output)

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