import gradio as gr
from transformers import pipeline
from PIL import Image
import traceback
import time
import threading

# Models
models = [
    ("Ateeqq/ai-vs-human-image-detector", "ateeq"),
    ("umm-maybe/AI-image-detector", "umm_maybe"),
    ("dima806/ai_vs_human_generated_image_detection", "dimma"),
]

pipes = []
for model_id, _ in models:
    try:
        pipes.append((model_id, pipeline("image-classification", model=model_id)))
        print(f"Loaded {model_id}")
    except Exception as e:
        print(f"Error loading {model_id}: {e}")

def predict_image(image: Image.Image):
    try:
        results = []
        for _, pipe in pipes:
            res = pipe(image)[0]
            results.append(res)

        final_result = results[0]
        label = final_result["label"].lower()
        score = final_result["score"] * 100

        if "ai" in label or "fake" in label:
            verdict = f"🧠 AI-Generated ({score:.1f}% confidence)"
            color = "#007BFF"
        else:
            verdict = f"🧍 Human-Made ({score:.1f}% confidence)"
            color = "#4CAF50"

        html = f"""
        <div class='result-box' style="
            background: linear-gradient(135deg, {color}33, #1a1a1a);
            border: 2px solid {color};
            border-radius: 15px;
            padding: 25px;
            text-align: center;
            color: white;
            font-size: 20px;
            font-weight: 600;
            box-shadow: 0 0 20px {color}55;
            animation: fadeIn 0.6s ease-in-out;
        ">
            {verdict}
        </div>
        """
        return html
    except Exception as e:
        traceback.print_exc()
        return f"<div style='color:red;'>Error analyzing image: {str(e)}</div>"

# CSS for sleek glowing pulse
css = """
body, .gradio-container {
    font-family: 'Poppins', sans-serif !important;
    background: transparent !important;
}
h1 {
    text-align: center;
    font-weight: 700;
    color: #007BFF;
    margin-bottom: 10px;
}
.gr-button-primary {
    background-color: #007BFF !important;
    color: white !important;
    font-weight: 600;
    border-radius: 10px;
    height: 45px;
}
.gr-button-secondary {
    background-color: #dc3545 !important;
    color: white !important;
    border-radius: 10px;
    height: 45px;
}
#pulse-loader {
    width: 100%;
    height: 4px;
    background: linear-gradient(90deg, #007BFF, #00C3FF);
    animation: pulse 1.2s infinite ease-in-out;
    border-radius: 2px;
    box-shadow: 0 0 10px #007BFF;
}
@keyframes pulse {
    0% { transform: scaleX(0.1); opacity: 0.6; }
    50% { transform: scaleX(1); opacity: 1; }
    100% { transform: scaleX(0.1); opacity: 0.6; }
}
@keyframes fadeIn {
    from { opacity: 0; transform: scale(0.95); }
    to { opacity: 1; transform: scale(1); }
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    gr.Markdown("<h1>🔍 AI Image Detector</h1>")

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Upload an image")
            analyze_button = gr.Button("Analyze", variant="primary")
            clear_button = gr.Button("Clear", variant="secondary")
            loader = gr.HTML("")
        with gr.Column(scale=1):
            output = gr.HTML(label="Result")

    def analyze(img):
        if img is None:
            return ("", "<div style='color:red;'>Please upload an image first!</div>")
        loader_html = "<div id='pulse-loader'></div>"
        yield (loader_html, "")  # instantly show loader

        # do analysis in background
        result = predict_image(img)
        yield ("", result)  # hide loader, show result

    analyze_button.click(analyze, inputs=image_input, outputs=[loader, output])
    clear_button.click(lambda: ("", ""), outputs=[loader, output])

demo.launch()