import spaces  # MUST come before torch / any CUDA-touching import

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

MODEL_ID = "KDLAI/KDL-Frontier-Parser-nano"

processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
    trust_remote_code=True,
).to("cuda").eval()

# Task prompts per the model card (note the leading newline).
TASK_PROMPTS = {
    "Layout Detection": "\nLayout Detection:",
    "Text Recognition": "\nText Recognition:",
    "Table Recognition": "\nTable Recognition:",
    "Formula Recognition": "\nFormula Recognition:",
    "Image Analysis": "\nImage Analysis:",
}

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


@spaces.GPU(duration=120)
def parse_document(image: Image.Image, task: str, custom_prompt: str) -> str:
    """Run the KDL-Frontier-Parser-nano VLM on a document image.

    Args:
        image: a page image (PNG/JPEG) of the document to parse.
        task: one of Layout Detection, Text Recognition, Table Recognition,
            Formula Recognition, Image Analysis.
        custom_prompt: optional free-text prompt appended after the task
            header; leave empty for the canonical task prompt.

    Returns:
        The model's parsed text (OTSL for tables, etc.).
    """
    if image is None:
        return "Please upload a document image."
    if image.mode != "RGB":
        image = image.convert("RGB")

    if task == "Custom" and custom_prompt.strip():
        prompt_text = custom_prompt
    else:
        prompt_text = TASK_PROMPTS.get(task, "\nText Recognition:")
        if custom_prompt.strip():
            prompt_text = prompt_text + " " + custom_prompt.strip()

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": prompt_text},
            ],
        }
    ]

    chat_text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    inputs = processor(
        text=chat_text,
        images=image,
        return_tensors="pt",
        padding=True,
    ).to("cuda")

    t0 = time.perf_counter()
    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            do_sample=False,
            temperature=0.0,
            top_k=1,
            top_p=1.0,
            max_new_tokens=4096,
        )
    elapsed = time.perf_counter() - t0

    generated = output_ids[:, inputs["input_ids"].shape[1]:]
    text = processor.batch_decode(
        generated,
        skip_special_tokens=False,
    )[0]

    footer = f"\n\n---\n*Inference: {elapsed:.1f}s · task: {task}*"
    return text + footer


with gr.Blocks() as demo:
    gr.Markdown(
        "# KDL-Frontier-Parser-nano\n"
        "A 1.2B-parameter document-parsing VLM by [KoreaDeep](https://www.koreadeep.com/). "
        "Upload a page image, pick a task, and run. See the "
        "[model card](https://huggingface.co/KDLAI/KDL-Frontier-Parser-nano) for details."
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_in = gr.Image(type="pil", label="Document image", sources=["upload"])
            task = gr.Radio(
                choices=list(TASK_PROMPTS.keys()) + ["Custom"],
                value="Layout Detection",
                label="Task",
            )
            custom_prompt = gr.Textbox(
                label="Custom prompt (optional)",
                placeholder="Used only when Task = Custom; otherwise appended to the task header.",
                lines=2,
            )
            run_btn = gr.Button("Run", variant="primary")
        with gr.Column(scale=1):
            output = gr.Textbox(label="Parsed output", lines=18)

    with gr.Accordion("Tips", open=False):
        gr.Markdown(
            "- **Layout Detection** returns the page layout (regions + types).\n"
            "- **Text Recognition** extracts text per region.\n"
            "- **Table Recognition** returns OTSL.\n"
            "- **Formula Recognition** parses formulas.\n"
            "- **Image Analysis** describes figure content.\n"
            "- Feed page images, not PDFs. Greedy decoding (temperature=0) is used."
        )

    gr.Examples(
        examples=[
            ["examples/astronaut.jpg", "Layout Detection", ""],
            ["examples/astronaut.jpg", "Text Recognition", ""],
            ["examples/noodles.jpg", "Image Analysis", ""],
        ],
        inputs=[image_in, task, custom_prompt],
        outputs=output,
        fn=parse_document,
        cache_examples=True,
        cache_mode="lazy",
    )

    run_btn.click(
        fn=parse_document,
        inputs=[image_in, task, custom_prompt],
        outputs=output,
        api_name="parse",
    )

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