"""NaviDC-OCR — document parsing across digital and camera-captured documents.

Two-stage pipeline, faithful to the authors' reference implementation
(https://github.com/caipeng328/NaviDC-OCR):

1. Layout stage — the page is resized to 1036x1036 and the model predicts
   reading-ordered blocks as `<box:...><label:...><angle>` (boxes in
   "Detection" mode, multi-point polygons in "Segmentation" mode, which is what
   the paper uses for curved / camera-captured pages).
2. Recognition stage — every block is cropped (polygon-masked when needed),
   de-rotated, and recognized with the block-type-specific prompt and sampling
   parameters from `NaviOCR/vlm_utils/NaviOCR_client.py`, then post-processed
   (OTSL tables -> HTML, LaTeX equation fixes) with the authors' post-processors.

A third mode skips layout and runs the authors' single-region path
(`NaviOCRClient.block_parse`) on the whole image, which is how the model card
demonstrates chart-to-table extraction, seal reading and table/formula crops.
"""

import base64
import io
import os
import re
import tempfile
import time
from dataclasses import asdict
from typing import Any

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces  # noqa: F401  (must precede torch / transformers)
import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoModelForImageTextToText, AutoProcessor

from NaviOCR.vlm_utils.NaviOCR_client import (
    DEFAULT_PROMPTS,
    DEFAULT_SAMPLING_PARAMS,
    LAYOUT_PROMPTS,
    NaviOCRClient,
)
from NaviOCR.vlm_utils.post_process.otsl2html import convert_otsl_to_html
from NaviOCR.vlm_utils.structs import ContentBlock
from NaviOCR.vlm_utils.vlm_client import SamplingParams

MODEL_ID = "StarDoc-AI/NaviDC-OCR"

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

client = NaviOCRClient(
    backend="transformers",
    model=model,
    processor=processor,
    prompts=DEFAULT_PROMPTS,
    sampling_params=DEFAULT_SAMPLING_PARAMS,
    batch_size=0,  # the authors' transformers backend default: one region at a time
    use_tqdm=True,
)

PARATEXT_TYPES = {"header", "footer", "page_number", "aside_text", "page_footnote", "unknown"}
CAPTION_TYPES = {
    "table_caption",
    "image_caption",
    "code_caption",
    "table_footnote",
    "image_footnote",
}
BLOCK_COLORS = {
    "title": (216, 27, 96),
    "text": (30, 136, 229),
    "table": (0, 137, 123),
    "table_caption": (0, 172, 193),
    "table_footnote": (0, 172, 193),
    "image": (245, 124, 0),
    "image_caption": (251, 192, 45),
    "image_footnote": (251, 192, 45),
    "equation": (142, 36, 170),
    "equation_block": (142, 36, 170),
    "code": (94, 53, 177),
    "code_caption": (121, 85, 72),
    "algorithm": (94, 53, 177),
    "list": (57, 73, 171),
    "ref_text": (109, 76, 65),
    "seal": (211, 47, 47),
    "char": (0, 121, 107),
}
DEFAULT_COLOR = (117, 117, 117)

# Block types the single-region mode exposes, with the authors' prompt keys.
REGION_TASKS = [
    ("Text", "text"),
    ("Table \u2192 HTML", "table"),
    ("Formula \u2192 LaTeX", "formula"),
    ("Code", "code"),
    ("Chart / scientific figure \u2192 table", "char"),
    ("Seal", "seal"),
]
# Prompt keys and block-type names differ for formulas ("formula" vs "equation").
TASK_BLOCK_TYPES = {"formula": "equation"}
# The model prefixes recognized code with its own language marker, e.g. `<_Python_>`.
CODE_LANG_RE = re.compile(r"^\s*<_([A-Za-z0-9+#._\- ]+)_>\s*")


def _sampling_params(task: str, max_new_tokens: int) -> SamplingParams:
    """Authors' per-task sampling params, with a bounded generation length."""
    base = DEFAULT_SAMPLING_PARAMS.get(task) or DEFAULT_SAMPLING_PARAMS["default"]
    fields = asdict(base)
    fields["max_new_tokens"] = int(max_new_tokens)
    return SamplingParams(**fields)


def _points(bbox, width: int, height: int) -> np.ndarray:
    pts = np.array(bbox, dtype=np.float32).reshape(-1, 2)
    pts[:, 0] *= width
    pts[:, 1] *= height
    return pts.astype(np.int32)


def _crop(image: Image.Image, bbox) -> Image.Image:
    pts = _points(bbox, image.width, image.height)
    x1, y1 = int(pts[:, 0].min()), int(pts[:, 1].min())
    x2, y2 = int(pts[:, 0].max()), int(pts[:, 1].max())
    x1, y1 = max(0, x1), max(0, y1)
    x2, y2 = min(image.width, max(x2, x1 + 1)), min(image.height, max(y2, y1 + 1))
    return image.crop((x1, y1, x2, y2))


def _data_uri(image: Image.Image, max_width: int = 900) -> str:
    if image.width > max_width:
        ratio = max_width / image.width
        image = image.resize((max_width, max(1, int(image.height * ratio))), Image.Resampling.LANCZOS)
    buffer = io.BytesIO()
    image.convert("RGB").save(buffer, format="JPEG", quality=88)
    return "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode()


def _font(size: int):
    try:
        return ImageFont.load_default(size=size)
    except TypeError:  # very old Pillow
        return ImageFont.load_default()


def draw_layout(image: Image.Image, blocks: list) -> Image.Image:
    """Overlay the predicted blocks, numbered in predicted reading order."""
    canvas = image.convert("RGB").copy()
    overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)
    line_width = max(2, round(min(canvas.size) / 400))
    font = _font(max(13, round(min(canvas.size) / 55)))

    for order, block in enumerate(blocks, start=1):
        color = BLOCK_COLORS.get(block.type, DEFAULT_COLOR)
        pts = _points(block.bbox, canvas.width, canvas.height)
        if len(pts) == 2:
            xy = [(int(pts[0][0]), int(pts[0][1])), (int(pts[1][0]), int(pts[1][1]))]
            draw.rectangle(xy, outline=color + (255,), width=line_width)
            anchor = xy[0]
        else:
            polygon = [(int(x), int(y)) for x, y in pts]
            draw.polygon(polygon, outline=color + (255,), fill=color + (28,), width=line_width)
            anchor = min(polygon, key=lambda p: (p[1], p[0]))

        label = f"{order} {block.type}"
        if block.angle:
            label += f" {block.angle}\u00b0"
        tx, ty = anchor[0], max(0, anchor[1] - font.size - 4)
        text_box = draw.textbbox((tx, ty), label, font=font)
        draw.rectangle(
            (text_box[0] - 2, text_box[1] - 2, text_box[2] + 2, text_box[3] + 2),
            fill=color + (235,),
        )
        draw.text((tx, ty), label, fill=(255, 255, 255, 255), font=font)

    return Image.alpha_composite(canvas.convert("RGBA"), overlay).convert("RGB")


def _fenced_code(content: str) -> str:
    match = CODE_LANG_RE.match(content)
    language = ""
    if match:
        language = match.group(1).strip().lower().replace(" ", "")
        content = content[match.end() :]
    return f"```{language}\n{content}\n```"


def blocks_to_markdown(image: Image.Image, blocks: list, drop_paratext: bool):
    """Assemble reading-ordered blocks into Markdown (raw + display variants)."""
    parts: list[str] = []
    figures: dict[str, Image.Image] = {}

    for block in blocks:
        block_type = block.type
        content = (block.content or "").strip()

        if drop_paratext and block_type in PARATEXT_TYPES:
            continue

        if block_type == "image":
            key = f"figure_{len(figures) + 1}.jpg"
            figures[key] = _crop(image, block.bbox)
            parts.append(f"![{key}]({key})")
            continue

        if not content:
            continue

        if block_type == "title":
            parts.append(f"## {content}")
        elif block_type == "table":
            parts.append(content)  # already OTSL -> HTML in post-processing
        elif block_type == "char":
            parts.append(convert_otsl_to_html(content) or content)
        elif block_type in {"code", "algorithm"}:
            parts.append(_fenced_code(content))
        elif block_type in CAPTION_TYPES:
            parts.append(f"*{content}*")
        elif block_type == "seal":
            parts.append(f"**[seal]** {content}")
        else:  # text, list, ref_text, equation, phonetic, header/footer, ...
            parts.append(content)

    raw_markdown = "\n\n".join(parts).strip()
    display_markdown = raw_markdown
    for key, crop in figures.items():
        display_markdown = display_markdown.replace(
            f"![{key}]({key})",
            f'<img src="{_data_uri(crop)}" style="max-width:100%;border-radius:6px" />',
        )
    return raw_markdown, display_markdown


def _write_markdown(markdown: str) -> str:
    directory = tempfile.mkdtemp(prefix="navidc_ocr_")
    path = os.path.join(directory, "navidc_ocr.md")
    with open(path, "w", encoding="utf-8") as handle:
        handle.write(markdown)
    return path


def _estimate_duration(*args, **kwargs) -> int:
    """Measured on ZeroGPU: single regions 5-13 s, a dense 31-region page 73 s.

    Runtime is dominated by generated tokens, so scale with the per-region cap
    (105 s at the default 2048, the measured worst case x1.4).
    """
    max_new_tokens = 2048
    if len(args) > 4:
        max_new_tokens = args[4]
    max_new_tokens = int(kwargs.get("max_new_tokens", max_new_tokens) or 2048)
    return int(min(180, 60 + 0.022 * max_new_tokens))


@spaces.GPU(duration=_estimate_duration)
def parse_document(
    image: Image.Image,
    layout_mode: str = "Detection",
    region_task: str = "text",
    drop_paratext: bool = True,
    max_new_tokens: int = 2048,
    progress=gr.Progress(track_tqdm=True),
) -> tuple[Image.Image, str, str, list[dict[str, Any]], str, str]:
    """Parse a document page into Markdown with NaviDC-OCR.

    Args:
        image: A document page — a digital page, a scan, or a camera photo.
        layout_mode: "Detection" for axis-aligned boxes (digital pages, flat
            scans), "Segmentation" for multi-point polygons (camera-captured,
            curved or crumpled pages), or "Region" to skip layout and recognize
            the whole image as one block.
        region_task: The block type used in "Region" mode — one of text, table,
            formula, code, char (chart/scientific figure), seal.
        drop_paratext: Drop headers, footers, page numbers and margin notes.
        max_new_tokens: Generation cap per region.

    Returns:
        The layout overlay, rendered Markdown, raw Markdown, the block list as
        JSON, a downloadable .md file, and a short run report.
    """
    if image is None:
        raise gr.Error("Please provide a document image first.")

    started = time.time()
    page = image.convert("RGB") if isinstance(image, Image.Image) else Image.open(image).convert("RGB")
    helper = client.helper
    mode = layout_mode if layout_mode in LAYOUT_PROMPTS else "Region"

    # ---- single-region mode: the authors' block_parse path ----------------
    if mode == "Region":
        task = region_task if region_task in DEFAULT_PROMPTS else "text"
        crop = helper.resize_by_need(page)
        output = client.client.predict(
            crop,
            DEFAULT_PROMPTS[task],
            _sampling_params(task, max_new_tokens),
        )
        block = ContentBlock(
            type=TASK_BLOCK_TYPES.get(task, task),
            bbox=[[0.0, 0.0], [1.0, 1.0]],
            content=output,
        )
        blocks = helper.post_process([block]) or [block]
        raw_markdown, display_markdown = blocks_to_markdown(page, blocks, False)
        seconds = time.time() - started
        report = (
            f"Single region recognized as `{task}` \u2014 {seconds:.1f}s.  \n"
            f"Switch to a full-page mode to run layout analysis first."
        )
        return (
            page,
            display_markdown,
            raw_markdown,
            [dict(item) for item in blocks],
            _write_markdown(raw_markdown),
            report,
        )

    # ---- stage 1: layout ------------------------------------------------
    layout_image = helper.prepare_for_layout(page)  # resized to 1036x1036
    raw_layout = client.client.predict(
        layout_image,
        LAYOUT_PROMPTS[mode],
        _sampling_params("layout", max(1024, int(max_new_tokens))),
    )
    blocks = helper.parse_layout_output(raw_layout)
    layout_seconds = time.time() - started

    if not blocks:
        report = (
            f"No layout blocks were parsed in **{mode}** mode "
            f"({layout_seconds:.1f}s). Raw layout output is in the *Blocks* tab."
        )
        return (
            page,
            "",
            "",
            [{"raw_layout_output": raw_layout}],
            _write_markdown(""),
            report,
        )

    # ---- stage 2: per-region recognition --------------------------------
    block_images, prompts, params, indices = helper.prepare_for_extract(page, blocks)
    params = [
        _sampling_params(blocks[idx].type, max_new_tokens) for idx in indices
    ]
    if block_images:
        outputs = client.client.batch_predict(block_images, prompts, params)
        for idx, output in zip(indices, outputs):
            blocks[idx].content = output

    blocks = helper.post_process(blocks)

    raw_markdown, display_markdown = blocks_to_markdown(page, blocks, drop_paratext)
    overlay = draw_layout(page, blocks)
    total_seconds = time.time() - started

    counts: dict[str, int] = {}
    for block in blocks:
        counts[block.type] = counts.get(block.type, 0) + 1
    summary = ", ".join(f"{count}\u00d7{name}" for name, count in sorted(counts.items()))
    report = (
        f"**{len(blocks)} regions** in `{mode}` mode \u2014 {summary}.  \n"
        f"Layout {layout_seconds:.1f}s \u00b7 total {total_seconds:.1f}s."
    )

    return (
        overlay,
        display_markdown,
        raw_markdown,
        [dict(block) for block in blocks],
        _write_markdown(raw_markdown),
        report,
    )


CSS = """
#col-container { max-width: 1400px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
#doc-md { overflow-x: auto; }
#doc-md table { border-collapse: collapse; }
#doc-md td, #doc-md th { border: 1px solid var(--border-color-primary); padding: 4px 8px; }
"""

LATEX = [
    {"left": "$$", "right": "$$", "display": True},
    {"left": "$", "right": "$", "display": False},
    {"left": "\\(", "right": "\\)", "display": False},
    {"left": "\\[", "right": "\\]", "display": True},
]

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="NaviDC-OCR") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# NaviDC-OCR — document parsing, digital *and* camera-captured

A 1.2B document-parsing VLM that reads layout, text, tables, formulas and code
off flat scans **and** photographed / crumpled pages, and returns Markdown.

[model](https://huggingface.co/StarDoc-AI/NaviDC-OCR) ·
[paper](https://huggingface.co/papers/2608.12898) ·
[code](https://github.com/caipeng328/NaviDC-OCR)
"""
        )
        with gr.Row():
            with gr.Column(scale=4):
                image = gr.Image(label="Document page", type="pil", height=460)
                layout_mode = gr.Radio(
                    choices=[
                        ("Full page, boxes — digital pages & flat scans", "Detection"),
                        (
                            "Full page, multi-point — photos, curved or crumpled pages",
                            "Segmentation",
                        ),
                        ("Single region — the image is one table / formula / …", "Region"),
                    ],
                    value="Detection",
                    label="Parsing mode",
                )
                region_task = gr.Dropdown(
                    choices=REGION_TASKS,
                    value="table",
                    label="Region type",
                    visible=False,
                )
                run_button = gr.Button("Parse document", variant="primary")
                report = gr.Markdown()
                with gr.Accordion("Advanced settings", open=False):
                    drop_paratext = gr.Checkbox(
                        value=True,
                        label="Full page: drop headers, footers, page numbers, margin notes",
                    )
                    max_new_tokens = gr.Slider(
                        256, 4096, value=2048, step=128, label="Max new tokens per region"
                    )
            with gr.Column(scale=6):
                with gr.Tabs():
                    with gr.Tab("Document"):
                        document = gr.Markdown(
                            latex_delimiters=LATEX,
                            elem_id="doc-md",
                            show_copy_button=True,
                        )
                    with gr.Tab("Markdown source"):
                        markdown_source = gr.Code(
                            language="markdown",
                            lines=28,
                            interactive=False,
                            label="Markdown",
                            wrap_lines=True,
                        )
                    with gr.Tab("Layout"):
                        overlay = gr.Image(label="Predicted regions (reading order)", height=620)
                    with gr.Tab("Blocks"):
                        blocks_json = gr.JSON(label="Blocks")
                markdown_file = gr.DownloadButton("Download Markdown")

        gr.Examples(
            examples=[
                ["examples/journal_page.jpg", "Detection", "table"],
                ["examples/crumpled_page.jpg", "Segmentation", "table"],
                ["examples/table.png", "Region", "table"],
                ["examples/formula.png", "Region", "formula"],
                ["examples/code.png", "Region", "code"],
                ["examples/scientific_figure.png", "Region", "char"],
            ],
            inputs=[image, layout_mode, region_task],
            outputs=[overlay, document, markdown_source, blocks_json, markdown_file, report],
            fn=parse_document,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples from the NaviDC-OCR model card",
        )

    layout_mode.change(
        fn=lambda mode: gr.update(visible=(mode == "Region")),
        inputs=[layout_mode],
        outputs=[region_task],
        show_api=False,
        queue=False,
    )

    gr.on(
        triggers=[run_button.click],
        fn=parse_document,
        inputs=[image, layout_mode, region_task, drop_paratext, max_new_tokens],
        outputs=[overlay, document, markdown_source, blocks_json, markdown_file, report],
    )

if __name__ == "__main__":
    demo.queue(max_size=16).launch(mcp_server=True)