import gradio as gr
import torch
from PIL import Image
from typing import Dict, Optional, Tuple
from transformers import AutoImageProcessor, SiglipForImageClassification

from trufor_runner import TruForEngine, TruForResult, TruForUnavailableError

MODEL_ID = "Ateeqq/ai-vs-human-image-detector"

# Use GPU when available so large batches stay responsive.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

try:
    processor = AutoImageProcessor.from_pretrained(MODEL_ID)
    model = SiglipForImageClassification.from_pretrained(MODEL_ID)
    model.to(device)
    model.eval()
except Exception as exc:  # pragma: no cover - surface loading issues early.
    raise RuntimeError(f"Failed to load model from {MODEL_ID}") from exc

try:
    TRUFOR_ENGINE: Optional[TruForEngine] = TruForEngine(device="cpu")
    TRUFOR_STATUS = TRUFOR_ENGINE.status_message
except TruForUnavailableError as exc:
    TRUFOR_ENGINE = None
    TRUFOR_STATUS = str(exc)


def analyze_ai_vs_human(image: Image.Image) -> Tuple[Dict[str, float], str]:
    """Run the Hugging Face detector and return confidences with a readable summary."""
    if image is None:
        empty_scores = {label: 0.0 for label in model.config.id2label.values()}
        return empty_scores, "No image provided."

    image = image.convert("RGB")
    inputs = processor(images=image, return_tensors="pt").to(device)

    with torch.no_grad():
        logits = model(**inputs).logits

    probabilities = torch.softmax(logits, dim=-1)[0]
    scores = {
        model.config.id2label[idx]: float(probabilities[idx])
        for idx in range(probabilities.size(0))
    }

    top_idx = int(probabilities.argmax().item())
    top_label = model.config.id2label[top_idx]
    top_score = scores[top_label]
    summary = f"**Predicted Label:** {top_label}  \
**Confidence:** {top_score:.4f}"

    return scores, summary


def analyze_trufor(image: Image.Image) -> Tuple[str, Optional[Image.Image], Optional[Image.Image]]:
    """Run TruFor inference when available, otherwise return diagnostics."""
    if TRUFOR_ENGINE is None:
        return TRUFOR_STATUS, None, None

    if image is None:
        return "Upload an image to run TruFor.", None, None

    try:
        result: TruForResult = TRUFOR_ENGINE.infer(image)
    except TruForUnavailableError as exc:
        return str(exc), None, None

    summary_lines = []
    if result.score is not None:
        summary_lines.append(f"**Tamper Score:** {result.score:.4f}")
    extras_dict = result.raw_scores.copy()
    if result.score is not None:
        extras_dict.pop("tamper_score", None)
    if extras_dict:
        extras = "  ".join(f"{key}: {value:.4f}" for key, value in extras_dict.items())
        summary_lines.append(f"`{extras}`")
    if not summary_lines:
        summary_lines.append("TruFor returned no scores for this image.")

    return "\n".join(summary_lines), result.map_overlay, result.confidence_overlay


def analyze_image(image: Image.Image) -> Tuple[Dict[str, float], str, str, Optional[Image.Image], Optional[Image.Image]]:
    ai_scores, ai_summary = analyze_ai_vs_human(image)
    trufor_summary, tamper_overlay, conf_overlay = analyze_trufor(image)
    return ai_scores, ai_summary, trufor_summary, tamper_overlay, conf_overlay


with gr.Blocks() as demo:
    gr.Markdown(
        """# Image Authenticity Workbench\nUpload an image to compare the AI-vs-human classifier with the TruFor forgery detector."""
    )

    status_box = gr.Markdown(f"`{TRUFOR_STATUS}`")

    image_input = gr.Image(label="Input Image", type="pil")
    analyze_button = gr.Button("Analyze", variant="primary", size="sm")

    with gr.Tabs():
        with gr.TabItem("AI vs Human"):
            ai_label_output = gr.Label(label="Prediction", num_top_classes=2)
            ai_summary_output = gr.Markdown("Upload an image to view the prediction.")
        with gr.TabItem("TruFor Forgery Detection"):
            trufor_summary_output = gr.Markdown("Configure TruFor assets to enable tamper analysis.")
            tamper_overlay_output = gr.Image(label="Tamper Heatmap", type="pil", interactive=False)
            conf_overlay_output = gr.Image(label="Confidence Heatmap", type="pil", interactive=False)

    output_components = [
        ai_label_output,
        ai_summary_output,
        trufor_summary_output,
        tamper_overlay_output,
        conf_overlay_output,
    ]

    analyze_button.click(
        fn=analyze_image,
        inputs=image_input,
        outputs=output_components,
    )

    image_input.change(
        fn=analyze_image,
        inputs=image_input,
        outputs=output_components,
    )


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