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

gr.Server backend + fully custom MRI-console frontend (frontend/index.html).
FLUX.2-klein-9B edits your photo into an X-ray radiograph; Qwen2.5-VL-3B
identifies the object so the scan knows what internals to render.

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

Set XRAY_MOCK=1 to run locally without GPU/models (replays cached scans).
"""

import base64
import io
import json
import os
import re
import time
from pathlib import Path

from fastapi.responses import FileResponse, HTMLResponse, Response
from gradio import Server
from PIL import Image

MOCK = os.getenv("XRAY_MOCK") == "1"
ROOT = Path(__file__).parent

if not MOCK:
    import spaces
    import torch
    from diffusers import Flux2KleinPipeline
    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)
else:
    class _Spaces:
        @staticmethod
        def GPU(*a, **k):
            def deco(f):
                return f
            return deco
    spaces = _Spaces()

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>
}"""


def vlm_chat(image, 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)


def b64_to_image(data_url: str) -> Image.Image:
    payload = data_url.split(",", 1)[-1]
    return Image.open(io.BytesIO(base64.b64decode(payload)))


def image_to_b64(image: Image.Image) -> str:
    buf = io.BytesIO()
    image.convert("RGB").save(buf, format="WEBP", quality=92)
    return "data:image/webp;base64," + base64.b64encode(buf.getvalue()).decode()


@spaces.GPU(duration=120)
def run_scan(image: Image.Image) -> dict:
    if MOCK:
        time.sleep(6)
        result = Image.open(ROOT / "examples/outputs/gpu_xray.webp")
        return {"object": "graphics card", "image": image_to_b64(result)}

    image = prep_image(image)
    info = parse_identify(vlm_chat(image, IDENTIFY_PROMPT, max_new_tokens=256))
    obj = info["object"]
    components = ", ".join(info["components"])

    edit_prompt = (
        f"X-ray radiograph scan of this exact {obj}, keeping its position "
        f"and silhouette. Dark black background, the outer casing rendered "
        f"as a translucent ghostly blue-white shell, and the internal "
        f"components clearly visible through it, glowing in bright "
        f"white-cyan: {components}. Airport security scanner / medical "
        f"X-ray imaging aesthetic, monochromatic blue-cyan palette, high "
        f"contrast, photorealistic radiograph, dense parts brighter than "
        f"hollow parts. No text, no words, no labels, no writing anywhere "
        f"in the image."
    )
    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]
    return {"object": obj, "image": image_to_b64(result)}


app = Server()


@app.api(name="xray")
def xray(image_b64: str) -> dict:
    """Scan a specimen: data-URL image in, {object, image (data-URL)} out."""
    image = b64_to_image(image_b64)
    return run_scan(image)


SAMPLES_DIR = (ROOT / "examples").resolve()


@app.get("/samples/{name}")
def sample(name: str):
    path = (SAMPLES_DIR / name).resolve()
    if not path.is_file() or SAMPLES_DIR not in path.parents:
        return Response(status_code=404)
    return FileResponse(path)


@app.get("/", response_class=HTMLResponse)
def index():
    return (ROOT / "frontend" / "index.html").read_text()


app.launch()