# app.py

import gradio as gr
from huggingface_hub import InferenceClient
from PIL import Image
import io
import os
import numpy as np

# -----------------------------
# Hugging Face API Setup
# -----------------------------
HF_API_TOKEN = os.environ.get("HF_API_TOKEN")  # Add your token in Hugging Face Secrets
client = InferenceClient(token=HF_API_TOKEN)

# -----------------------------
# Model Names
# -----------------------------
MODEL_1 = "prithivMLmods/deepfake-detector-model-v1"   # Deepfake detector
MODEL_2 = "microsoft/dit-base-finetuned-aigc-detection"  # AIGC detector
MODEL_3 = "zhipeixu/fakeshield-v1-22b"                 # Forgery detector

# -----------------------------
# Helper function: overlay mask on image
# -----------------------------
def overlay_mask(image, mask):
    if mask is None:
        return image
    mask = np.array(mask.convert("L"))  # Convert mask to grayscale
    mask = (mask > 128).astype(np.uint8) * 255  # Binary mask
    overlay = Image.new("RGBA", image.size, (255,0,0,100))  # Red overlay
    img_rgba = image.convert("RGBA")
    img_rgba.paste(overlay, mask=Image.fromarray(mask))
    return img_rgba

# -----------------------------
# Main function: Analyze image using 3 models
# -----------------------------
def analyze_image(image):
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    buf.seek(0)

    # -------- MODEL 1: Deepfake Detector --------
    try:
        out1 = client.image_classification(model=MODEL_1, inputs=buf)
        label1 = out1[0]["label"]
        score1 = round(out1[0]["score"] * 100, 2)
    except:
        label1, score1 = "Error", 0
    buf.seek(0)

    # -------- MODEL 2: AIGC Detector --------
    try:
        out2 = client.image_classification(model=MODEL_2, inputs=buf)
        label2 = out2[0]["label"]
        score2 = round(out2[0]["score"] * 100, 2)
    except:
        label2, score2 = "Error", 0
    buf.seek(0)

    # -------- MODEL 3: Forgery / Mask Detector --------
    try:
        out3 = client(inputs=buf, model=MODEL_3)
        explanation = out3.get("explanation", "No manipulation detected")
        mask = out3.get("mask", None)
    except:
        explanation, mask = "Error detecting forgery", None

    # -------- FINAL DECISION (Fusion) --------
    ai_votes = 0
    if "fake" in label1.lower() or "ai" in label1.lower():
        ai_votes += 1
    if "ai" in label2.lower() or "generated" in label2.lower():
        ai_votes += 1

    if ai_votes == 2:
        final_label = "AI-GENERATED"
    elif ai_votes == 1:
        final_label = "Possibly AI-GENERATED"
    else:
        final_label = "REAL IMAGE"

    # Overlay mask if exists
    output_image = overlay_mask(image, mask)

    # Return outputs
    return (
        output_image,
        f"{final_label}",
        f"Deepfake Model: {label1} ({score1}%)\nAIGC Model: {label2} ({score2}%)\nForgery Detector: {explanation}"
    )

# -----------------------------
# Gradio Interface
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("<h2 style='text-align:center'>AI DeepFake & Manipulation Detector</h2>")
    
    with gr.Row():
        inp = gr.Image(type="pil", label="Upload Image")
        out_img = gr.Image(type="pil", label="Result Image with Mask Overlay")
    
    out_text = gr.Textbox(label="Detection Result & Explanation", lines=8)
    
    btn = gr.Button("Analyze Image")
    btn.click(fn=analyze_image, inputs=[inp], outputs=[out_img, out_text, out_text])

# -----------------------------
# Launch Gradio App
# -----------------------------
demo.launch()