import os

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

import spaces  # noqa: E402  (must precede torch / CUDA-touching imports)
import torch  # noqa: E402
import gradio as gr  # noqa: E402
from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration  # noqa: E402

MODEL_ID = "ATH-MaaS/OvisOCR2"

# The exact page-parsing system prompt from the OvisOCR2 model card.
OCR_PROMPT = (
    "\nExtract all readable content from the image in natural human reading order "
    "and output the result as a single Markdown document. For charts or images, "
    'represent them using an HTML image tag: <img src="images/bbox_{left}_{top}_{right}_{bottom}.jpg" />, '
    "where left, top, right, bottom are bounding box coordinates scaled to [0, 1000). "
    "Format formulas as LaTeX. Format tables as HTML: <table>...</table>. "
    "Transcribe all other text as standard Markdown. Preserve the original text "
    "without translation or paraphrasing."
)

# Match the reference implementation's pixel budget.
MIN_PIXELS = 448 * 448
MAX_PIXELS = 2880 * 2880

# ----------------------------------------------------------------------------
# Load model + processor at module scope, eagerly on CUDA (ZeroGPU pattern).
# ----------------------------------------------------------------------------
processor = AutoProcessor.from_pretrained(
    MODEL_ID,
    min_pixels=MIN_PIXELS,
    max_pixels=MAX_PIXELS,
)
model = Qwen3_5ForConditionalGeneration.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()


def _clean_truncated_repeats(
    text: str,
    min_text_len: int = 8000,
    max_period: int = 200,
    min_period: int = 1,
    min_repeat_chars: int = 100,
    min_repeat_times: int = 5,
) -> str:
    """Remove truncated repetition artifacts from generated text (from model card)."""
    n = len(text)
    if n < min_text_len:
        return text

    max_period = min(max_period, n - 1)
    for unit_len in range(min_period, max_period + 1):
        if text[n - 1] != text[n - 1 - unit_len]:
            continue

        match_len = 1
        idx = n - 2
        while idx >= unit_len and text[idx] == text[idx - unit_len]:
            match_len += 1
            idx -= 1

        total_len = match_len + unit_len
        repeat_times = total_len // unit_len
        tail_len = total_len % unit_len

        if repeat_times >= min_repeat_times and total_len >= min_repeat_chars:
            return text[: n - total_len + unit_len] + text[n - tail_len:]

    return text


def _filter_img_tags(text: str) -> str:
    """Drop the placeholder <img src="images/bbox_..."> blocks (from model card)."""
    return "\n\n".join(
        block
        for block in text.split("\n\n")
        if not block.strip().startswith('<img src="images/bbox_')
    )


def _estimate_duration(image=None, filter_image_tags=True, max_new_tokens=4096, *args, **kwargs):
    # Pure-torch gated-delta-net path is ~30ms/token; scale the GPU reservation
    # with the token budget and keep it under the ZeroGPU per-call ceiling.
    return int(min(280, 45 + float(max_new_tokens) * 0.03))


@spaces.GPU(duration=_estimate_duration)
def parse_document(image, filter_image_tags: bool = True, max_new_tokens: int = 4096):
    """Parse a document image into Markdown with OvisOCR2.

    Args:
        image: A PIL document/page image to transcribe.
        filter_image_tags: Drop placeholder <img> bbox blocks the model emits for figures.
        max_new_tokens: Maximum number of tokens to generate for the transcription.

    Returns:
        The extracted document content as a Markdown string.
    """
    if image is None:
        return ""

    image = image.convert("RGB")

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

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
        enable_thinking=False,
    ).to(model.device)

    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=False,
            temperature=None,
            top_p=None,
            top_k=None,
        )

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

    if filter_image_tags:
        text = _filter_img_tags(text)
    text = _clean_truncated_repeats(text)
    return text


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

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 📄 OvisOCR2 — Document → Markdown

            End-to-end, page-level document parsing with
            [**ATH-MaaS/OvisOCR2**](https://huggingface.co/ATH-MaaS/OvisOCR2) — a compact
            0.9B vision-language model (Qwen3.5 backbone). Upload a document page and it
            transcribes the content into a single **Markdown** document, formatting
            formulas as LaTeX and tables as HTML.
            """
        )
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(type="pil", label="Document image", height=460)
                run = gr.Button("Extract Markdown", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    filter_tags = gr.Checkbox(
                        value=True,
                        label="Filter figure placeholder <img> tags",
                    )
                    max_tokens = gr.Slider(
                        512, 8192, value=4096, step=512,
                        label="Max new tokens",
                    )
            with gr.Column():
                md_out = gr.Code(
                    label="Extracted Markdown",
                    language="markdown",
                    lines=22,
                )
                with gr.Accordion("Rendered preview", open=False):
                    md_render = gr.Markdown()

        def _run(image, filter_image_tags, max_new_tokens):
            text = parse_document(image, filter_image_tags, max_new_tokens)
            return text, text

        def _run_example(image):
            text = parse_document(image, True, 4096)
            return text, text

        run.click(
            _run,
            inputs=[image_in, filter_tags, max_tokens],
            outputs=[md_out, md_render],
            api_name="parse",
        )

        gr.Examples(
            examples=[
                ["examples/academic_paper.jpg"],
                ["examples/math_textbook.png"],
                ["examples/magazine_page.png"],
            ],
            inputs=[image_in],
            outputs=[md_out, md_render],
            fn=_run_example,
            cache_examples=False,
            run_on_click=True,
        )

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