"""
GenAI-DeepDetect — Gradio Space entry point.
Hardware: ZeroGPU (A10G, 40GB VRAM)
M1: SyncNet lip-sync  |  M2: CLIP fingerprint  |  M3: ViT temporal  |  M5: Gemini explainability
"""
import os
import time

import gradio as gr
import spaces  # HuggingFace ZeroGPU

from modules.m1_lipsync import LipSyncModule
from modules.m2_fingerprint import FingerprintModule
from modules.m3_fallback import M3FallbackModule   # swap → m3_sstgnn post L40S
from modules.m5_fusion import FusionModule
from modules.m5_explain import ExplainModule

CACHE = "/data/model_cache" if os.path.exists("/data") else "./cache"
os.makedirs(CACHE, exist_ok=True)

# All models load on CPU at startup — GPU not allocated yet
print("Loading M1 SyncNet…")
m1 = LipSyncModule(cache_dir=CACHE)
print("Loading M2 Fingerprint…")
m2 = FingerprintModule(cache_dir=CACHE)
print("Loading M3 ViT fallback…")
m3 = M3FallbackModule(cache_dir=CACHE)
m5_fusion = FusionModule(weights_path="weights/fusion_mlp.pt")
m5_explain = ExplainModule()
print("All modules ready. GPU allocated per-request via ZeroGPU.")


@spaces.GPU(duration=120)
def analyze(video_file):
    if video_file is None:
        return "⚠️ Please upload a video.", "", "", ""

    start = time.time()

    # Move to A10G for this request
    m1.to_gpu()
    m2.to_gpu()
    m3.to_gpu()

    try:
        r1 = m1.score(video_file)
        r2 = m2.score(video_file)
        r3 = m3.score(video_file)
    finally:
        m1.to_cpu()
        m2.to_cpu()
        m3.to_cpu()

    fusion = m5_fusion.fuse(r1["s1"], r2["s2"], r3["s3"])
    explanation = m5_explain.explain(
        fakescore=fusion["FakeScore"],
        s1=r1["s1"],
        s2=r2["s2"],
        s3=r3["s3"],
        weights=fusion["weights"],
        attribution=r2["attribution"],
        segments=r1.get("segments", []),
        top_generator=r2["top_generator"],
    )

    elapsed = time.time() - start
    verdict = "FAKE" if fusion["FakeScore"] > 0.5 else "REAL"
    icon = "🔴" if verdict == "FAKE" else "🟢"

    verdict_md = f"## {icon} {verdict}\n**FakeScore: {fusion['FakeScore']:.3f}**"

    scores_md = f"""### Per-Module Scores
| Module | Score | Weight |
|--------|-------|--------|
| 🎤 Lip-Sync (SyncNet) | `{r1['s1']:.3f}` | {fusion['weights']['lip_sync']:.2f} |
| 🖼️ Fingerprint (CLIP) | `{r2['s2']:.3f}` | {fusion['weights']['fingerprint']:.2f} |
| 🕸️ Temporal (ViT) | `{r3['s3']:.3f}` | {fusion['weights']['graph_gnn']:.2f} |

**⏱️ Time:** {elapsed:.1f}s  |  **💻 Hardware:** A10G (ZeroGPU)"""

    attr_md = "### Generator Attribution\n"
    if r2["attribution"]:
        for gen, prob in sorted(r2["attribution"].items(), key=lambda x: -x[1])[:5]:
            bar = "█" * int(prob * 25) + "░" * (25 - int(prob * 25))
            attr_md += f"- **{gen}**: {prob * 100:.1f}% `{bar}`\n"
        attr_md += f"\n**Top match:** {r2['top_generator']}"
    else:
        attr_md += "_Classified as real — attribution skipped._"

    # Lip-sync anomaly timestamps
    if r1.get("segments"):
        scores_md += "\n\n**⚠️ Desync segments:**\n"
        for seg in r1["segments"][:5]:
            scores_md += f"- t={seg['time']}s (score={seg['score']:.2f})\n"

    return verdict_md, scores_md, attr_md, explanation


# ── UI ────────────────────────────────────────────────────────────────────────

with gr.Blocks(
    title="GenAI-DeepDetect",
    theme=gr.themes.Base(
        primary_hue="red",
        font=["DM Sans", "ui-sans-serif", "sans-serif"],
    ),
    css="""
    .verdict-box { border-radius: 12px; padding: 16px; }
    footer { display: none !important; }
    """,
) as demo:
    gr.Markdown(
        """# 🔍 GenAI-DeepDetect
### Multimodal Deepfake Detection & Attribution
**Modules:** SyncNet (lip-sync) · CLIP (fingerprint) · ViT (temporal) · Gemini explainability
**Hardware:** ZeroGPU A10G (40GB) · **Paper:** SRM IST 2026"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            vid = gr.Video(label="Upload Video", height=280)
            btn = gr.Button("🔍 Analyze", variant="primary", size="lg")
            if os.path.exists("test_assets/real_sample.mp4"):
                gr.Examples(
                    examples=[["test_assets/real_sample.mp4"], ["test_assets/fake_sample.mp4"]],
                    inputs=[vid],
                    label="Try sample videos",
                )

        with gr.Column(scale=2):
            verdict_out = gr.Markdown(label="Verdict", elem_classes=["verdict-box"])
            scores_out = gr.Markdown(label="Module Scores")

    with gr.Row():
        attr_out = gr.Markdown(label="Generator Attribution")
        expl_out = gr.Markdown(label="AI Forensic Explanation")

    btn.click(
        fn=analyze,
        inputs=[vid],
        outputs=[verdict_out, scores_out, attr_out, expl_out],
    )

    gr.Markdown(
        "---\n*GenAI-DeepDetect · Akshat Agarwal, Dev Chopda · SRM IST · "
        "[GitHub](https://github.com/akagtag/genai-deepdetect)*"
    )

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