"""X-Ray Anything — photograph any object and see inside it.

A kid points a camera at a toaster and gets a cutaway illustration of its
insides (FLUX.2-klein-9B editing their actual photo) plus an age-tuned
explanation of how it works (Qwen2.5-VL-3B).

Build Small Hackathon 2026 — models: 9B image + 3B VLM, both <= 32B.
"""

import json
import re

import gradio as gr
import spaces
import torch
from diffusers import Flux2KleinPipeline
from PIL import Image
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

DEVICE = "cuda"
DTYPE = torch.bfloat16

FLUX_ID = "black-forest-labs/FLUX.2-klein-9B"
VLM_ID = "Qwen/Qwen2.5-VL-3B-Instruct"

flux_pipe = Flux2KleinPipeline.from_pretrained(FLUX_ID, torch_dtype=DTYPE)
flux_pipe.to(DEVICE)

vlm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    VLM_ID, torch_dtype=DTYPE, device_map=DEVICE
)
vlm_processor = AutoProcessor.from_pretrained(VLM_ID)

IDENTIFY_PROMPT = """Look at this photo. Identify the single main object in it.

Reply with ONLY a JSON object, no other text:
{
  "object": "<short name of the object>",
  "components": ["<4-6 main internal parts/components inside this object>"],
  "is_mechanical_or_electrical": <true/false>
}"""

EXPLAIN_PROMPT = """You are a warm, funny science explainer for a {age}-year-old child.
The child just photographed a {object} and is looking at a cutaway illustration
showing its insides: {components}.

Write for a {age}-year-old:
1. One exciting opening line about what's hiding inside.
2. A short, simple explanation of how it works (3-5 sentences), walking through
   the parts above in the order energy/material flows through them.
3. Three fun facts, each starting with an emoji.

Use simple words appropriate for age {age}. Be accurate — no made-up science.
Do not use headings or the word 'cutaway'."""


def vlm_chat(image: Image.Image | None, prompt: str, max_new_tokens: int = 512) -> str:
    content = []
    if image is not None:
        content.append({"type": "image", "image": image})
    content.append({"type": "text", "text": prompt})
    messages = [{"role": "user", "content": content}]
    text = vlm_processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = vlm_processor(
        text=[text],
        images=[image] if image is not None else None,
        return_tensors="pt",
    ).to(DEVICE)
    with torch.inference_mode():
        out = vlm.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    out = out[:, inputs.input_ids.shape[1]:]
    return vlm_processor.batch_decode(out, skip_special_tokens=True)[0].strip()


def parse_identify(raw: str) -> dict:
    match = re.search(r"\{.*\}", raw, re.DOTALL)
    if match:
        try:
            data = json.loads(match.group(0))
            if data.get("object") and data.get("components"):
                return data
        except json.JSONDecodeError:
            pass
    return {"object": "object", "components": ["inner workings"],
            "is_mechanical_or_electrical": True}


def prep_image(image: Image.Image, max_side: int = 1024) -> Image.Image:
    image = image.convert("RGB")
    w, h = image.size
    scale = max_side / max(w, h)
    if scale < 1:
        image = image.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
    # klein wants multiples of 16
    w, h = image.size
    return image.resize((w - w % 16, h - h % 16), Image.LANCZOS)


@spaces.GPU(duration=120)
def xray(image: Image.Image, age: int, progress=gr.Progress()):
    if image is None:
        raise gr.Error("Take or upload a photo first!")
    image = prep_image(image)

    progress(0.1, desc="🔍 Figuring out what this is...")
    info = parse_identify(vlm_chat(image, IDENTIFY_PROMPT, max_new_tokens=256))
    obj = info["object"]
    components = ", ".join(info["components"])

    progress(0.35, desc=f"🩻 X-raying the {obj}...")
    edit_prompt = (
        f"Educational cutaway illustration of this exact {obj}, keeping its "
        f"position, colors and background, but with the outer casing partially "
        f"cut away to reveal the internal components inside: {components}. "
        f"Children's science encyclopedia style, clearly visible internal parts "
        f"with thin white label lines and small text labels naming each part, "
        f"bright friendly colors, crisp detailed technical illustration."
    )
    with torch.inference_mode():
        result = flux_pipe(
            image=image,
            prompt=edit_prompt,
            height=image.height,
            width=image.width,
            guidance_scale=1.0,
            num_inference_steps=4,
            generator=torch.Generator(device=DEVICE).manual_seed(42),
        ).images[0]

    progress(0.75, desc="📖 Writing your explanation...")
    explanation = vlm_chat(
        None,
        EXPLAIN_PROMPT.format(age=age, object=obj, components=components),
        max_new_tokens=600,
    )

    header = f"## 🩻 Inside your {obj}!\n\n"
    footer = (
        "\n\n---\n*🎨 This is an artist's imagination of the inside, made by a "
        "small AI — real engineers' diagrams may differ!*"
    )
    return result, header + explanation + footer


with gr.Blocks(title="X-Ray Anything", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🩻 X-Ray Anything
        ### Photograph any object — see what's hiding inside it, and learn how it works!
        Built with **FLUX.2-klein-9B** (image x-raying) and **Qwen2.5-VL-3B** (eyes & explanations).
        """
    )
    with gr.Row():
        with gr.Column():
            image_in = gr.Image(
                sources=["webcam", "upload"], type="pil", label="📷 Your object"
            )
            age = gr.Slider(4, 14, value=8, step=1, label="🎂 Explain it for age...")
            btn = gr.Button("🩻 X-Ray it!", variant="primary", size="lg")
        with gr.Column():
            image_out = gr.Image(label="🔬 Inside view", interactive=False)
    explanation_out = gr.Markdown()

    btn.click(xray, inputs=[image_in, age], outputs=[image_out, explanation_out])

    gr.Markdown(
        "*Build Small Hackathon 2026 · everything under 32B parameters · "
        "[Barath](https://huggingface.co/Barath)*"
    )

demo.launch()