"""Hugging Face Space app for the ai-image-detector model.

This Space loads the merged CLIP ViT-B/16 LoRA model from this repository
(`model.safetensors`) and exposes a Gradio interface for uploading an image
(or pointing at a directory) and getting a real/fake/uncertain prediction.
"""
from __future__ import annotations

import argparse
import csv
import json
import tempfile
from pathlib import Path
from typing import Any

import gradio as gr
import torch
from PIL import Image
from torch.amp import autocast

# NOTE: This Space uses the standalone model loader (_load_model_from_space)
# and does not depend on the source training package. All config is read from
# config.json at runtime.

IMAGE_EXTENSIONS = {".bmp", ".jpeg", ".jpg", ".png", ".tif", ".tiff", ".webp"}
BATCH_HEADERS = ["file", "prediction", "confidence", "p_real", "p_fake", "fake_threshold", "real_threshold", "error"]
APP_CSS = """
.gradio-container { max-width: 1040px !important; }
.detector-header { align-items: center; display: flex; justify-content: space-between; margin-bottom: 16px; }
.detector-title { font-size: 26px; font-weight: 700; line-height: 1.15; }
.detector-meta { color: var(--body-text-color-subdued); font-size: 13px; text-align: right; }
.result-card { background: var(--background-fill-secondary); border-radius: 8px; padding: 14px 16px; }
.result-heading { display: flex; gap: 12px; justify-content: space-between; margin-bottom: 10px; }
.result-label { font-size: 24px; font-weight: 800; line-height: 1.1; }
.result-pill { border-radius: 999px; color: white; font-size: 12px; font-weight: 800; height: fit-content; padding: 5px 10px; }
.result-real { background: #15803d; }
.result-fake { background: #b91c1c; }
.result-uncertain { background: #b45309; }
.metric-grid { display: grid; gap: 12px; grid-template-columns: repeat(auto-fit, minmax(120px, 1fr)); margin-top: 8px; }
.metric { display: flex; flex-direction: column; gap: 2px; }
.metric-name { color: var(--body-text-color-subdued); font-size: 12px; }
.metric-value { font-size: 18px; font-weight: 700; }
"""


def _load_model_from_space():
    """Load the merged weights from model.safetensors in this Space/repo."""
    import timm
    from safetensors.torch import load_file

    weights_path = Path(__file__).resolve().parent / "model.safetensors"
    config_path = Path(__file__).resolve().parent / "config.json"
    cfg = json.loads(config_path.read_text())

    model = timm.create_model(cfg["backbone"], pretrained=False, num_classes=1, img_size=cfg["image_size"])
    state = load_file(str(weights_path))
    missing, unexpected = model.load_state_dict(state, strict=False)
    model.eval()
    return model, cfg


class SpacePredictor:
    """Lightweight predictor that loads the merged model directly from safetensors."""

    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model, self.cfg = _load_model_from_space()
        self.model = self.model.to(self.device)
        self.real_threshold = float(self.cfg.get("real_threshold", 0.93))
        self.fake_threshold = float(self.cfg.get("fake_threshold", 0.91))
        self.temperature = float(self.cfg.get("temperature", 1.0))
        self.img_size = int(self.cfg.get("image_size", 256))
        mean = self.cfg.get("normalization_mean", [0.481, 0.458, 0.408])
        std = self.cfg.get("normalization_std", [0.269, 0.261, 0.276])
        from torchvision import transforms
        self.transform = transforms.Compose([
            transforms.Resize((self.img_size, self.img_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std),
        ])

    @torch.inference_mode()
    def predict(self, image: Image.Image) -> dict[str, Any]:
        if image is None:
            return {}
        image = image.convert("RGB")
        x = self.transform(image).unsqueeze(0).to(self.device)
        logit = self.model(x).reshape(-1)
        if self.temperature and self.temperature > 0:
            logit = logit / self.temperature
        p_real = float(torch.sigmoid(logit).item())
        if p_real < self.fake_threshold:
            prediction, confidence = "fake", 1.0 - p_real
        elif p_real >= self.real_threshold:
            prediction, confidence = "real", p_real
        else:
            prediction, confidence = "uncertain", max(p_real - self.fake_threshold, self.real_threshold - p_real)
        return {
            "prediction": prediction,
            "confidence": confidence,
            "real_probability": p_real,
            "fake_probability": 1.0 - p_real,
            "fake_threshold": self.fake_threshold,
            "real_threshold": self.real_threshold,
            "img_size": self.img_size,
            "score_semantics": "sigmoid output is p(real); fake=0, real=1",
            "temperature": self.temperature,
            "decision_rule": "fake if p(real) < fake_threshold; real if p(real) >= real_threshold; otherwise uncertain",
        }


predictor = SpacePredictor()


def _result_card(m: dict[str, Any]) -> str:
    prediction = m.get("prediction", "unknown")
    pill = {"real": "result-real", "fake": "result-fake"}.get(prediction, "result-uncertain")
    return f"""
<div class="result-card">
  <div class="result-heading">
    <div>
      <div class="metric-name">Prediction</div>
      <div class="result-label">{prediction.upper()}</div>
    </div>
    <div class="result-pill {pill}">{prediction.upper()}</div>
  </div>
  <div class="metric-grid">
    <div class="metric"><div class="metric-name">Confidence</div><div class="metric-value">{m.get('confidence', 0):.4f}</div></div>
    <div class="metric"><div class="metric-name">p(real)</div><div class="metric-value">{m.get('real_probability', 0):.4f}</div></div>
    <div class="metric"><div class="metric-name">p(fake)</div><div class="metric-value">{m.get('fake_probability', 0):.4f}</div></div>
    <div class="metric"><div class="metric-name">Fake threshold</div><div class="metric-value">{m.get('fake_threshold', 0)}</div></div>
    <div class="metric"><div class="metric-name">Real threshold</div><div class="metric-value">{m.get('real_threshold', 0)}</div></div>
  </div>
</div>
"""


def predict_single(image: Image.Image | None):
    if image is None:
        return '<div class="result-card">Please upload an image.</div>', {}
    m = predictor.predict(image)
    return _result_card(m), m


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="AI Image Detector", css=APP_CSS) as demo:
        gr.HTML(
            f"""
<div class="detector-header">
  <div class="detector-title">AI Image Detector</div>
  <div class="detector-meta">threshold {predictor.real_threshold:.2f}<br>CLIP ViT-B/16 + LoRA</div>
</div>
"""
        )
        with gr.Tab("Single image"):
            with gr.Row():
                with gr.Column():
                    img_in = gr.Image(type="pil", label="Image")
                    btn = gr.Button("Predict")
                with gr.Column():
                    out_card = gr.HTML(label="Result")
                    raw = gr.Json(label="Raw scores")
            btn.click(predict_single, inputs=[img_in], outputs=[out_card, raw])
            gr.Examples(
                examples=[[str(p)] for p in sorted((Path(__file__).resolve().parent / "examples").glob("*.jpg"))],
                inputs=[img_in],
                outputs=[out_card, raw],
                fn=predict_single,
                cache_examples=False,
            )
    return demo


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