import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  must precede torch / transformers

import json
import re
import html
import tempfile

import torch
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoProcessor, AutoModelForImageTextToText

MODEL_ID = "P1n3/sdg-detector-grpo"

# ---------------------------------------------------------------------------
# Model (loaded once at module scope, moved to CUDA eagerly for ZeroGPU)
# ---------------------------------------------------------------------------
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()

# ---------------------------------------------------------------------------
# Prompt (the SDG detector's structured-grounding question template)
# ---------------------------------------------------------------------------
QUESTION_TEMPLATE = """You are an AI image quality evaluator. You will be given **one image** to analyze.

### Definitions

**Misalignment**: Areas where the image content does NOT match the text caption, including:
- Missing objects: Objects mentioned in caption but not present in image
- Extra objects: Objects present in image but not mentioned in caption
- Wrong attributes: Incorrect color, size, material, count, or other properties
- Wrong spatial relationships: Incorrect positions, orientations, or arrangements

**Artifact**: Visual defects in the generated image, including:
- Distorted anatomy: Malformed hands, extra/missing limbs, wrong number of fingers
- Duplicated/missing parts: Repeated or absent body parts, objects
- Warped geometry: Perspective errors, impossible shapes
- Texture issues: Melted, smeared, or overly smooth textures
- Unnatural edges: Jagged, broken, or blurry boundaries
- Garbled text: Unreadable or malformed text/letters
- Lighting inconsistencies: Wrong shadows, reflections, or light sources

Text Caption: {caption}

**Goal**: Produce a detailed analysis of the image quality and output bounding boxes for all detected issues.

### Strict Output Rules
Output **TWO blocks in this exact order**:
1) `<think>` - Your detailed analysis
2) `<answer>` - JSON list of bounding boxes

### Answer Format (for <answer>)
Return a JSON list:
[
    {{"box_2d": [x0, y0, x1, y1], "label": "misalignment"|"artifact", "description": "brief description of the issue", "importance": 1-100}}
]

Bounding box coordinates are in normalized 0-1000 space: [x0, y0, x1, y1].
The "importance" is an integer from 1 (minor) to 100 (severe) rating how much the defect hurts image quality.
If there are no issues, output an empty list.

Now analyze the image and produce your output:
"""

ARTIFACT_COLOR = (239, 68, 68)       # red
MISALIGNMENT_COLOR = (37, 99, 235)   # blue


def _load_font(size):
    for path in (
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    ):
        try:
            return ImageFont.truetype(path, size)
        except Exception:
            continue
    return ImageFont.load_default()


def parse_answer(response: str):
    """Extract the structured defect list from the model's <answer> block."""
    m = re.search(r"<answer>\s*(.*?)\s*</answer>", response, re.DOTALL)
    payload = m.group(1) if m else response
    start, end = payload.find("["), payload.rfind("]")
    if start == -1 or end == -1 or end <= start:
        return []
    try:
        data = json.loads(payload[start:end + 1])
    except Exception:
        return []
    if not isinstance(data, list):
        return []

    defects = []
    for item in data:
        if not isinstance(item, dict):
            continue
        box = item.get("box_2d")
        if not (isinstance(box, list) and len(box) == 4):
            continue
        try:
            box = [float(v) for v in box]
        except Exception:
            continue
        label = str(item.get("label", "")).lower()
        if label not in ("artifact", "misalignment"):
            label = "artifact"
        desc = item.get("description") or item.get("desc") or ""
        imp = item.get("importance")
        try:
            imp = int(imp)
        except Exception:
            imp = None
        defects.append({
            "box_2d": box,
            "label": label,
            "description": str(desc),
            "importance": imp,
        })
    return defects


def extract_think(response: str) -> str:
    m = re.search(r"<think>\s*(.*?)\s*</think>", response, re.DOTALL)
    return m.group(1).strip() if m else ""


def draw_defects(image: Image.Image, defects):
    """Overlay defect bounding boxes (coords normalized 0-1000, Qwen order)."""
    image = image.convert("RGB").copy()
    w, h = image.size
    draw = ImageDraw.Draw(image)
    line_w = max(3, round(min(w, h) / 250))
    font = _load_font(max(16, round(min(w, h) / 45)))

    for idx, d in enumerate(defects, start=1):
        x0, y0, x1, y1 = d["box_2d"]
        px0, py0 = x0 / 1000 * w, y0 / 1000 * h
        px1, py1 = x1 / 1000 * w, y1 / 1000 * h
        if px1 < px0:
            px0, px1 = px1, px0
        if py1 < py0:
            py0, py1 = py1, py0
        color = ARTIFACT_COLOR if d["label"] == "artifact" else MISALIGNMENT_COLOR
        draw.rectangle([px0, py0, px1, py1], outline=color, width=line_w)

        imp = d["importance"]
        tag = f"{idx}" + (f" \u00b7 {imp}" if imp is not None else "")
        bbox = draw.textbbox((0, 0), tag, font=font)
        tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
        ly = max(0, py0 - th - 4)
        draw.rectangle([px0, ly, px0 + tw + 8, ly + th + 4], fill=color)
        draw.text((px0 + 4, ly + 2), tag, fill=(255, 255, 255), font=font)

    return image


def defects_to_markdown(defects):
    if not defects:
        return ("### ✅ No defects detected\n\n"
                "The SDG detector did not ground any localized artifacts or "
                "caption misalignments in this image.")
    n_art = sum(1 for d in defects if d["label"] == "artifact")
    n_mis = sum(1 for d in defects if d["label"] == "misalignment")
    lines = [
        f"### Found {len(defects)} defect(s) — "
        f"🔴 {n_art} artifact, 🔵 {n_mis} misalignment\n",
        "| # | Type | Importance | Where / What / Why |",
        "|---|------|-----------|--------------------|",
    ]
    for idx, d in enumerate(defects, start=1):
        emoji = "🔴" if d["label"] == "artifact" else "🔵"
        imp = d["importance"] if d["importance"] is not None else "—"
        box = ", ".join(str(int(v)) for v in d["box_2d"])
        desc = html.escape(d["description"]).replace("\n", " ").replace("|", "\\|")
        lines.append(
            f"| {idx} | {emoji} {d['label']} | {imp} | "
            f"{desc} <br><sub>box (0–1000): [{box}]</sub> |"
        )
    return "\n".join(lines)


@spaces.GPU(duration=90)
def detect(image, caption, max_new_tokens=1024):
    """Detect and ground structured defects in a text-to-image generation.

    Args:
        image: the generated image to inspect for defects.
        caption: the text prompt the image was generated from (used to spot
            caption/image misalignments). Optional.
        max_new_tokens: generation budget for the detector's reasoning + answer.

    Returns:
        The image annotated with defect boxes, a structured defect table, and
        the detector's raw reasoning.
    """
    if image is None:
        raise gr.Error("Please provide an image to analyze.")
    caption = (caption or "").strip()

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": QUESTION_TEMPLATE.format(caption=caption)},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
        )
    trimmed = generated[:, inputs["input_ids"].shape[1]:]
    response = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]

    defects = parse_answer(response)
    annotated = draw_defects(image, defects)
    table = defects_to_markdown(defects)
    think = extract_think(response)
    reasoning = think if think else response.strip()
    return annotated, table, reasoning


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

INTRO = """# 🔍 Structured Defect Grounding (SDG)

Detect **localized defects** in text-to-image generations and get instance-level
feedback — *where* the defect is, *what* type it is (🔴 artifact / 🔵 misalignment),
*why* it's wrong, and *how important* it is (1–100).

Powered by the **SDG Detector** ([`P1n3/sdg-detector-grpo`](https://huggingface.co/P1n3/sdg-detector-grpo)),
a Qwen3-VL-4B model trained with SFT + GRPO from the paper
[*Where, What, Why, and Importance: Structured Defect Grounding for Text-to-Image Feedback*](https://huggingface.co/papers/2606.06113).

Upload a generated image and (optionally) the caption it was generated from,
then press **Analyze**.
"""

with gr.Blocks(title="Structured Defect Grounding") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)
        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(type="pil", label="Generated image", height=380)
                caption_in = gr.Textbox(
                    label="Caption / prompt (optional)",
                    placeholder="The text prompt the image was generated from…",
                    lines=2,
                )
                run = gr.Button("Analyze", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    max_tokens = gr.Slider(
                        256, 2048, value=1024, step=64,
                        label="Max new tokens",
                        info="Generation budget for the detector's reasoning + answer.",
                    )
            with gr.Column(scale=1):
                image_out = gr.Image(type="pil", label="Detected defects", height=380)
                table_out = gr.Markdown(label="Structured feedback")
        with gr.Accordion("Detector reasoning (<think>)", open=False):
            reasoning_out = gr.Textbox(label="Raw reasoning", lines=8)

        gr.Examples(
            examples=[
                ["examples/throne_dystopia.png", "a jung male sitting down on a throne in a dystopian world, digital art, epic"],
                ["examples/sign_bianca_buda.png", "A sign that says Bianca Buda"],
                ["examples/pikachu_mario.png", "Pikachu Ninja turtle Mewtwo super Mario"],
                ["examples/car_fruit_stand.png", "a car driving through a fruit stand, movie action scene, fruits flying everywhere"],
            ],
            inputs=[image_in, caption_in],
            outputs=[image_out, table_out, reasoning_out],
            fn=detect,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        detect,
        inputs=[image_in, caption_in, max_tokens],
        outputs=[image_out, table_out, reasoning_out],
        api_name="detect",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)