"""Gradio demo for Indra OCR (datalab-to/indra-ocr).

A 1M-parameter ONNX OCR model that turns PDFs and images into markdown, plain
text, or JSON — with a bounding box for every line and every layout block, and
reading order — entirely on CPU.
"""

import json
import threading
import time
from pathlib import Path

import gradio as gr
import indra_ocr
from indra_ocr.document import MAX_PAGE_PIXELS
from PIL import Image, ImageDraw

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

TAG_COLORS = {
    "Text": (9, 105, 218),            # blue
    "Section-Header": (221, 74, 104),  # red
    "Page-Header": (160, 160, 160),   # gray
    "Page-Footer": (160, 160, 160),   # gray
    "List-Group": (63, 185, 80),     # green
    "Table": (255, 140, 26),          # orange
    "Caption": (188, 141, 224),      # purple
    "Footnote": (130, 170, 255),     # light blue
    "Table-Of-Contents": (255, 214, 10),  # yellow
    "Form": (46, 160, 170),          # teal
}

MAX_PAGES = 5  # demo cap: fine for a look; the library reads whole files too

_converters = {}
_lock = threading.Lock()


def _converter(mode: str, furniture: bool):
    """One Converter per (mode, furniture) combination, built once and reused."""
    key = (mode, furniture)
    with _lock:
        if key not in _converters:
            _converters[key] = indra_ocr.load(
                threads=4, mode=mode, furniture=furniture
            )
        return _converters[key]


print("[indra] loading models (first boot downloads ~76 MB from datalab-to/indra-ocr)...")
_t0 = time.perf_counter()
_converter("auto", False)  # eager: the default pipeline is ready before the UI
print(f"[indra] models ready in {time.perf_counter() - _t0:.1f}s")


def _is_pdf(path: str) -> bool:
    try:
        with open(path, "rb") as f:
            return b"%PDF-" in f.read(1024)
    except OSError:
        return False


def _pdf_page_scale(dpi: float, page) -> float:
    """The same render scale Indra's Converter uses (dpi, lowered for huge pages)."""
    w, h = page.get_size()
    k = dpi / 72.0
    return min(k, (MAX_PAGE_PIXELS / max(w * h, 1.0)) ** 0.5)


def _rotate(img: Image.Image, degrees_ccw: int) -> Image.Image:
    if degrees_ccw % 360 == 0:
        return img
    turns = {90: Image.Transpose.ROTATE_270, 180: Image.Transpose.ROTATE_180}
    turns[270] = Image.Transpose.ROTATE_90
    key = degrees_ccw % 360
    return img.transpose(turns[key]) if key in turns else img


def _draw_preview(img: Image.Image, page) -> Image.Image:
    """Page rendered upright with layout blocks (colored, labeled) and lines."""
    img = _rotate(img, page.rotation)
    draw = ImageDraw.Draw(img)
    for ln in page.lines:  # faint blue boxes: every detected line
        draw.rectangle(ln.box, outline=(185, 210, 255), width=2)
    for b in page.blocks:  # colored boxes with labels: layout blocks
        color = TAG_COLORS.get(b["label"], (110, 110, 110))
        box = b["bbox"]
        draw.rectangle(box, outline=color, width=4)
        label = f" {b['label']} "
        lx, ly = box[0], max(0.0, box[1] - 18)
        tw = draw.textlength(label)
        draw.rectangle((lx, ly, lx + tw, ly + 17), fill=color)
        draw.text((lx, ly + 1), label, fill=(255, 255, 255))
    return img


def _page_dict(page) -> dict:
    """A Page in the same shape as the CLI's `--format json` output."""
    return {
        "markdown": page.markdown,
        "text": page.text,
        "source": page.source,
        "rotation": page.rotation,
        "skew": round(page.skew, 2),
        "note": page.note,
        "timings": {k: round(v, 1) for k, v in page.timings.items()},
        "lines": [
            {"box": [round(c, 1) for c in ln.box], "text": ln.text,
             "conf": round(ln.conf, 3)}
            for ln in page.lines
        ],
        "blocks": [
            {"label": b["label"], "bbox": [round(c, 1) for c in b["bbox"]],
             "text": b["text"]}
            for b in page.blocks
        ],
    }


def _status_line(i: int, page) -> str:
    bits = [f"page {i}", f"source: {page.source}",
            f"{round(page.timings.get('total_ms', 0))} ms",
            f"{len(page.lines)} lines · {len(page.blocks)} blocks"]
    if page.rotation:
        bits.append(f"turned {page.rotation}°")
    if page.skew:
        bits.append(f"deskewed {page.skew:.1f}°")
    if page.note:
        bits.append(f"({page.note})")
    return " · ".join(bits)


def convert_file(
    file_path: str,
    mode: str = "auto",
    output_format: str = "markdown",
    dpi: float = 200.0,
    keep_furniture: bool = False,
    draw_boxes: bool = True,
):
    """Read a PDF or image document: markdown / plain text / JSON, with layout.

    Indra OCR detects every text line and every layout block (headings, lists,
    tables, captions, footnotes...), orders them for reading, and writes the
    page as markdown. Digital PDFs are read through their own text layer where
    it passes Indra's checks; scans and images go through the OCR models.

    Args:
        file_path: the uploaded PDF or image file.
        mode: "auto" reads the PDF text layer where it passes the checks and
            OCRs the rest; "ocr" always OCRs; "text" trusts the text layer.
        output_format: "markdown", "plain text", or "json" (lines, boxes,
            layout blocks, confidences, timings).
        dpi: render resolution for PDFs.
        keep_furniture: keep page headers and footers in the output.
        draw_boxes: overlay layout blocks and line boxes on the page preview.
    """
    if not file_path:
        raise gr.Error("Please upload a PDF or image file first.")
    conv = _converter(mode, bool(keep_furniture))

    try:
        pages = list(conv.convert(file_path, max_pages=MAX_PAGES))
    except indra_ocr.InputError as e:
        raise gr.Error(str(e))
    if not pages:
        raise gr.Error("No pages could be read from that file.")

    # ---- preview of page 0, drawn from the page's own boxes -------------
    first_i, first_page = pages[0]
    if _is_pdf(file_path):
        import pypdfium2 as pdfium

        doc = pdfium.PdfDocument(file_path)
        try:
            pg = doc[first_i]
            preview = pg.render(
                scale=_pdf_page_scale(dpi, pg), grayscale=False
            ).to_pil()
        finally:
            doc.close()
    else:
        preview = Image.open(file_path).convert("RGB")
    if draw_boxes:
        preview = _draw_preview(preview, first_page)

    # ---- result in the chosen format -----------------------------------
    if output_format == "json":
        result = json.dumps(
            {"pages": [_page_dict(p) for _, p in pages]}, ensure_ascii=False, indent=1
        )
    elif output_format == "plain text":
        result = "\n\n".join(p.text for _, p in pages)
    else:
        result = "\n\n---\n\n".join(p.markdown or "" for _, p in pages)

    rendered = "\n\n---\n\n".join(p.markdown or "" for _, p in pages)
    status = "\n".join(_status_line(i, p) for i, p in pages)
    if len(pages) >= MAX_PAGES:
        status += f"\n(first {MAX_PAGES} pages only — the library reads whole files)"

    return preview, result, rendered, status


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 📄 Indra OCR
        **[datalab-to/indra-ocr](https://huggingface.co/datalab-to/indra-ocr)** — a 1M-parameter ONNX OCR model that reads PDFs and images
        and turns them into **markdown**, plain text, or **JSON with layout blocks and line boxes** — all on CPU. It reads handwriting, 80+ languages, and pages at any angle.

        Drop in a document below, or click one of the example pages.
        [Model card](https://huggingface.co/datalab-to/indra-ocr) · [code](https://github.com/datalab-to/indra) (Apache-2.0; weights under a modified OpenRAIL-M license)
        """
    )

    with gr.Row():
        with gr.Column(scale=5):
            file_in = gr.File(
                label="PDF or image",
                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp", ".bmp",
                            ".tif", ".tiff"],
                type="filepath",
            )
            run_btn = gr.Button("Convert", variant="primary")
            with gr.Accordion("Advanced settings", open=False):
                mode = gr.Radio(
                    ["auto", "ocr", "text"],
                    value="auto",
                    label="Page reading mode",
                    info="auto: PDF text layer where it passes the checks, OCR the rest · ocr: never read the text layer · text: trust the text layer",
                )
                output_format = gr.Radio(
                    ["markdown", "plain text", "json"],
                    value="markdown",
                    label="Output format",
                )
                dpi = gr.Slider(100, 300, value=200, step=10, label="PDF render DPI")
                keep_furniture = gr.Checkbox(
                    value=False, label="Keep page headers and footers (furniture)"
                )
                draw_boxes = gr.Checkbox(
                    value=True,
                    label="Draw layout blocks and line boxes on the preview",
                )
        with gr.Column(scale=4):
            page_preview = gr.Image(
                label="Detected layout (faint blue: lines · colored: blocks)",
                type="pil",
                interactive=False,
            )

    with gr.Tabs():
        with gr.Tab("Raw output"):
            result_out = gr.Textbox(
                label="Result",
                lines=20,
                max_lines=45,
                buttons=["copy"],
            )
        with gr.Tab("Rendered markdown"):
            rendered_out = gr.Markdown(label="Rendered markdown", value="")
    status_out = gr.Textbox(
        label="Per-page status and timing", lines=4, interactive=False
    )

    run_btn.click(
        convert_file,
        inputs=[file_in, mode, output_format, dpi, keep_furniture, draw_boxes],
        outputs=[page_preview, result_out, rendered_out, status_out],
        api_name="convert",
        concurrency_limit=1,
    )
    file_in.upload(
        convert_file,
        inputs=[file_in, mode, output_format, dpi, keep_furniture, draw_boxes],
        outputs=[page_preview, result_out, rendered_out, status_out],
        api_name="convert_upload",
        concurrency_limit=1,
    )

    gr.Examples(
        examples=[
            ["examples/handwriting_letter.png"],
            ["examples/korean_letter.png"],
            ["examples/thai_calendar.png"],
            ["examples/tilted_memo.png"],
            ["examples/table_pmc.png"],
            ["examples/multi_column.png"],
            ["examples/chinese_essay.png"],
            ["examples/rotated_letter.png"],
        ],
        inputs=[file_in],
        fn=convert_file,
        outputs=[page_preview, result_out, rendered_out, status_out],
        cache_examples=True,
        cache_mode="lazy",
        label="Example documents (the OmniParseBench pages used in the model card)",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)