import os

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

import spaces  # must be imported before torch / CUDA-touching libraries
import gradio as gr
import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor

from otsl_to_html import convert_otsl_to_html

MODEL_ID = "openbmb/MiniCPM-OCR-1.0"

# Task -> (official prompt, default max_new_tokens, convert OTSL output to HTML).
# Prompts are taken verbatim from the MiniCPM-OCR 1.0 model card.
TASKS = {
    "Text": ("Text Recognition:", 4096, False),
    "Table": ("Table Recognition:", 8192, True),
    "Formula": ("Formula Recognition:", 2048, False),
    "Seal": ("Seal Recognition:", 2048, False),
    "Chart": ("Chart to Table:", 4096, True),
}
TASK_CHOICES = list(TASKS.keys())


def _normalize_latex_delimiters(text):
    """Rewrite \\(..\\) / \\[..\\] to $..$ / $$..$$ so KaTeX renders them.

    MiniCPM-OCR emits LaTeX with backslash delimiters; gr.Markdown renders
    math written with dollar delimiters. Same LaTeX, standard delimiters.
    """
    if not text:
        return text
    return (
        text.replace("\\[", "$$").replace("\\]", "$$")
        .replace("\\(", "$").replace("\\)", "$")
    )

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


def _gpu_duration(
    image=None, task="Text", custom_prompt="", max_new_tokens=4096, slice_nums=9,
    *args, **kwargs,  # Gradio injects its Progress object positionally
):
    """Reserve GPU time proportional to the generation budget.

    Measured on ZeroGPU: ~18-20 ms per generated token on the torch-reference
    linear-attention path, plus ~15 s for weight unpack + image slicing/prefill
    of a 9-slice A4 page. The budget is the only advance signal of output
    length, so scale with it and cap at 240 s.
    """
    return min(240, 15 + int(max_new_tokens) * 0.02)


@spaces.GPU(duration=_gpu_duration)
def run_ocr(
    image: Image.Image,
    task: str = "Text",
    custom_prompt: str = "",
    max_new_tokens: int = 4096,
    slice_nums: int = 9,
):
    """Recognize content in a document image with MiniCPM-OCR 1.0.

    Args:
        image: the document region or page image (PNG/JPG) to parse.
        task: recognition task — one of "Text", "Table", "Formula", "Seal", "Chart".
        custom_prompt: optional free-form instruction (e.g. key information
            extraction); when non-empty it replaces the task prompt.
        max_new_tokens: generation budget; raise it if the output looks truncated.
        slice_nums: max number of visual slices for large images (up to 9).

    Returns:
        (rendered_markdown, raw_model_output): parsed content — tables become
        HTML, formulas LaTeX — plus the raw model text.
    """
    if image is None:
        raise gr.Error("Please upload an image first.")
    if not isinstance(image, Image.Image):
        raise gr.Error("Image input could not be read.")

    if custom_prompt and custom_prompt.strip():
        prompt = custom_prompt.strip()
    else:
        prompt = TASKS[task][0]
    budget = int(max_new_tokens) if max_new_tokens else TASKS[task][1]

    conversation = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        enable_thinking=False,  # OCR mode: the template skips its reasoning block
        max_slice_nums=int(slice_nums),
    )
    inputs = {k: (v.to(model.device) if hasattr(v, "to") else v) for k, v in inputs.items()}

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=budget,
            do_sample=False,
        )

    # Strip the prompt tokens; only decode what the model generated.
    generated = output_ids[:, inputs["input_ids"].shape[1]:]
    text = processor.batch_decode(generated, skip_special_tokens=True)[0].strip()

    raw = text
    rendered = text
    wants_otsl = TASKS.get(task, ("", 0, False))[2] if not (custom_prompt and custom_prompt.strip()) else False
    if wants_otsl and ("<fcel>" in text or "<ecel>" in text):
        html_table = convert_otsl_to_html(text)
        if html_table:
            rendered = html_table
    if (
        task == "Formula"
        and rendered
        and "$" not in rendered
        and "\\(" not in rendered
        and "\\[" not in rendered
    ):
        rendered = f"$${rendered}$$"

    return _normalize_latex_delimiters(rendered), raw


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

THEME = gr.themes.Citrus()

with gr.Blocks(title="MiniCPM-OCR 1.0") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # MiniCPM-OCR 1.0 — Document Image → Text / Table / Formula

            OCR-specialist 1.3B vision-language model ([openbmb/MiniCPM-OCR-1.0](https://huggingface.co/openbmb/MiniCPM-OCR-1.0)):
            SigLIP vision encoder + Qwen3.5 language model. It recognizes text, tables, formulas,
            seals, and charts from document crops, scans, and photos, in Chinese and English.

            Upload a **document region or page** and pick the task — the prompt tells the model what to extract.
            Tables come back as HTML, formulas as LaTeX, plain text as text.
            """
        )

        with gr.Row():
            image_in = gr.Image(
                type="pil", label="Document image", height=340, sources=["upload", "clipboard"]
            )
            with gr.Column():
                task = gr.Radio(
                    TASK_CHOICES, value="Text", label="Task", info="What should be extracted from the image?"
                )
                run_btn = gr.Button("Run OCR", variant="primary")

        rendered_out = gr.Markdown(label="Result", value="Output will appear here.")
        with gr.Accordion("Raw model output", open=False):
            raw_out = gr.Textbox(
                label="Raw output (OTSL / LaTeX / JSON as emitted by the model)",
                lines=8,
                max_lines=20,
                interactive=False,
            )
        with gr.Accordion("Advanced settings", open=False):
            custom_prompt = gr.Textbox(
                value="",
                label="Custom prompt (overrides the task)",
                placeholder=(
                    'e.g. Extract the "store_name", "amount". Format the output as a JSON object '
                    'with keys: "store_name", "amount". If an item is not in the image, use "###".'
                ),
                lines=2,
            )
            max_new_tokens = gr.Slider(
                256, 8192, value=4096, step=256, label="Max new tokens",
                info="Raise if the output is truncated.",
            )
            slice_nums = gr.Slider(
                1, 9, value=9, step=1, label="Max visual slices",
                info="How many tiles a large image may be split into (9 = max detail).",
            )

        gr.Examples(
            examples=[
                ["examples/text_block.png", "Text"],
                ["examples/formula_block.png", "Formula"],
                ["examples/table_block.png", "Table"],
            ],
            inputs=[image_in, task],
            outputs=[rendered_out, raw_out],
            fn=run_ocr,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples",
        )
        gr.Markdown(
            """
            <sub>Model: [openbmb/MiniCPM-OCR-1.0](https://huggingface.co/openbmb/MiniCPM-OCR-1.0) (Apache-2.0).
            Example images are region crops from [OmniDocBench v1.6](https://huggingface.co/datasets/opendatalab/OmniDocBench),
            shipped with the [MiniCPM-OCR repo](https://github.com/OpenBMB/MiniCPM-OCR).
            Tables are emitted in OTSL and converted to HTML with the official SDK converter.</sub>
            """
        )

    run_btn.click(
        fn=run_ocr,
        inputs=[image_in, task, custom_prompt, max_new_tokens, slice_nums],
        outputs=[rendered_out, raw_out],
        api_name="ocr",
    )

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