"""Gradio demo for ATH-MaaS/OvisOCR3 — a unified 0.8B OCR / document-understanding
vision-language model covering 5 capability domains selected via task prompts:
document parsing, general OCR, key information extraction, figure captioning,
and cross-page table reconstruction.

The inference path follows the model card exactly (task prompts from the repo's
`prompts/` directory, literal `.replace()` substitution, greedy decoding,
`enable_thinking=False`, 448²–2880² pixel budget) and is served through
transformers' `Qwen3_5ForConditionalGeneration` on ZeroGPU.
"""

import json
import os
import re
from pathlib import Path

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

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

MODEL_ID = "ATH-MaaS/OvisOCR3"
PROMPTS_DIR = Path(__file__).parent / "prompts"

# Pixel budget from the model card's `images_kwargs`.
MIN_PIXELS = 448 * 448
MAX_PIXELS = 2880 * 2880

MAX_NEW_TOKENS_CEILING = 16384

# ---------------------------------------------------------------------------
# Model + processor, loaded eagerly at module scope 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()


# ---------------------------------------------------------------------------
# Prompt loading — mirrors the model card's `load_prompt`.
# ---------------------------------------------------------------------------
def load_prompt(name: str) -> str:
    """Read a task prompt shipped with the Space from `prompts/<name>.txt`.

    Args:
        name: Prompt file stem, e.g. "document_markdown".

    Returns:
        The raw prompt text exactly as shipped in the model repo.
    """
    return (PROMPTS_DIR / f"{name}.txt").read_text(encoding="utf-8")


def _general_ocr_prompt(name: str) -> str:
    """Load a General OCR task prompt, stripped except for position-aware VQA.

    Args:
        name: Task prompt stem (see model card's `general_ocr_prompt`).

    Returns:
        The task prompt, stripped unless it is `position_aware_vqa`.
    """
    prompt = load_prompt(name)
    return prompt if name == "position_aware_vqa" else prompt.strip()


# ---------------------------------------------------------------------------
# Post-processing helpers — copied from the model card.
# ---------------------------------------------------------------------------
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 (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 placeholder `<img src="images/bbox_...">` blocks (model card)."""
    return "\n\n".join(
        block
        for block in text.split("\n\n")
        if not block.strip().startswith('<img src="images/bbox_')
    )


def _parse_region(region_text: str) -> str:
    """Normalize a user-supplied region into the card's compact JSON form.

    Accepts `[x1,y1,x2,y2]`, a JSON array, or whitespace-separated numbers.

    Args:
        region_text: Region input from the UI.

    Returns:
        A compact JSON array string like `"[100,100,900,300]"`.

    Raises:
        ValueError: If the input does not yield four integer coordinates.
    """
    text = region_text.strip()
    if not text:
        raise ValueError("Region is empty. Provide four integers in [0, 1000).")
    try:
        values = json.loads(text)
    except json.JSONDecodeError:
        values = [tok for tok in re.split(r"[\s,]+", text) if tok]
    if not isinstance(values, list) or len(values) != 4:
        raise ValueError("Region must contain exactly four integers in [0, 1000).")
    region = [int(round(float(v))) for v in values]
    for v in region:
        if not 0 <= v <= 999:
            raise ValueError("Region coordinates must be integers in [0, 1000).")
    return json.dumps(region, separators=(",", ":"))


# ---------------------------------------------------------------------------
# Core inference — same contract as the card's `infer(...)`.
# ---------------------------------------------------------------------------
def _generate(images: list, prompt: str, max_new_tokens: int) -> str:
    """Run one greedy generation request over one or two images.

    Args:
        images: One or two RGB-converted PIL images (one request).
        prompt: Fully substituted task prompt.
        max_new_tokens: Generation token budget.

    Returns:
        The raw decoded completion text (unmodified).
    """
    content = [{"type": "image", "image": img.convert("RGB")} for img in images]
    content.append({"type": "text", "text": prompt})
    messages = [{"role": "user", "content": content}]

    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]:]
    return processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0].strip()


def _estimate_duration(*args, max_new_tokens=4096, **kwargs):
    """ZeroGPU duration estimator sized from live measurements (~30ms/token
    plus prefill overhead), covering the full requested token budget."""
    return int(min(280, 15 + float(max_new_tokens) * 0.032))


# ---------------------------------------------------------------------------
# Tab 1: single-image task engine (16 tasks across 4 domains).
# ---------------------------------------------------------------------------
TASKS = {
    # key: (label, prompt stem, task-input keys required)
    "document_markdown": ("Document → Markdown", "document_markdown", ()),
    "document_structured": ("Document → structured JSON", "document_structured", ()),
    "document_layout": ("Layout detection (JSON)", "document_layout", ()),
    "plain_ocr": ("Plain OCR", "plain_ocr", ()),
    "word_spotting": ("Word-level text spotting (JSON)", "word_spotting", ()),
    "reference_grounding": ("Reference grounding (JSON)", "reference_grounding", ("references",)),
    "region_ocr": ("Region OCR", "region_ocr", ("region",)),
    "formula_recognition": ("Formula recognition (LaTeX)", "formula_recognition", ()),
    "chart_parsing": ("Chart parsing (JSON)", "chart_parsing", ()),
    "text_translation": ("Text translation (region)", "text_translation", ("region", "language")),
    "image_classification": ("Image classification", "image_classification", ("target_description", "labels")),
    "visual_question_answering": ("Visual question answering", "visual_question_answering", ("question",)),
    "position_aware_vqa": ("Position-aware VQA (JSON)", "position_aware_vqa", ("question",)),
    "kie_open": ("KIE — open extraction", "kie_open", ()),
    "kie_fields": ("KIE — specified fields", "kie_fields", ("keys",)),
    "kie_schema": ("KIE — schema-based", "kie_schema", ("schema",)),
    "figure_captioning": ("Figure captioning (JSON)", "figure_captioning", ()),
}

# Task inputs shown per selected task; (input key, gr.Component) pairs.
TASK_INPUT_KEYS = (
    "question",
    "references",
    "region",
    "language",
    "labels",
    "target_description",
    "keys",
    "schema",
)

DEFAULT_KEYS_JSON = '["invoice_number", "date", "total"]'
DEFAULT_SCHEMA_JSON = json.dumps(
    {
        "invoice_number": "",
        "date": "",
        "seller": {"name": "", "address": ""},
        "items": [{"name": "", "quantity": "", "unit_price": "", "amount": ""}],
        "total": "",
    },
    ensure_ascii=False,
)
DEFAULT_LABELS_JSON = '["financial report", "invoice", "letter", "academic paper"]'
NO_TABLE_RESPONSE = '[{"source_pages": [], "html": "No cross-page table exists between these pages."}]'


def _build_task_prompt(
    task: str,
    question: str,
    references: str,
    region: str,
    language: str,
    labels: str,
    target_description: str,
    keys: str,
    schema: str,
) -> str:
    """Build the fully substituted task prompt for a single-image task.

    Uses literal `.replace()` (never `str.format`) because prompts contain
    JSON and LaTeX braces.

    Args:
        task: Task key from TASKS.
        question: Question text for VQA tasks.
        references: JSON array of reference strings (grounding).
        region: Region spec `[x1,y1,x2,y2]` in [0,1000) coordinates.
        language: Target language for translation.
        labels: JSON array of candidate labels (classification).
        target_description: What to classify (classification).
        keys: JSON array of field names (KIE fields).
        schema: JSON output template (KIE schema).

    Returns:
        The fully substituted prompt string.
    """
    if task not in TASKS:
        raise ValueError(f"Unknown task: {task}")

    if task == "visual_question_answering":
        if not question.strip():
            raise ValueError("Provide a question for VQA tasks.")
        prompt = _general_ocr_prompt("visual_question_answering").replace("{question}", question.strip())
        return re.sub(r"\n{3,}", "\n\n", prompt).strip()

    prompt = _general_ocr_prompt(task) if task in {
        "plain_ocr", "word_spotting", "reference_grounding", "region_ocr",
        "formula_recognition", "chart_parsing", "text_translation",
        "image_classification", "position_aware_vqa",
    } else load_prompt(task).strip()

    if task == "reference_grounding":
        refs = references.strip() or '["Invoice number", "Total"]'
        refs = json.dumps(json.loads(refs), ensure_ascii=False, separators=(",", ":"))
        prompt = prompt.replace("{references}", refs)
    elif task == "region_ocr":
        prompt = prompt.replace("{target_region}", _parse_region(region))
    elif task == "text_translation":
        if not language.strip():
            raise ValueError("Provide a target language for translation.")
        prompt = prompt.replace("{target_region}", _parse_region(region))
        prompt = prompt.replace("{target_language}", language.strip())
    elif task == "image_classification":
        if not target_description.strip():
            raise ValueError("Provide a target description for classification.")
        labs = json.dumps(
            json.loads(labels) if labels.strip() else DEFAULT_LABELS_JSON,
            ensure_ascii=False,
        )
        prompt = prompt.replace("{target description}", target_description.strip())
        prompt = prompt.replace("{candidate_labels}", labs)
    elif task == "position_aware_vqa":
        if not question.strip():
            raise ValueError("Provide a question for VQA tasks.")
        prompt = prompt.replace("{question}", question.strip())
    elif task == "kie_fields":
        if not keys.strip():
            raise ValueError("Provide a JSON array of field keys.")
        prompt = prompt.replace("{keys_json}", json.dumps(json.loads(keys), ensure_ascii=False))
    elif task == "kie_schema":
        if not schema.strip():
            raise ValueError("Provide a JSON output template.")
        prompt = prompt.replace("{schema_json}", json.dumps(json.loads(schema), ensure_ascii=False))

    return prompt


@spaces.GPU(duration=_estimate_duration)
def run_task(
    image,
    task: str = "document_markdown",
    question: str = "What is the title of this document?",
    references: str = '["Invoice number", "Total"]',
    region: str = "[0,0,999,999]",
    language: str = "English",
    labels: str = DEFAULT_LABELS_JSON,
    target_description: str = "the document type",
    keys: str = DEFAULT_KEYS_JSON,
    schema: str = DEFAULT_SCHEMA_JSON,
    filter_image_tags: bool = True,
    max_new_tokens: int = 4096,
):
    """Run one of OvisOCR3's single-image document tasks on an uploaded image.

    Args:
        image: The document/page image to analyze.
        task: Which of the 16 supported tasks to run.
        question: Question text (VQA tasks).
        references: JSON array of references to ground (reference grounding).
        region: `[x1,y1,x2,y2]` region in [0,1000) coordinates (region OCR / translation).
        language: Target language (text translation).
        labels: JSON array of candidate labels (image classification).
        target_description: What to classify (image classification).
        keys: JSON array of field names (KIE specified fields).
        schema: JSON output template (KIE schema mode).
        filter_image_tags: Drop placeholder `<img src="images/bbox_...">` blocks (Markdown task only).
        max_new_tokens: Maximum number of tokens to generate.

    Returns:
        The model output: Markdown for document parsing tasks, JSON for
        structured tasks, plain text for OCR tasks.
    """
    if image is None:
        raise ValueError("Upload a document image first.")
    max_new_tokens = max(256, min(int(max_new_tokens), MAX_NEW_TOKENS_CEILING))

    prompt = _build_task_prompt(
        task, question, references, region, language, labels,
        target_description, keys, schema,
    )
    raw = _generate([image], prompt, max_new_tokens)

    if task == "document_markdown":
        text = raw.strip()
        if filter_image_tags:
            text = _filter_img_tags(text)
        return _clean_truncated_repeats(text)
    return _clean_truncated_repeats(raw)


def _run_task_ui(
    image,
    task: str = "document_markdown",
    question: str = "What is the title of this document?",
    references: str = '["Invoice number", "Total"]',
    region: str = "[0,0,999,999]",
    language: str = "English",
    labels: str = DEFAULT_LABELS_JSON,
    target_description: str = "the document type",
    keys: str = DEFAULT_KEYS_JSON,
    schema: str = DEFAULT_SCHEMA_JSON,
    filter_image_tags: bool = True,
    max_new_tokens: int = 4096,
):
    """UI wrapper for `run_task` that also fills the rendered Markdown preview."""
    text = run_task(
        image, task, question, references, region, language, labels,
        target_description, keys, schema, filter_image_tags, max_new_tokens,
    )
    return text, text


# ---------------------------------------------------------------------------
# Tab 2: cross-page table reconstruction (two-stage, per model card).
# ---------------------------------------------------------------------------
def _prepare_page_markdown(raw: str, index: int) -> str:
    """Turn a raw page-parse into the page-local Markdown used by Stage 2.

    Applies the card's pipeline: repeated-tail cleanup, code-fence stripping,
    and qualifying image references with the page prefix.

    Args:
        raw: Raw `document_markdown` generation for one page.
        index: 1-based page number.

    Returns:
        Page Markdown with page-qualified image references.
    """
    markdown = _clean_truncated_repeats(raw.strip())
    markdown = re.sub(r'^```(?:markdown|html)?\s*\n|\n```\s*$', '', markdown.strip(), flags=re.I)
    if not markdown:
        raise ValueError(f"Page {index}: empty Markdown")
    return re.sub(
        r'(?<![\w/])(?:page_\d+/)?(images/bbox_\d+_\d+_\d+_\d+\.jpg)',
        rf'page_{index}/\1',
        markdown,
    )


@spaces.GPU(duration=_estimate_duration)
def parse_page(image, max_new_tokens: int = 4096):
    """Stage 1 helper: parse one page to Markdown (for cross-page reconstruction).

    Args:
        image: One page image.
        max_new_tokens: Maximum number of tokens to generate.

    Returns:
        Raw Markdown for the page (before Stage 2 post-processing).
    """
    if image is None:
        raise ValueError("Upload a page image first.")
    max_new_tokens = max(256, min(int(max_new_tokens), MAX_NEW_TOKENS_CEILING))
    return _generate([image], load_prompt("document_markdown").strip(), max_new_tokens)


def _parse_page1(image, max_new_tokens):
    """UI handler: parse page 1 and store its Stage-2-ready Markdown."""
    raw = parse_page(image, max_new_tokens)
    return _prepare_page_markdown(raw, 1)


def _parse_page2(image, max_new_tokens):
    """UI handler: parse page 2 and store its Stage-2-ready Markdown."""
    raw = parse_page(image, max_new_tokens)
    return _prepare_page_markdown(raw, 2)


@spaces.GPU(duration=_estimate_duration)
def reconstruct_crosspage_table(
    page1_image,
    page2_image,
    page1_markdown: str,
    page2_markdown: str,
    max_new_tokens: int = 4096,
):
    """Stage 2: reconstruct a table continued across two adjacent pages.

    Args:
        page1_image: The earlier page image.
        page2_image: The following page image.
        page1_markdown: Stage-1 Markdown for page 1 (from `parse_page`).
        page2_markdown: Stage-1 Markdown for page 2 (from `parse_page`).
        max_new_tokens: Maximum number of tokens to generate.

    Returns:
        A JSON array; each record has `source_pages` and an `html` table.
        When no cross-page table exists, the card's negative prediction is
        returned verbatim.
    """
    max_new_tokens = max(256, min(int(max_new_tokens), MAX_NEW_TOKENS_CEILING))

    pages = [img for img in (page1_image, page2_image) if img is not None]
    if len(pages) != 2:
        raise ValueError("Upload both page images for cross-page reconstruction.")
    if not page1_markdown.strip() or not page2_markdown.strip():
        raise ValueError("Parse both pages first so their Markdown can be edited/reused.")

    md1 = _prepare_page_markdown(page1_markdown, 1)
    md2 = _prepare_page_markdown(page2_markdown, 2)

    has_table = re.compile(r'<table\b[^>]*>.*?</table\s*>', re.I | re.S)
    if not (has_table.search(md1) or has_table.search(md2)):
        return NO_TABLE_RESPONSE

    page_context = "\n\n".join(
        f'<INPUT_PAGE index="{index}">\n{markdown}\n</INPUT_PAGE>'
        for index, markdown in enumerate((md1, md2), start=1)
    )
    prompt = load_prompt("crosspage_table").removesuffix("\n") + "\n\nInput Markdown:\n\n" + page_context
    raw = _generate(pages, prompt, max_new_tokens)
    return _clean_truncated_repeats(raw)


def _reconstruct_ui(page1_image, page2_image, page1_markdown, page2_markdown, max_new_tokens):
    """UI wrapper for cross-page reconstruction; returns the JSON result text."""
    return reconstruct_crosspage_table(
        page1_image, page2_image, page1_markdown, page2_markdown, max_new_tokens
    )


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

TASK_TO_INPUTS = {key: spec[2] for key, spec in TASKS.items()}


def _toggle_task_inputs(task: str):
    """Return per-task-input visibility flags for the selected task.

    Args:
        task: Task key from the dropdown.

    Returns:
        A list of `gr.update` visibility values, ordered as TASK_INPUT_KEYS.
    """
    required = set(TASK_TO_INPUTS.get(task, ()))
    return [
        gr.update(visible=key in required) for key in TASK_INPUT_KEYS
    ]


with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 📑 OvisOCR3 — Unified Document Understanding

            Interactive demo for [**ATH-MaaS/OvisOCR3**](https://huggingface.co/ATH-MaaS/OvisOCR3),
            a compact 0.8B vision-language model that covers **document parsing**, **general
            OCR**, **key information extraction**, **figure captioning**, and **cross-page
            table reconstruction** — all selected through task prompts on the same checkpoint.

            Upload a page and pick a task. Task inputs appear only where the task needs them.
            """
        )

        with gr.Tabs():
            # ------------------------- Tab 1: tasks -------------------------
            with gr.TabItem("Document tasks"):
                with gr.Row():
                    with gr.Column(scale=1):
                        image_in = gr.Image(type="pil", label="Document image", height=460)
                        task_dd = gr.Dropdown(
                            choices=[(label, key) for key, (label, _, _) in TASKS.items()],
                            value="document_markdown",
                            label="Task",
                            filterable=False,
                        )
                        run_btn = gr.Button("Run task", variant="primary")

                        with gr.Accordion("Task inputs", open=True):
                            question_tb = gr.Textbox(
                                label="Question (VQA tasks)",
                                value="What is the title of this document?",
                                lines=2,
                            )
                            references_tb = gr.Textbox(
                                label="References — JSON array (reference grounding)",
                                value='["Invoice number", "Total"]',
                                lines=2,
                            )
                            region_tb = gr.Textbox(
                                label="Region — [x1,y1,x2,y2] in [0,1000) (region OCR / translation)",
                                value="[0,0,999,999]",
                                lines=1,
                            )
                            language_tb = gr.Textbox(
                                label="Target language (translation)",
                                value="English",
                                lines=1,
                            )
                            labels_tb = gr.Textbox(
                                label="Candidate labels — JSON array (classification)",
                                value=DEFAULT_LABELS_JSON,
                                lines=2,
                            )
                            target_desc_tb = gr.Textbox(
                                label="Target description (classification)",
                                value="the document type",
                                lines=1,
                            )
                            keys_tb = gr.Textbox(
                                label=f"Field keys — JSON array (KIE fields). Example: {DEFAULT_KEYS_JSON}",
                                value=DEFAULT_KEYS_JSON,
                                lines=2,
                            )
                            schema_tb = gr.Textbox(
                                label="Output template — JSON (KIE schema). Empty strings are value placeholders.",
                                value=DEFAULT_SCHEMA_JSON,
                                lines=6,
                            )

                        with gr.Accordion("Advanced settings", open=False):
                            filter_tags = gr.Checkbox(
                                value=True,
                                label="Filter figure placeholder <img> tags (Markdown task only)",
                            )
                            max_tokens = gr.Slider(
                                512, MAX_NEW_TOKENS_CEILING, value=4096, step=512,
                                label="Max new tokens",
                            )

                    with gr.Column(scale=1):
                        task_out = gr.Code(label="Output", lines=24)
                        with gr.Accordion("Rendered Markdown preview", open=False):
                            task_render = gr.Markdown()

                task_dd.change(
                    _toggle_task_inputs,
                    inputs=[task_dd],
                    outputs=[
                        question_tb, references_tb, region_tb, language_tb,
                        labels_tb, target_desc_tb, keys_tb, schema_tb,
                    ],
                )
                # Apply visibility for the default task on load.
                demo.load(
                    _toggle_task_inputs,
                    inputs=[task_dd],
                    outputs=[
                        question_tb, references_tb, region_tb, language_tb,
                        labels_tb, target_desc_tb, keys_tb, schema_tb,
                    ],
                )

                task_inputs = [
                    image_in, task_dd, question_tb, references_tb, region_tb,
                    language_tb, labels_tb, target_desc_tb, keys_tb, schema_tb,
                    filter_tags, max_tokens,
                ]
                run_btn.click(
                    _run_task_ui,
                    inputs=task_inputs,
                    outputs=[task_out, task_render],
                    api_name="run_task",
                )

                gr.Examples(
                    examples=[
                        ["examples/paper_page.png", "document_markdown"],
                        ["examples/financial_table.jpg", "kie_open"],
                        ["examples/handwriting.jpg", "plain_ocr"],
                        ["examples/math_textbook.png", "formula_recognition"],
                        ["examples/magazine_page.png", "figure_captioning"],
                        ["examples/paper_page.png", "document_structured"],
                    ],
                    inputs=[image_in, task_dd],
                    outputs=[task_out, task_render],
                    fn=_run_task_ui,
                    cache_examples=True,
                    cache_mode="lazy",
                )

            # --------------------- Tab 2: cross-page tables ---------------------
            with gr.TabItem("Cross-page table"):
                gr.Markdown(
                    """
                    **Two-stage cross-page table reconstruction** (one request per stage):
                    1. Parse each page to Markdown (Stage 1 — buttons below, one call per page).
                    2. Reconstruct the table that continues across the page boundary (Stage 2 —
                    one two-image request with both pages and their Markdown).

                    Both page images are required. You can edit the parsed Markdown before
                    Stage 2.
                    """
                )
                with gr.Row():
                    with gr.Column():
                        cp_page1_img = gr.Image(type="pil", label="Page 1 (earlier page)", height=320)
                        cp_parse1_btn = gr.Button("Parse page 1 → Markdown", variant="secondary")
                        cp_page1_md = gr.Code(label="Page 1 Markdown", language="markdown", lines=10)
                    with gr.Column():
                        cp_page2_img = gr.Image(type="pil", label="Page 2 (following page)", height=320)
                        cp_parse2_btn = gr.Button("Parse page 2 → Markdown", variant="secondary")
                        cp_page2_md = gr.Code(label="Page 2 Markdown", language="markdown", lines=10)
                with gr.Accordion("Advanced settings", open=False):
                    cp_max_tokens = gr.Slider(
                        512, MAX_NEW_TOKENS_CEILING, value=4096, step=512,
                        label="Max new tokens (Stage 1 parsing)",
                    )
                    cp_max_tokens2 = gr.Slider(
                        512, MAX_NEW_TOKENS_CEILING, value=4096, step=512,
                        label="Max new tokens (Stage 2 reconstruction)",
                    )
                cp_run_btn = gr.Button("Reconstruct cross-page table", variant="primary")
                cp_out = gr.Code(label="Reconstructed table (JSON)", language="json", lines=12)

                cp_parse1_btn.click(
                    _parse_page1,
                    inputs=[cp_page1_img, cp_max_tokens],
                    outputs=[cp_page1_md],
                    api_name="parse_page1",
                )
                cp_parse2_btn.click(
                    _parse_page2,
                    inputs=[cp_page2_img, cp_max_tokens],
                    outputs=[cp_page2_md],
                    api_name="parse_page2",
                )
                cp_run_btn.click(
                    _reconstruct_ui,
                    inputs=[cp_page1_img, cp_page2_img, cp_page1_md, cp_page2_md, cp_max_tokens2],
                    outputs=[cp_out],
                    api_name="reconstruct_crosspage_table",
                )

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