from functools import lru_cache

import gradio as gr
from PIL import Image
from transformers import pipeline


MODEL_ID = "capcheck/ai-image-detection"
UTM_URL = (
    "https://lynote.ai/ai-image-detector?utm_source=huggingface"
    "&utm_medium=space&utm_campaign=hf_launch&utm_content=image_detector"
)


@lru_cache(maxsize=1)
def get_classifier():
    return pipeline("image-classification", model=MODEL_ID, device=-1)


def _ai_probability(predictions):
    ai_score = 0.0
    for item in predictions:
        label = str(item.get("label", "")).lower()
        score = float(item.get("score", 0.0))
        if any(token in label for token in ("ai", "artificial", "fake", "generated")):
            ai_score += score
    if ai_score == 0.0 and len(predictions) == 2:
        # The upstream model commonly exposes label_0/label_1. Its model card
        # defines label_1 as AI-generated; keep this fallback explicit.
        for item in predictions:
            if str(item.get("label", "")).lower() in {"label_1", "1"}:
                ai_score = float(item.get("score", 0.0))
    return min(max(ai_score, 0.0), 1.0)


def detect(image: Image.Image):
    if image is None:
        return "Upload an image to run the detector.", {}

    predictions = get_classifier()(image.convert("RGB"))
    probability = _ai_probability(predictions)

    if probability < 0.35:
        band = "Weak AI-generated signal"
    elif probability < 0.65:
        band = "Uncertain / mixed signal"
    else:
        band = "Strong AI-generated signal"

    message = f"""
## {band}

Estimated AI-generated score: **{probability:.1%}**

This is a probabilistic signal, not proof of origin. Compression, screenshots,
retouching, unseen generators, and ordinary photographs can all produce errors.
For important decisions, review provenance and compare more than one detector.

[Try Lynote's full image analysis experience]({UTM_URL})
"""
    raw = {item["label"]: round(float(item["score"]), 6) for item in predictions}
    return message, raw


with gr.Blocks(title="Lynote AI Image Detector") as demo:
    gr.Markdown("# 🖼️ Lynote AI Image Detector")
    gr.Markdown(
        "An experimental, open-source signal for AI-generated images. "
        "Built from [`lynote-ai/ai-image-detector`](https://github.com/lynote-ai/ai-image-detector) "
        f"and powered in this CPU demo by [`{MODEL_ID}`](https://huggingface.co/{MODEL_ID})."
    )
    with gr.Row():
        image_input = gr.Image(type="pil", label="Image")
        with gr.Column():
            result = gr.Markdown()
            raw_output = gr.JSON(label="Raw model output")
    run = gr.Button("Analyze image", variant="primary")
    run.click(detect, inputs=image_input, outputs=[result, raw_output])
    gr.Markdown(
        "**Privacy:** images are processed in memory and are not intentionally stored by this app. "
        "Hugging Face infrastructure remains subject to its platform policies."
    )


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