"""
PaddleOCR-VL Doc Editor
=======================

Interactive document-parsing workbench:

  PDF / image  ->  layout zones (PP-DocLayout, CPU)
               ->  per-zone recognition (PaddleOCR-VL-0.9B VLM, ZeroGPU)
               ->  click a zone to edit its text
               ->  rebuild Markdown / JSON

The two ML stacks are deliberately kept apart:
  * Layout detection uses PaddlePaddle on CPU (no CUDA -> no conflict with the
    ZeroGPU torch hijack).
  * The VLM recognizer uses transformers on the ZeroGPU GPU.
"""

import spaces  # must be imported before torch / any CUDA-touching import

import io
import json
import os
import tempfile
import traceback

import fitz  # PyMuPDF
import numpy as np
import torch
import gradio as gr
from PIL import Image, ImageDraw, ImageFont

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
MODEL_PATH = "PaddlePaddle/PaddleOCR-VL"
# Standalone layout model exposed through paddleocr.LayoutDetection.
LAYOUT_MODELS = ["PP-DocLayoutV2", "PP-DocLayout_plus-L", "PP-DocLayout-L"]
# Text-line detector (for dense OCR of scattered text / numbers in images).
TEXTDET_MODELS = ["PP-OCRv5_mobile_det", "PP-OCRv5_server_det", "PP-OCRv4_mobile_det"]
MAX_DENSE_ZONES = 120     # cap text-line zones per page to keep GPU time bounded
RENDER_DPI = 150          # PDF rasterization DPI
MAX_SIDE = 2200           # cap the long side of a page (px)
SAMPLE_PATH = "examples/paddleocr_vl_demo.png"

# On ZeroGPU the GPU is not attached at import time, so torch.cuda.is_available()
# is False here; the `spaces` package intercepts `.to("cuda")` at module scope and
# streams the weights in on the first @spaces.GPU call. So trust the ZeroGPU env
# var, not is_available(), when deciding to place weights on cuda.
ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
USE_CUDA = ON_ZEROGPU or torch.cuda.is_available()
DEVICE = "cuda" if USE_CUDA else "cpu"

# Prompt per PaddleOCR-VL task.
PROMPTS = {
    "ocr": "OCR:",
    "table": "Table Recognition:",
    "formula": "Formula Recognition:",
    "chart": "Chart Recognition:",
}

# Max new tokens per task (tables/charts can be long).
MAX_TOKENS = {"ocr": 1024, "table": 1536, "formula": 512, "chart": 1024}

# Free ZeroGPU caps per-call duration (the scheduler scales the declared duration
# up by a safety factor, so keep it modest) AND caps total GPU-seconds/day per
# visitor. Process zones in chunks: few enough allocations to be quota-friendly,
# short enough per call to stay under the duration cap.
GPU_DURATION = 55
CHUNK = 12

# Map a layout label to a recognition task. Anything not listed -> "ocr".
LABEL_TO_TASK = {
    "table": "table",
    "formula": "formula",
    "equation": "formula",
    "chart": "chart",
    "figure": "figure",   # image region, skip text recognition
    "image": "figure",
    "seal": "figure",
    "stamp": "figure",
}

# Distinct-ish colors for zone overlays.
PALETTE = [
    (239, 68, 68), (59, 130, 246), (16, 185, 129), (245, 158, 11),
    (139, 92, 246), (236, 72, 153), (14, 165, 233), (34, 197, 94),
    (249, 115, 22), (168, 85, 247), (20, 184, 166), (99, 102, 241),
]

# ---------------------------------------------------------------------------
# Models (lazy singletons)
# ---------------------------------------------------------------------------
_vlm = {"model": None, "processor": None}
_layout = {"model": None, "name": None}
_textdet = {"model": None, "name": None}


def _patch_transformers_masking():
    """Bridge a kwarg-name mismatch in the model's remote code.

    PaddleOCR-VL's modeling code calls
        create_causal_mask(inputs_embeds=...)
    but transformers 4.x names that parameter `input_embeds` (no 's'). Wrap the
    function so the remote code loads and runs unchanged. Must run before the
    remote modeling module does `from transformers.masking_utils import ...`.
    """
    import functools

    try:
        import transformers.masking_utils as mu
    except Exception as e:  # noqa: BLE001
        print("[patch] masking_utils unavailable:", e, flush=True)
        return
    orig = mu.create_causal_mask
    if getattr(orig, "_paddleocr_shim", False):
        return

    @functools.wraps(orig)
    def shim(*args, **kwargs):
        if "inputs_embeds" in kwargs and "input_embeds" not in kwargs:
            kwargs["input_embeds"] = kwargs.pop("inputs_embeds")
        return orig(*args, **kwargs)

    shim._paddleocr_shim = True
    mu.create_causal_mask = shim
    print("[patch] create_causal_mask kwarg shim installed", flush=True)


def _load_vlm():
    """Load the PaddleOCR-VL VLM at module import so ZeroGPU can pack the weights."""
    if _vlm["model"] is not None:
        return
    _patch_transformers_masking()
    from transformers import AutoModelForCausalLM, AutoProcessor

    print("[vlm] loading", MODEL_PATH, flush=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    ).eval()
    if DEVICE == "cuda":
        model = model.to("cuda")
    processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True)
    _vlm["model"] = model
    _vlm["processor"] = processor
    print("[vlm] ready", flush=True)


def _get_layout():
    """Lazily build the CPU layout detector, trying model names in order."""
    if _layout["model"] is not None:
        return _layout["model"]
    from paddleocr import LayoutDetection

    last_err = None
    for name in LAYOUT_MODELS:
        try:
            print("[layout] trying", name, flush=True)
            _layout["model"] = LayoutDetection(model_name=name)
            _layout["name"] = name
            print("[layout] using", name, flush=True)
            return _layout["model"]
        except Exception as e:  # noqa: BLE001
            last_err = e
            print(f"[layout] {name} failed: {e}", flush=True)
    raise RuntimeError(f"Could not init any layout model: {last_err}")


def _get_textdet():
    """Lazily build the CPU text-line detector (PP-OCR detection)."""
    if _textdet["model"] is not None:
        return _textdet["model"]
    from paddleocr import TextDetection

    last_err = None
    for name in TEXTDET_MODELS:
        try:
            print("[textdet] trying", name, flush=True)
            _textdet["model"] = TextDetection(model_name=name)
            _textdet["name"] = name
            print("[textdet] using", name, flush=True)
            return _textdet["model"]
        except Exception as e:  # noqa: BLE001
            last_err = e
            print(f"[textdet] {name} failed: {e}", flush=True)
    raise RuntimeError(f"Could not init any text-detection model: {last_err}")


# Load the VLM eagerly at import (ZeroGPU best practice).
try:
    _load_vlm()
except Exception as e:  # noqa: BLE001
    print("[vlm] eager load failed, will retry lazily:", e, flush=True)


# ---------------------------------------------------------------------------
# PDF / image helpers
# ---------------------------------------------------------------------------
def _cap_size(img: Image.Image) -> Image.Image:
    w, h = img.size
    m = max(w, h)
    if m > MAX_SIDE:
        s = MAX_SIDE / m
        img = img.resize((int(w * s), int(h * s)), Image.LANCZOS)
    return img


def load_pages(file_path):
    """Return a list of PIL.Image pages from a PDF or single image file."""
    if not file_path:
        return []
    ext = os.path.splitext(file_path)[1].lower()
    pages = []
    if ext == ".pdf":
        doc = fitz.open(file_path)
        zoom = RENDER_DPI / 72.0
        mat = fitz.Matrix(zoom, zoom)
        for page in doc:
            pix = page.get_pixmap(matrix=mat, alpha=False)
            img = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB")
            pages.append(_cap_size(img))
        doc.close()
    else:
        pages.append(_cap_size(Image.open(file_path).convert("RGB")))
    return pages


# ---------------------------------------------------------------------------
# Layout detection (CPU)
# ---------------------------------------------------------------------------
def _finalize(zones):
    """Sort into rough reading order and assign ids."""
    zones.sort(key=lambda z: (z["box"][1], z["box"][0]))
    for i, z in enumerate(zones):
        z["id"] = i
    return zones


def detect_layout_zones(img: Image.Image):
    """PP-DocLayout: region-level zones (text, title, table, formula, chart, figure)."""
    model = _get_layout()
    arr = np.array(img)
    out = model.predict(arr, batch_size=1)
    res = list(out)[0] if not isinstance(out, dict) else out
    boxes = res["boxes"] if "boxes" in res else res.get("boxes", [])

    zones = []
    W, H = img.size
    for b in boxes:
        coord = b.get("coordinate") or b.get("bbox") or b.get("box")
        if coord is None:
            continue
        x1, y1, x2, y2 = [float(v) for v in coord[:4]]
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(W, x2), min(H, y2)
        if x2 - x1 < 3 or y2 - y1 < 3:
            continue
        label = str(b.get("label", "text")).lower().strip()
        task = LABEL_TO_TASK.get(label, "ocr")
        zones.append(
            {
                "box": [int(x1), int(y1), int(x2), int(y2)],
                "label": label,
                "task": task,
                "score": float(b.get("score", 0.0)),
            }
        )
    return _finalize(zones)


def detect_text_line_zones(img: Image.Image):
    """PP-OCR text detection: one zone per detected text line / number."""
    model = _get_textdet()
    arr = np.array(img)
    out = model.predict(arr, batch_size=1)
    res = list(out)[0] if not isinstance(out, dict) else out
    polys = res.get("dt_polys", res.get("polys", []))
    scores = res.get("dt_scores", res.get("scores", []))

    zones = []
    W, H = img.size
    for i, p in enumerate(polys):
        pts = np.array(p).reshape(-1, 2)
        x1, y1 = pts[:, 0].min(), pts[:, 1].min()
        x2, y2 = pts[:, 0].max(), pts[:, 1].max()
        x1, y1 = max(0, float(x1)), max(0, float(y1))
        x2, y2 = min(W, float(x2)), min(H, float(y2))
        bw, bh = x2 - x1, y2 - y1
        if bw < 2 or bh < 2:
            continue
        # A much taller-than-wide box is vertical text -> keep it vertical.
        vertical = bh > bw * 1.8
        angle = 90 if vertical else 0
        # pad along the reading direction to help the recognizer
        if vertical:
            pad = int(bw * 0.15)
            box = [max(0, int(x1 - pad)), int(y1), min(W, int(x2 + pad)), int(y2)]
        else:
            pad = int(bh * 0.15)
            box = [int(x1), max(0, int(y1 - pad)), int(x2), min(H, int(y2 + pad))]
        zones.append(
            {
                "box": box,
                "label": "text-line",
                "task": "ocr",
                "score": float(scores[i]) if i < len(scores) else 0.0,
                "dense": True,
                "angle": angle,
            }
        )
    if len(zones) > MAX_DENSE_ZONES:
        # keep the highest-confidence lines
        zones.sort(key=lambda z: z["score"], reverse=True)
        zones = zones[:MAX_DENSE_ZONES]
    return _finalize(zones)


def detect_zones(img: Image.Image, mode: str = "Document layout"):
    if mode.startswith("Dense"):
        return detect_text_line_zones(img)
    return detect_layout_zones(img)


# ---------------------------------------------------------------------------
# VLM recognition (GPU)
# ---------------------------------------------------------------------------
def _vlm_one(crop: Image.Image, task: str, max_tokens: int) -> str:
    model = _vlm["model"]
    processor = _vlm["processor"]
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": crop},
                {"type": "text", "text": PROMPTS[task]},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)
    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            do_sample=False,
            use_cache=True,
        )
    # Decode only the newly generated tokens, not the prompt echo.
    in_len = inputs["input_ids"].shape[1]
    gen = out[:, in_len:]
    text = processor.batch_decode(gen, skip_special_tokens=True)[0]
    # Belt-and-suspenders: if a chat template still leaks through, keep the tail.
    for marker in ("Assistant:", "assistant\n", "assistant:"):
        if marker in text:
            text = text.split(marker)[-1]
    return text.strip()


@spaces.GPU(duration=GPU_DURATION)
def _recognize_chunk(img: Image.Image, chunk, ocr_figures: bool):
    """Recognize a small batch of zones inside one short GPU allocation.

    Runs in the ZeroGPU worker process, so results are returned (not mutated in
    place) to cross the process boundary.
    """
    _load_vlm()
    for z in chunk:
        task = z["task"]
        if task == "figure":
            if not ocr_figures:
                z["content"] = ""
                continue
            task = "ocr"  # OCR the figure region instead of skipping it
        # Dense single-line boxes need few tokens; regions can need many.
        max_tokens = 256 if z.get("dense") else MAX_TOKENS.get(task, 1024)
        crop = img.crop(tuple(z["box"]))
        # Rotate a vertical crop to horizontal so the VLM reads it correctly; the
        # field/text is placed back at the same angle when the PDF is built.
        if int(z.get("angle", 0) or 0) == 90:
            crop = crop.transpose(Image.ROTATE_90)
        try:
            z["content"] = _vlm_one(crop, task, max_tokens)
        except Exception as e:  # noqa: BLE001
            z["content"] = f"[recognition error: {e}]"
    return chunk


def recognize_all(img: Image.Image, zones, ocr_figures: bool = True):
    """Recognize all zones, chunked across several short GPU allocations.

    ocr_figures: if True, image/figure regions are also OCR'd (pulls out text /
    numbers embedded in diagrams); if False they are left blank.
    """
    for i in range(0, len(zones), CHUNK):
        done = _recognize_chunk(img, zones[i : i + CHUNK], ocr_figures)
        zones[i : i + CHUNK] = done
    return zones


# ---------------------------------------------------------------------------
# Rendering helpers
# ---------------------------------------------------------------------------
def _font(size=22):
    try:
        return ImageFont.truetype("DejaVuSans-Bold.ttf", size)
    except Exception:  # noqa: BLE001
        return ImageFont.load_default()


def render_base(img: Image.Image, zones):
    """Draw numbered rectangles on a copy of the page for display."""
    base = img.convert("RGB").copy()
    draw = ImageDraw.Draw(base)
    f = _font(max(16, img.size[0] // 60))
    for z in zones:
        color = PALETTE[z["id"] % len(PALETTE)]
        x1, y1, x2, y2 = z["box"]
        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
        tag = str(z["id"] + 1)
        tw = draw.textlength(tag, font=f)
        draw.rectangle([x1, y1, x1 + tw + 8, y1 + 26], fill=color)
        draw.text((x1 + 4, y1 + 2), tag, fill=(255, 255, 255), font=f)
    return base


def build_annotations(img: Image.Image, zones):
    """AnnotatedImage value: (base_image, [(mask, label), ...])."""
    base = render_base(img, zones)
    W, H = img.size
    anns = []
    for z in zones:
        mask = np.zeros((H, W), dtype=np.uint8)
        x1, y1, x2, y2 = z["box"]
        mask[y1:y2, x1:x2] = 1
        anns.append((mask, f"{z['id'] + 1}. {z['label']}"))
    return (base, anns)


def zone_choices(zones):
    return [(f"{z['id'] + 1}. {z['label']}", z["id"]) for z in zones]


# ---------------------------------------------------------------------------
# Document rebuild
# ---------------------------------------------------------------------------
def rebuild_markdown(zones):
    ordered = sorted(zones, key=lambda z: (z["box"][1], z["box"][0]))
    parts = []
    for z in ordered:
        c = (z.get("content") or "").strip()
        if z["task"] == "figure" and not c:
            parts.append(f"*[figure — zone {z['id'] + 1}]*")
            continue
        if not c:
            continue
        if z["label"] in ("title", "doc_title", "paragraph_title"):
            parts.append(f"## {c}")
        elif z["task"] == "formula":
            parts.append(f"$$\n{c}\n$$")
        else:
            parts.append(c)
    return "\n\n".join(parts)


# ---------------------------------------------------------------------------
# PDF rebuild
# ---------------------------------------------------------------------------
def _zones_for(zones_pages, i):
    zp = zones_pages or {}
    return zp.get(i) or zp.get(str(i)) or []


def _add_form_field(page, z, text, page_index):
    """Add an editable AcroForm text field at the zone, pre-filled with `text`."""
    x1, y1, x2, y2 = z["box"]
    rect = fitz.Rect(x1, y1, x2, y2)
    # Paint an opaque white rectangle first to hide the original scanned text, so
    # the editable field's text is not shown on top of a duplicate underneath.
    page.draw_rect(rect, fill=(1, 1, 1), color=(0.7, 0.75, 0.85), width=0.5)
    w = fitz.Widget()
    w.rect = rect
    # Unique across the whole document: fields sharing a name are one field in
    # PDF editors, so editing page 1 would also change page 2.
    w.field_name = f"p{page_index + 1}_zone_{z['id'] + 1}"
    w.field_type = fitz.PDF_WIDGET_TYPE_TEXT
    w.field_flags = fitz.PDF_TX_FIELD_IS_MULTILINE
    w.field_value = text
    w.text_fontsize = 0  # auto-size to the box
    w.text_font = "Helv"
    # Transparent field (the white rectangle above already provides the clean
    # background); no border so the box reads as plain white.
    w.fill_color = None
    w.border_color = None
    annot = page.add_widget(w)
    angle = int(z.get("angle", 0) or 0) % 360
    if angle:
        # Keep the original orientation (vertical labels stay vertical). PDF
        # field rotation lives in /MK /R; PyMuPDF's Widget has no rotation
        # setting, so set the key and regenerate the appearance.
        doc = page.parent
        doc.xref_set_key(annot.xref, "MK/R", str(angle))
        placed = next((x for x in page.widgets() if x.xref == annot.xref), None)
        if placed is not None:
            placed.update()
        if "/R" not in (doc.xref_get_key(annot.xref, "MK")[1] or ""):
            doc.xref_set_key(annot.xref, "MK/R", str(angle))


def _finalize_form(doc):
    """Make the fields editable in any PDF editor.

    Editors redraw a field from its /DA font when you type and resolve that font
    through the form's /DR resources; declare it, and ask viewers to regenerate
    appearances so values always render.
    """
    catalog = doc.pdf_catalog()
    if doc.xref_get_key(catalog, "AcroForm/DR/Font/Helv")[0] == "null":
        xref = doc.get_new_xref()
        doc.update_object(
            xref, "<</Type/Font/Subtype/Type1/BaseFont/Helvetica/Encoding/WinAnsiEncoding>>"
        )
        doc.xref_set_key(catalog, "AcroForm/DR/Font/Helv", f"{xref} 0 R")
    if doc.xref_get_key(catalog, "AcroForm/DA")[0] == "null":
        doc.xref_set_key(catalog, "AcroForm/DA", "(/Helv 0 Tf 0 g)")
    try:
        doc.need_appearances(True)
    except Exception:  # noqa: BLE001
        pass


def build_pdf(pages, zones_pages, style="Searchable (over original)"):
    """Rebuild a PDF from the pages + recognized/edited zones.

    Searchable: original page image as background with an invisible, selectable
    text layer placed at each zone (looks identical to the input, but the text is
    now real and reflects your edits).
    Editable form: original page image with a fillable AcroForm text field at each
    zone, pre-filled with the recognized text — edit it in any PDF viewer.
    Text only: white pages with the recognized text drawn at each zone's position.
    """
    searchable = style.startswith("Searchable")
    form = style.startswith("Editable")
    doc = fitz.open()
    for i, img in enumerate(pages):
        W, H = img.size
        page = doc.new_page(width=W, height=H)
        if searchable or form:
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            page.insert_image(fitz.Rect(0, 0, W, H), stream=buf.getvalue())
        for z in _zones_for(zones_pages, i):
            text = (z.get("content") or "").strip()
            if not text:
                continue
            if form:
                try:
                    _add_form_field(page, z, text, i)
                except Exception:  # noqa: BLE001
                    traceback.print_exc()
                continue
            x1, y1, x2, y2 = z["box"]
            rect = fitz.Rect(x1, y1, x2, y2)
            bh = max(1, y2 - y1)
            if z.get("dense"):
                fs = max(6, min(int(bh * 0.7), 40))
            else:
                fs = 10
            render_mode = 3 if searchable else 0  # 3 = invisible (search layer)
            rotate = int(z.get("angle", 0) or 0)
            for fname in ("china-s", None):
                try:
                    kw = {"fontname": fname} if fname else {}
                    page.insert_textbox(
                        rect, text, fontsize=fs, render_mode=render_mode,
                        align=0, rotate=rotate, **kw,
                    )
                    break
                except Exception:  # noqa: BLE001
                    continue
    if form and any(_zones_for(zones_pages, i) for i in range(len(pages))):
        _finalize_form(doc)
    out = os.path.join(tempfile.mkdtemp(), "document.pdf")
    doc.save(out, deflate=True, garbage=3)
    doc.close()
    return out


# ---------------------------------------------------------------------------
# Gradio callbacks
# ---------------------------------------------------------------------------
def on_upload(file_path):
    pages = load_pages(file_path)
    if not pages:
        return (
            gr.update(choices=[], value=None),
            [],
            {},
            [],
            None,
            gr.update(choices=[], value=None),
            "",
            "",
            "Upload a PDF or image to begin.",
        )
    choices = [(f"Page {i + 1}", i) for i in range(len(pages))]
    first = pages[0]
    return (
        gr.update(choices=choices, value=0),
        pages,
        {},
        [],
        (first, []),
        gr.update(choices=[], value=None),
        "",
        "",
        f"Loaded {len(pages)} page(s). Click **Process pages & export PDF** to do them all, "
        "or pick a page and click **Detect & recognize this page**.",
    )


def load_sample():
    if not os.path.exists(SAMPLE_PATH):
        return on_upload(None)
    return on_upload(SAMPLE_PATH)


def on_page_change(page_idx, pages, zones_pages):
    if not pages or page_idx is None or int(page_idx) >= len(pages):
        return (None, [], gr.update(choices=[], value=None), "", "")
    i = int(page_idx)
    img = pages[i]
    zones = _zones_for(zones_pages, i)
    if zones:
        return (
            build_annotations(img, zones),
            zones,
            gr.update(choices=zone_choices(zones), value=zones[0]["id"]),
            zones[0]["label"],
            zones[0].get("content", ""),
        )
    return ((img, []), [], gr.update(choices=[], value=None), "", "")


def _ensure_pages(pages, file_path):
    """Pages from session state, or straight from the uploaded file.

    Clicking a button right after uploading can run before the upload handler has
    stored the pages, so fall back to reading the file itself.
    """
    if pages:
        return pages
    if file_path:
        loaded = load_pages(file_path)
        if loaded:
            return loaded
    raise gr.Error("Upload a document (or click 'Load sample') first.")


def _page_choices(pages, page_idx):
    return gr.update(choices=[(f"Page {i + 1}", i) for i in range(len(pages))], value=page_idx)


def analyze(file_path, page_idx, pages, mode, ocr_figures, zones_pages):
    pages = _ensure_pages(pages, file_path)
    if page_idx is None or int(page_idx) >= len(pages):
        page_idx = 0
    page_idx = int(page_idx)
    img = pages[page_idx]
    try:
        zones = detect_zones(img, mode)
    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        raise gr.Error(f"Detection failed: {e}")
    zp = dict(zones_pages or {})
    if not zones:
        zp[page_idx] = []
        return (
            pages,
            _page_choices(pages, page_idx),
            (img, []),
            [],
            zp,
            gr.update(choices=[], value=None),
            "",
            "",
            "No zones detected on this page. Try the other detection mode.",
        )
    zones = recognize_all(img, zones, bool(ocr_figures))
    zp[page_idx] = zones
    ann = build_annotations(img, zones)
    status = (
        f"Recognized {len(zones)} zone(s) on page {page_idx + 1}. Click a box to "
        "edit, then **Export PDF**."
    )
    return (
        pages,
        _page_choices(pages, page_idx),
        ann,
        zones,
        zp,
        gr.update(choices=zone_choices(zones), value=zones[0]["id"]),
        zones[0]["label"],
        zones[0].get("content", ""),
        status,
    )


def _select_zone(zones, zid):
    for z in zones:
        if z["id"] == zid:
            return z["label"], z.get("content", "")
    return "", ""


def on_box_select(zones, evt: gr.SelectData):
    if not zones:
        return gr.update(), "", ""
    idx = evt.index
    if isinstance(idx, (list, tuple)):
        idx = idx[0]
    # annotations are built in zone-id order
    if idx is None or idx < 0 or idx >= len(zones):
        return gr.update(), "", ""
    z = zones[idx]
    label, content = _select_zone(zones, z["id"])
    return gr.update(value=z["id"]), label, content


def on_dropdown(zones, zid):
    if zid is None or not zones:
        return "", ""
    return _select_zone(zones, zid)


def save_edit(zones, zid, text, page_idx, zones_pages):
    if zid is None or not zones:
        return zones, zones_pages, "Select a zone first."
    for z in zones:
        if z["id"] == zid:
            z["content"] = text
            break
    else:
        return zones, zones_pages, "Zone not found."
    zp = dict(zones_pages or {})
    if page_idx is not None:
        zp[int(page_idx)] = zones
    return zones, zp, f"Saved zone {zid + 1}. Click **Export PDF** when done."


def api_parse_sample(mode: str = "Document layout", ocr_figures: bool = True):
    """API-only: parse the bundled sample end to end and return the result.

    Lets the pipeline be smoke-tested (and called via API/MCP) without a file
    upload. Returns detected zones with their recognized content and the rebuilt
    Markdown.
    """
    pages = load_pages(SAMPLE_PATH)
    if not pages:
        return {"error": "sample not found"}
    img = pages[0]
    zones = detect_zones(img, mode)
    zones = recognize_all(img, zones, bool(ocr_figures))
    return {
        "mode": mode,
        "n_zones": len(zones),
        "markdown": rebuild_markdown(zones),
        "zones": [
            {"id": z["id"], "label": z["label"], "box": z["box"],
             "content": z.get("content", "")}
            for z in zones
        ],
    }


def export_pdf(pages, zones_pages, style):
    if not pages:
        return None, "Upload a document first."
    if not any(_zones_for(zones_pages, i) for i in range(len(pages))):
        return None, "No recognized zones yet — click **Process pages & export PDF** first."
    try:
        path = build_pdf(pages, zones_pages, style)
    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        raise gr.Error(f"PDF export failed: {e}")
    n = sum(1 for i in range(len(pages)) if _zones_for(zones_pages, i))
    return path, f"PDF ready ({n} recognized page(s)). Download below. ⬇️"


def parse_page_spec(spec: str, page_count: int) -> list[int]:
    """'all' / '' -> every page; '1-5, 8, 10-12' -> those pages (0-based, sorted)."""
    spec = (spec or "").strip().lower()
    if spec in ("", "all", "*"):
        return list(range(page_count))
    picked: set[int] = set()
    for part in spec.replace(";", ",").split(","):
        part = part.strip()
        if not part:
            continue
        try:
            if "-" in part:
                a, b = part.split("-", 1)
                start, end = int(a), int(b)
                if start > end:
                    start, end = end, start
            else:
                start = end = int(part)
        except ValueError:
            raise gr.Error(f"Could not read the page list near '{part}'. Use e.g. all, 1-5, 8")
        picked.update(range(max(1, start), min(page_count, end) + 1))
    if not picked:
        raise gr.Error(f"No pages in range — this document has {page_count} page(s).")
    return sorted(p - 1 for p in picked)


def _is_quota_error(exc: Exception) -> bool:
    msg = str(exc).lower()
    return "quota" in msg or ("zerogpu" in msg and "limit" in msg)


def analyze_pages(file_path, page_spec, pages, mode, ocr_figures, zones_pages, style, page_idx,
                  progress=gr.Progress()):
    """Detect + recognize several pages in one run, then export one PDF."""
    pages = _ensure_pages(pages, file_path)
    if page_idx is None or int(page_idx) >= len(pages):
        page_idx = 0
    targets = parse_page_spec(page_spec, len(pages))
    zp = dict(zones_pages or {})
    done, failed, stopped_at = [], [], None
    for k, i in enumerate(targets):
        progress((k, len(targets)), desc=f"Page {i + 1} ({k + 1} of {len(targets)})")
        img = pages[i]
        try:
            zones = detect_zones(img, mode)
            zones = recognize_all(img, zones, bool(ocr_figures)) if zones else []
        except Exception as e:  # noqa: BLE001
            traceback.print_exc()
            if _is_quota_error(e):
                # Out of free GPU time: keep what finished and export that.
                stopped_at = i + 1
                break
            failed.append(i + 1)
            continue
        zp[i] = zones
        done.append(i + 1)
    progress((len(targets), len(targets)), desc="Building PDF")

    path = None
    if any(_zones_for(zp, i) for i in range(len(pages))):
        path = build_pdf(pages, zp, style)

    parts = [f"Processed {len(done)} of {len(targets)} page(s)."]
    if path:
        parts.append("The PDF below has every page; processed pages have editable fields. ⬇️")
    if failed:
        parts.append(f"Could not read page(s) {', '.join(map(str, failed))} — left unchanged.")
    if stopped_at:
        parts.append(
            f"Stopped at page {stopped_at}: the free GPU quota ran out. Sign in to Hugging "
            "Face for more quota, then process the remaining pages — finished pages are kept."
        )
    if not path:
        parts.append("No text was found to export.")

    # Refresh the page currently on screen so its new boxes are clickable.
    view = on_page_change(page_idx, pages, zp)
    return (pages, _page_choices(pages, int(page_idx)), zp, path, " ".join(parts), *view)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
LATEX_DELIMS = [
    {"left": "$$", "right": "$$", "display": True},
    {"left": "$", "right": "$", "display": False},
]

with gr.Blocks(title="PaddleOCR-VL Doc Editor", theme=gr.themes.Soft()) as demo:
    pages_state = gr.State([])
    zones_state = gr.State([])
    zones_pages = gr.State({})

    gr.Markdown(
        "# 📝 PaddleOCR-VL Doc Editor\n"
        "Upload a **PDF or image** → detected zones are recognized by "
        "[PaddleOCR-VL](https://huggingface.co/PaddlePaddle/PaddleOCR-VL) → "
        "**click any box to edit** its text → **export a corrected PDF**."
    )

    with gr.Row():
        with gr.Column(scale=3):
            file_in = gr.File(
                label="Upload PDF or image",
                file_types=[".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".webp"],
                type="filepath",
            )
            with gr.Row():
                page_sel = gr.Dropdown(label="Page", choices=[], interactive=True, scale=2)
                sample_btn = gr.Button("Load sample", scale=1)
                analyze_btn = gr.Button("Detect & recognize this page", variant="primary", scale=3)
            with gr.Row():
                page_spec = gr.Textbox(
                    label="Pages to process",
                    value="all",
                    placeholder="all, or e.g. 1-5, 8, 10-12",
                    info="Process several pages in one run and download one PDF.",
                    scale=2,
                )
                batch_btn = gr.Button("⚡ Process pages & export PDF", variant="primary", scale=3)
            with gr.Row():
                mode = gr.Radio(
                    choices=["Document layout", "Dense text lines"],
                    value="Document layout",
                    label="Detection mode",
                    info=(
                        "Document layout = region zones (text/table/formula/chart) for "
                        "structured docs. Dense text lines = one box per text/number, "
                        "best for diagrams, drawings and scattered labels."
                    ),
                    scale=3,
                )
                ocr_figures = gr.Checkbox(
                    value=True,
                    label="OCR image/figure regions too",
                    scale=1,
                )
            annotated = gr.AnnotatedImage(
                label="Detected zones (click a box to edit)",
                height=640,
            )
        with gr.Column(scale=2):
            status = gr.Markdown("Upload a PDF or image to begin.")
            zone_dd = gr.Dropdown(label="Zone", choices=[], interactive=True)
            zone_label = gr.Textbox(label="Zone type", interactive=False)
            zone_editor = gr.Textbox(label="Recognized content (editable)", lines=12)
            save_btn = gr.Button("💾 Save edit")
            pdf_style = gr.Radio(
                choices=[
                    "Editable form (fillable fields)",
                    "Searchable (over original)",
                    "Text only (white pages)",
                ],
                value="Editable form (fillable fields)",
                label="PDF output",
                info=(
                    "Editable form = a fillable field at each zone you edit directly "
                    "in any PDF viewer. Searchable = original pages with a corrected "
                    "text layer. Text only = white pages with the recognized text."
                ),
            )
            export_btn = gr.Button("📄 Export PDF", variant="primary")
            dl_pdf = gr.File(label="PDF download")

    # --- wiring ---
    up_outputs = [
        page_sel, pages_state, zones_pages, zones_state, annotated,
        zone_dd, zone_label, zone_editor, status,
    ]
    file_in.change(on_upload, inputs=file_in, outputs=up_outputs)
    sample_btn.click(load_sample, inputs=None, outputs=up_outputs)
    page_sel.change(
        on_page_change,
        inputs=[page_sel, pages_state, zones_pages],
        outputs=[annotated, zones_state, zone_dd, zone_label, zone_editor],
    )
    analyze_btn.click(
        analyze,
        inputs=[file_in, page_sel, pages_state, mode, ocr_figures, zones_pages],
        outputs=[pages_state, page_sel, annotated, zones_state, zones_pages, zone_dd, zone_label,
                 zone_editor, status],
    )
    annotated.select(
        on_box_select,
        inputs=[zones_state],
        outputs=[zone_dd, zone_label, zone_editor],
    )
    zone_dd.change(
        on_dropdown,
        inputs=[zones_state, zone_dd],
        outputs=[zone_label, zone_editor],
    )
    save_btn.click(
        save_edit,
        inputs=[zones_state, zone_dd, zone_editor, page_sel, zones_pages],
        outputs=[zones_state, zones_pages, status],
    )
    export_btn.click(
        export_pdf,
        inputs=[pages_state, zones_pages, pdf_style],
        outputs=[dl_pdf, status],
    )
    batch_btn.click(
        analyze_pages,
        inputs=[file_in, page_spec, pages_state, mode, ocr_figures, zones_pages, pdf_style, page_sel],
        outputs=[pages_state, page_sel, zones_pages, dl_pdf, status, annotated, zones_state, zone_dd,
                 zone_label, zone_editor],
    )

    # Hidden API endpoint: parse the bundled sample (for smoke tests / API use).
    with gr.Row(visible=False):
        api_btn = gr.Button("api")
        api_out = gr.JSON()
    api_btn.click(
        api_parse_sample,
        inputs=[mode, ocr_figures],
        outputs=api_out,
        api_name="parse_sample",
    )

    demo.queue(default_concurrency_limit=1, max_size=16)


if __name__ == "__main__":
    # SSR (experimental) proxies requests through a Node server; it intermittently
    # returns 502 on file uploads, so serve the app directly.
    demo.launch(ssr_mode=False)