import spaces
import gradio as gr
import torch
import tempfile
import os
from PIL import Image
from transformers import pipeline

# ── Model registry ───────────────────────────────────────────────────────────
MODELS = {
    ("fast", "kodf"):        "KoreaPeter/ms-eff-gcvit-deepfake-b0-kodf",
    ("fast", "celeb-df-v2"): "KoreaPeter/ms-eff-gcvit-deepfake-b0-celeb-df-v2",
    ("fast", "ff++"):        "KoreaPeter/ms-eff-gcvit-deepfake-b0-ff-plus-plus",
    ("pro",  "kodf"):        "KoreaPeter/ms-eff-gcvit-deepfake-b5-kodf",
    ("pro",  "celeb-df-v2"): "KoreaPeter/ms-eff-gcvit-deepfake-b5-celeb-df-v2",
    ("pro",  "ff++"):        "KoreaPeter/ms-eff-gcvit-deepfake-b5-ff-plus-plus",
}

_cache: dict = {}

def _get_pipeline(variant: str, dataset: str):
    key = (variant, dataset)
    if key not in _cache:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        _cache[key] = pipeline(
            "image-classification",
            model=MODELS[key],
            trust_remote_code=True,
            device=device,
        )
    return _cache[key]


# ── Inference ─────────────────────────────────────────────────────────────────
@spaces.GPU
def predict(
    image: Image.Image,
    variant: str,
    dataset: str,
    margin_ratio: float,
    conf_thres: float,
    min_face_ratio: float,
    tta_hflip: float,
    top_k: int,
):
    if image is None:
        return None, "⚠️ Please upload an image."

    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
        image.save(f.name)
        tmp_path = f.name

    try:
        pipe = _get_pipeline(variant, dataset)
        results = pipe(
            tmp_path,
            margin_ratio=margin_ratio,
            conf_thres=conf_thres,
            min_face_ratio=min_face_ratio,
            tta_hflip=tta_hflip,
            top_k=top_k,
        )
    except Exception as e:
        return None, f"❌ Pipeline error: {e}"
    finally:
        os.unlink(tmp_path)

    if not results:
        return None, "⚠️ No result returned."

    top = results[0]
    if top.get("label") == "no_face_detected":
        return (
            {"REAL": 0.0, "FAKE": 0.0},
            "😶 No face detected.\nPlease use an image with a clearly visible face.",
        )

    scores = {r["label"].upper(): round(r["score"], 4) for r in results
              if r.get("label") in ("fake", "real")}

    fake_prob = scores.get("FAKE", 0.0)

    if fake_prob >= 0.7:
        verdict = f"🔴 FAKE  ({fake_prob * 100:.1f}% confidence)"
    elif fake_prob >= 0.4:
        verdict = f"🟡 UNCERTAIN  (FAKE {fake_prob * 100:.1f}%)"
    else:
        verdict = f"🟢 REAL  ({(1 - fake_prob) * 100:.1f}% confidence)"

    return scores, verdict


# ── Hero banner ───────────────────────────────────────────────────────────────
HERO = """
<div style="
  background: linear-gradient(135deg, #c2185b 0%, #e91e63 45%, #880e4f 100%);
  border-radius: 12px; padding: 36px 40px; margin-bottom: 8px;
  border: 1px solid #ad1457; box-shadow: 0 4px 24px rgba(194,24,91,0.4);
">
  <div style="display:flex; align-items:center; gap:16px; margin-bottom:12px;">
    <span style="font-size:2.4rem;">🕵️</span>
    <div>
      <h1 style="margin:0; color:#fce4ec; font-size:1.9rem; font-weight:700; letter-spacing:-0.5px;">DeepGuard</h1>
      <p style="margin:4px 0 0; color:#f48fb1; font-size:1rem;">Image Deepfake Detection</p>
    </div>
  </div>
  <p style="margin:0 0 16px; color:#f8bbd0; font-size:0.95rem; line-height:1.6;">
    Deepfake detection powered by <strong style="color:#fce4ec;">MS-EffGCViT</strong>
    &nbsp;·&nbsp; 6 domain-specific models &nbsp;·&nbsp; Asian · Celebrity · Western faces
  </p>
  <div style="display:flex; gap:12px; flex-wrap:wrap;">
    <a href="https://github.com/HanMoonSub/DeepGuard" target="_blank"
       style="background:#880e4f; color:#f48fb1; padding:7px 16px; border-radius:6px;
              text-decoration:none; font-size:0.85rem; border:1px solid #e91e63;">🔗 GitHub</a>
    <a href="https://huggingface.co/KoreaPeter" target="_blank"
       style="background:#880e4f; color:#f48fb1; padding:7px 16px; border-radius:6px;
              text-decoration:none; font-size:0.85rem; border:1px solid #e91e63;">🤗 HuggingFace</a>
    <a href="https://huggingface.co/spaces/KoreaPeter/DeepFake-Video-Detection" target="_blank"
       style="background:#880e4f; color:#f48fb1; padding:7px 16px; border-radius:6px;
              text-decoration:none; font-size:0.85rem; border:1px solid #e91e63;">🎬 Video Detection</a>
    <a href="https://huggingface.co/spaces/KoreaPeter/DeepFake-Detection-XAI" target="_blank"
       style="background:#880e4f; color:#f48fb1; padding:7px 16px; border-radius:6px;
              text-decoration:none; font-size:0.85rem; border:1px solid #e91e63;">🔬 DeepFake XAI</a>
  </div>
</div>
"""

# ── Theme & CSS ───────────────────────────────────────────────────────────────
THEME = gr.themes.Base(
    primary_hue="pink",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
).set(
    body_background_fill="#1a0010",
    body_text_color="#fce4ec",
    block_background_fill="#2d0020",
    block_border_color="#ad1457",
    input_background_fill="#1a0010",
    button_primary_background_fill="#c2185b",
    button_primary_text_color="#fce4ec",
    button_primary_border_color="#e91e63",
    button_primary_background_fill_hover="#e91e63",
    button_primary_text_color_hover="#ffffff",
)

CSS = """
.result-box textarea { font-size:1.1rem !important; font-weight:bold !important; text-align:center; }
footer { display:none !important; }
"""

# ── Build UI ──────────────────────────────────────────────────────────────────
with gr.Blocks(title="DeepGuard Image") as demo:

    gr.HTML(HERO)

    with gr.Row():
        with gr.Column(scale=1, min_width=300):
            image_input = gr.Image(type="pil", label="Upload Image")
            with gr.Accordion("🖼️ Example Images — (click to expand)", open=False):
                gr.Examples(
                    examples=[
                        "example1.JPG",
                        "example2.JPG",
                        "example3.JPG",
                        "example4.JPG",
                    ],
                    inputs=image_input,
                    label="All samples are FAKE",
                )

            gr.Markdown("### ⚡ Model Variant")
            variant_radio = gr.Radio(
                choices=["fast", "pro"],
                value="fast",
                label=None,
                info="Fast (b0): 8.7M params  |  Pro (b5): 50.3M params",
            )

            gr.Markdown("### 🌍 Training Domain")
            dataset_radio = gr.Radio(
                choices=[
                    ("🇰🇷 KoDF — East Asian faces",      "kodf"),
                    ("🌟 Celeb-DF-v2 — Celebrity faces",  "celeb-df-v2"),
                    ("🌍 FF++ — Western faces",            "ff++"),
                ],
                value="ff++",
                label=None,
            )

            run_btn = gr.Button("🔍 Analyze", variant="primary", size="lg")

        with gr.Column(scale=1, min_width=300):
            verdict_output = gr.Textbox(
                label="Result",
                interactive=False,
                lines=2,
                elem_classes=["result-box"],
            )
            score_output = gr.Label(label="Class Probabilities", num_top_classes=2)

            gr.Markdown("""
---
**Model Guide**

| | ⚡ Fast (b0) | 🔥 Pro (b5) |
|---|---|---|
| **Params** | 8.7M | 50.3M |
| **Resolution** | 224×224 | 384×384 |
| **Speed** | Faster | Slower |
| **Accuracy** | High | Higher |

**Domain Guide**

| Domain | Training Data | Best For |
|---|---|---|
| KoDF | 62K real / 175K fake (Korean) | East Asian faces |
| Celeb-DF-v2 | 890 real / 5.6K fake (Celebrities) | Celebrity faces |
| FF++ | 1K real / 6K fake (4 manipulation types) | Western faces |
""")

    with gr.Accordion("⚙️ Advanced Settings", open=True):
        gr.Markdown("Adjust inference parameters. Defaults are recommended for most cases.")
        with gr.Row():
            with gr.Column():
                margin_ratio_slider = gr.Slider(
                    minimum=0.0, maximum=0.5, value=0.2, step=0.05,
                    label="Margin Ratio",
                    info="Margin around the detected face bounding box",
                )
                conf_thres_slider = gr.Slider(
                    minimum=0.1, maximum=0.9, value=0.5, step=0.05,
                    label="Face Detection Confidence",
                    info="YOLO face detection confidence threshold",
                )
            with gr.Column():
                min_face_ratio_slider = gr.Slider(
                    minimum=0.001, maximum=0.1, value=0.01, step=0.001,
                    label="Min Face Ratio",
                    info="Minimum face-to-frame area ratio to process",
                )
                tta_hflip_slider = gr.Slider(
                    minimum=0.0, maximum=1.0, value=0.0, step=0.1,
                    label="TTA Horizontal Flip",
                    info="⚠️ > 0.0 doubles inference time",
                )
                top_k_slider = gr.Slider(
                    minimum=1, maximum=2, value=2, step=1,
                    label="Top K",
                    info="Number of top predictions to return",
                )

    run_btn.click(
        fn=predict,
        inputs=[
            image_input, variant_radio, dataset_radio,
            margin_ratio_slider, conf_thres_slider,
            min_face_ratio_slider, tta_hflip_slider, top_k_slider,
        ],
        outputs=[score_output, verdict_output],
    )
    image_input.change(
        fn=lambda img, v, d, mr, ct, mfr, tta, tk: (
            predict(img, v, d, mr, ct, mfr, tta, tk) if img is not None else (None, "")
        ),
        inputs=[
            image_input, variant_radio, dataset_radio,
            margin_ratio_slider, conf_thres_slider,
            min_face_ratio_slider, tta_hflip_slider, top_k_slider,
        ],
        outputs=[score_output, verdict_output],
    )

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS)