"""
OCR Compare -- classic PaddleOCR vs PaddleOCR-VL, side by side
==============================================================

  PDF / image page
     ├─► classic PaddleOCR (PP-OCRv5 det + rec, CPU)   ─┐
     └─► PaddleOCR-VL-0.9B (VLM, ZeroGPU)              ─┤  same {text, box} shape
                                                        ▼
         two page panels with invisible click-to-edit boxes + a diff table

Everything runs on Hugging Face: PaddlePaddle on the Space CPU, the VLM on
ZeroGPU. The box overlay and comparison live in overlay.py.

The VLM loading code (version pins, kwarg shim, chunked GPU calls, prompt
slicing) is carried over from the PaddleOCR-VL-Doc-Editor Space, where each of
those was needed to make it run on ZeroGPU at all.
"""

try:
    import spaces  # must be imported before torch / any CUDA-touching import
except ImportError:  # local run without the HF `spaces` package

    class spaces:  # noqa: N801
        @staticmethod
        def GPU(*args, **kwargs):
            if args and callable(args[0]):
                return args[0]
            return lambda fn: fn


import io
import os
import time
import traceback

import fitz  # PyMuPDF
import gradio as gr
import numpy as np
from PIL import Image

from overlay import (
    CSS,
    HEAD_JS,
    TABLE_HEADERS,
    compare,
    image_to_data_uri,
    overlay_html,
    placeholder_html,
    quad_to_box,
    reading_order,
    table_rows,
)
from pdf_export import build_edited_pdf, parse_edits

try:
    import torch
except ImportError:  # local UI testing without the VLM stack
    torch = None

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
VL_MODEL = "PaddlePaddle/PaddleOCR-VL"
CLASSIC_MODELS = {
    "mobile (fast)": ("PP-OCRv5_mobile_det", "PP-OCRv5_mobile_rec"),
    "server (accurate)": ("PP-OCRv5_server_det", "PP-OCRv5_server_rec"),
}
LAYOUT_MODELS = ["PP-DocLayoutV2", "PP-DocLayout_plus-L", "PP-DocLayout-L"]

VL_SAME_LINES = "Same text lines as classic (fair comparison)"
VL_LAYOUT = "VL's own layout zones (text / table / formula)"

RENDER_DPI = 150
MAX_SIDE = 2200          # cap the long side of a rendered page (px)
MAX_PAGES = 30
MAX_VL_ITEMS = 150       # bound GPU time per page
SAMPLE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples", "sample.png")

# Free ZeroGPU caps each call's duration and each visitor's GPU-seconds per day,
# so recognition runs in several short allocations rather than one long one.
GPU_DURATION = 55
CHUNK_LINES = 24         # short single-line crops
CHUNK_ZONES = 10         # whole regions (can be long)

PROMPTS = {"ocr": "OCR:", "table": "Table Recognition:",
           "formula": "Formula Recognition:", "chart": "Chart Recognition:"}
MAX_TOKENS = {"ocr": 1024, "table": 1536, "formula": 512, "chart": 1024}
LABEL_TO_TASK = {"table": "table", "formula": "formula", "equation": "formula",
                 "chart": "chart"}

ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
USE_CUDA = torch is not None and (ON_ZEROGPU or torch.cuda.is_available())
DEVICE = "cuda" if USE_CUDA else "cpu"

# ---------------------------------------------------------------------------
# Models (lazy singletons)
# ---------------------------------------------------------------------------
_classic: dict = {}
_layout = {"model": None}
_vlm = {"model": None, "processor": None, "error": None}


def _get_classic(size: str):
    if size in _classic:
        return _classic[size]
    from paddleocr import PaddleOCR

    det, rec = CLASSIC_MODELS[size]
    print(f"[classic] loading {det} + {rec}", flush=True)
    engine = PaddleOCR(
        text_detection_model_name=det,
        text_recognition_model_name=rec,
        # Keep the page exactly as rendered so boxes line up with the image.
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,
    )
    _classic[size] = engine
    return engine


def _get_layout():
    if _layout["model"] is not None:
        return _layout["model"]
    from paddleocr import LayoutDetection

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


def _patch_transformers_masking():
    """PaddleOCR-VL's remote code calls create_causal_mask(inputs_embeds=...),
    but transformers 4.55 names that parameter `input_embeds`. Rename it."""
    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


def _load_vlm():
    if _vlm["model"] is not None:
        return
    if torch is None:
        raise RuntimeError("PyTorch is not installed, so PaddleOCR-VL is unavailable here.")
    _patch_transformers_masking()
    from transformers import AutoModelForCausalLM, AutoProcessor

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


# Load the VLM at import so ZeroGPU can pack the weights (ZeroGPU best practice).
try:
    _load_vlm()
except Exception as e:  # noqa: BLE001
    _vlm["error"] = str(e)
    print("[vlm] eager load failed, will retry on first use:", e, flush=True)

# ---------------------------------------------------------------------------
# Pages
# ---------------------------------------------------------------------------


def _cap(img: Image.Image) -> Image.Image:
    w, h = img.size
    if max(w, h) > MAX_SIDE:
        s = MAX_SIDE / max(w, h)
        img = img.resize((int(w * s), int(h * s)), Image.LANCZOS)
    return img


def load_pages(path):
    if not path:
        return []
    if os.path.splitext(path)[1].lower() == ".pdf":
        doc = fitz.open(path)
        zoom = RENDER_DPI / 72.0
        pages = []
        for i, page in enumerate(doc):
            if i >= MAX_PAGES:
                break
            pix = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
            pages.append(_cap(Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB")))
        doc.close()
        return pages
    return [_cap(Image.open(path).convert("RGB"))]


# ---------------------------------------------------------------------------
# Engine 1: classic PaddleOCR (CPU)
# ---------------------------------------------------------------------------


def _result_dict(res):
    """PaddleOCR 3.x results are dict-like; .json gives plain lists."""
    try:
        j = res.json
        if isinstance(j, dict):
            return j.get("res", j)
    except Exception:  # noqa: BLE001
        pass
    return res


def run_classic(img: Image.Image, size: str):
    """-> (items, detected_boxes). detected_boxes includes lines the recognizer
    returned empty, so the VL 'same lines' mode sees every detection."""
    engine = _get_classic(size)
    bgr = np.ascontiguousarray(np.array(img.convert("RGB"))[:, :, ::-1])  # Paddle expects BGR
    res = _result_dict(list(engine.predict(bgr))[0])
    texts = list(res.get("rec_texts") or [])
    scores = list(res.get("rec_scores") or [])
    polys = list(res.get("rec_polys") or [])
    items = []
    for i, text in enumerate(texts):
        if i >= len(polys) or not str(text).strip():
            continue
        items.append({
            "text": str(text).strip(),
            "box": quad_to_box(polys[i]),
            "conf": float(scores[i]) * 100 if i < len(scores) else None,
            "label": "text line",
        })
    det = [quad_to_box(p) for p in (res.get("dt_polys") or polys)]
    return reading_order(items), det


# ---------------------------------------------------------------------------
# Engine 2: PaddleOCR-VL (ZeroGPU)
# ---------------------------------------------------------------------------


def _vlm_one(crop: Image.Image, task: str, max_tokens: int) -> str:
    model, processor = _vlm["model"], _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 new tokens; otherwise the prompt is echoed back.
    text = processor.batch_decode(out[:, inputs["input_ids"].shape[1]:],
                                  skip_special_tokens=True)[0]
    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: list[dict]) -> list[dict]:
    """One short GPU allocation. Returns the chunk (ZeroGPU runs this in a
    separate process, so in-place mutation would not come back)."""
    _load_vlm()
    W, H = img.size
    for z in chunk:
        x1, y1, x2, y2 = z["box"]
        if z.get("line"):
            # a little vertical padding helps the recognizer on tight line boxes
            pad = max(2, int((y2 - y1) * 0.15))
            y1, y2 = max(0, y1 - pad), min(H, y2 + pad)
            max_tokens = 128
        else:
            max_tokens = MAX_TOKENS.get(z["task"], 1024)
        try:
            z["text"] = _vlm_one(img.crop((x1, y1, x2, y2)), z["task"], max_tokens)
        except Exception as e:  # noqa: BLE001
            z["text"] = f"[recognition error: {e}]"
    return chunk


def _layout_zones(img: Image.Image) -> list[dict]:
    out = _get_layout().predict(np.array(img), batch_size=1)
    res = _result_dict(list(out)[0])
    W, H = img.size
    zones = []
    for b in res.get("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, x2, y2 = max(0, x1), max(0, y1), min(W, x2), min(H, y2)
        if x2 - x1 < 3 or y2 - y1 < 3:
            continue
        label = str(b.get("label", "text")).lower().strip()
        zones.append({"box": [int(x1), int(y1), int(x2), int(y2)], "label": label,
                      "task": LABEL_TO_TASK.get(label, "ocr")})
    return reading_order(zones)


def run_vl(img: Image.Image, mode: str, line_boxes: list[list[int]]):
    if mode == VL_SAME_LINES:
        zones = [{"box": b, "label": "text line", "task": "ocr", "line": True}
                 for b in line_boxes]
        chunk = CHUNK_LINES
    else:
        zones = _layout_zones(img)
        chunk = CHUNK_ZONES
    zones = reading_order(zones)[:MAX_VL_ITEMS]
    for i in range(0, len(zones), chunk):
        zones[i : i + chunk] = _recognize_chunk(img, zones[i : i + chunk])
    return [{"text": z.get("text", ""), "box": z["box"], "conf": None, "label": z["label"]}
            for z in zones]


# ---------------------------------------------------------------------------
# Gradio callbacks
# ---------------------------------------------------------------------------

EMPTY_TABLE = gr.update(value=[], headers=TABLE_HEADERS)


def on_upload(path):
    pages = load_pages(path)
    if not pages:
        return (gr.update(choices=[], value=None), [], None, placeholder_html("Upload a PDF or image."),
                placeholder_html("Upload a PDF or image."), [], EMPTY_TABLE, "", "[]", None)
    choices = [(f"Page {i + 1}", i) for i in range(len(pages))]
    msg = f"Loaded {len(pages)} page(s). Pick a page and click **Run both**."
    ready = placeholder_html("Click “Run both” to OCR this page.")
    return (gr.update(choices=choices, value=0), pages, path, ready, ready, [], EMPTY_TABLE, msg,
            "[]", None)


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


def _summary(t_classic, n_classic, t_vl, n_vl, stats, vl_mode, note=""):
    lines = [f"**PaddleOCR (classic):** {n_classic} lines in {t_classic:.1f}s"]
    if t_vl is not None:
        mode = "same lines" if vl_mode == VL_SAME_LINES else "layout zones"
        lines.append(f"**PaddleOCR-VL:** {n_vl} boxes ({mode}) in {t_vl:.1f}s")
    if stats:
        c = stats["counts"]
        mean = f"{stats['mean_ratio'] * 100:.1f}%" if stats["mean_ratio"] is not None else "—"
        lines.append(
            f"**Agreement:** {c['same']} same · {c['minor']} minor · {c['differs']} differ · "
            f"{c['only-classic']} only classic · {c['only-vl']} only VL "
            f"(mean similarity {mean})"
        )
    if note:
        lines.append(note)
    return "  \n".join(lines)


def run_both(page_idx, pages, classic_size, vl_mode, use_vl, only_diff):
    if not pages:
        raise gr.Error("Upload a document (or click “Load sample”) first.")
    pno = int(page_idx or 0)
    img = pages[pno]
    uri = image_to_data_uri(img)  # encode once, reuse for both panels
    waiting = placeholder_html("Waiting for classic PaddleOCR…")

    yield (placeholder_html("Running classic PaddleOCR on CPU…"), waiting, [],
           EMPTY_TABLE, "Running classic PaddleOCR…")

    t0 = time.perf_counter()
    try:
        classic, det_boxes = run_classic(img, classic_size)
    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        raise gr.Error(f"Classic PaddleOCR failed: {e}")
    t_classic = time.perf_counter() - t0
    classic_html = overlay_html(img, classic, "classic", uri, pno)

    if not use_vl:
        yield (classic_html, placeholder_html("PaddleOCR-VL is switched off."), [],
               EMPTY_TABLE, _summary(t_classic, len(classic), None, 0, None, vl_mode))
        return

    if vl_mode == VL_SAME_LINES and not det_boxes:
        yield (classic_html, placeholder_html("Classic found no text lines to re-read."),
               [], EMPTY_TABLE, _summary(t_classic, len(classic), None, 0, None, vl_mode))
        return

    yield (classic_html, placeholder_html("Running PaddleOCR-VL on ZeroGPU…"), [],
           EMPTY_TABLE,
           _summary(t_classic, len(classic), None, 0, None, vl_mode) + "  \nRunning PaddleOCR-VL…")

    t0 = time.perf_counter()
    try:
        vl = run_vl(img, vl_mode, det_boxes)
    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        yield (classic_html, placeholder_html(f"PaddleOCR-VL failed: {e}"), [], EMPTY_TABLE,
               _summary(t_classic, len(classic), None, 0, None, vl_mode))
        return
    t_vl = time.perf_counter() - t0

    rows, stats = compare(classic, vl)  # also tags boxes with their diff row
    note = ""
    if vl_mode == VL_SAME_LINES and len(det_boxes) > MAX_VL_ITEMS:
        note = f"_VL re-read the first {MAX_VL_ITEMS} of {len(det_boxes)} lines (GPU budget)._"
    yield (
        overlay_html(img, classic, "classic", uri, pno),
        overlay_html(img, vl, "vl", uri, pno),
        rows,
        gr.update(value=table_rows(rows, only_diff), headers=TABLE_HEADERS),
        _summary(t_classic, len(classic), t_vl, len(vl), stats, vl_mode, note),
    )


def on_filter(rows, only_diff):
    return gr.update(value=table_rows(rows or [], only_diff), headers=TABLE_HEADERS)


def export_pdf(edits_json, src_path, pages):
    if not src_path or not pages:
        raise gr.Error("Upload a document (or click “Load sample”) first.")
    edits = parse_edits(edits_json)
    if not edits:
        raise gr.Error("No saved edits yet. Click a box, change the text, then press Enter or Save.")
    try:
        path, applied, skipped = build_edited_pdf(src_path, pages, edits)
    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        raise gr.Error(f"Could not build the PDF: {e}")
    msg = f"Edited PDF ready: {applied} change(s) applied"
    if skipped:
        msg += f", {skipped} skipped (text too long for its box, or on an unloaded page)"
    return path, msg + "."


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="OCR Compare", theme=gr.themes.Soft(), css=CSS, head=HEAD_JS) as demo:
    pages_state = gr.State([])
    rows_state = gr.State([])
    src_state = gr.State(None)

    gr.Markdown(
        "# 🔍 OCR Compare — PaddleOCR vs PaddleOCR-VL\n"
        "Upload a scan, run both engines, and compare them side by side. "
        "**Every word or line is an invisible box — click one to edit its text, then "
        "download the edited PDF.** The other engine's reading for the same spot is "
        "shown under the editor."
    )

    with gr.Row():
        file_in = gr.File(label="PDF or image", type="filepath",
                          file_types=[".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tif", ".tiff"],
                          scale=3)
        with gr.Column(scale=2):
            page_sel = gr.Dropdown(label="Page", choices=[], interactive=True)
            with gr.Row():
                sample_btn = gr.Button("Load sample")
                run_btn = gr.Button("Run both", variant="primary")

    with gr.Accordion("Options", open=False):
        with gr.Row():
            classic_size = gr.Radio(list(CLASSIC_MODELS), value="mobile (fast)",
                                    label="Classic PaddleOCR models")
            vl_mode = gr.Radio([VL_SAME_LINES, VL_LAYOUT], value=VL_SAME_LINES,
                               label="PaddleOCR-VL boxes",
                               info="Same lines isolates recognition quality; layout "
                                    "zones shows how VL would read the page on its own.")
            use_vl = gr.Checkbox(value=True, label="Run PaddleOCR-VL (uses ZeroGPU quota)")

    with gr.Row():
        show_boxes = gr.Checkbox(value=False, label="Show all boxes (colored by agreement)")
        show_text = gr.Checkbox(value=False, label="Show recognized text on the page")

    status = gr.Markdown("Upload a PDF or image, or click **Load sample**.")

    with gr.Row(equal_height=False):
        with gr.Column():
            gr.Markdown("### PaddleOCR (classic)")
            classic_panel = gr.HTML(placeholder_html("No page yet."), padding=False)
        with gr.Column():
            gr.Markdown("### PaddleOCR-VL")
            vl_panel = gr.HTML(placeholder_html("No page yet."), padding=False)

    with gr.Row(equal_height=True):
        with gr.Column(scale=3):
            gr.HTML('<span id="ocrc-count">No edits yet</span>', padding=False)
            export_msg = gr.Markdown("")
        export_btn = gr.Button("⬇️ Download edited PDF", variant="primary", scale=1)
        pdf_out = gr.File(label="Edited PDF", scale=2)
    # Filled by the page script (overlay.HEAD_JS) with the saved edits as JSON;
    # hidden with CSS rather than visible=False so it stays in the DOM.
    edits_box = gr.Textbox("[]", elem_id="ocrc-edits", show_label=False, container=False)

    gr.Markdown(
        "### Differences\n"
        "Green = identical · yellow = minor · red = different · blue/purple = found by "
        "one engine only. Row numbers match the “diff row” shown when you click a box."
    )
    only_diff = gr.Checkbox(value=True, label="Only show rows that differ")
    diff_table = gr.Dataframe(headers=TABLE_HEADERS, value=[], wrap=True, interactive=False,
                              column_widths=["6%", "14%", "8%", "36%", "36%"])

    # --- wiring ---
    upload_out = [page_sel, pages_state, src_state, classic_panel, vl_panel, rows_state,
                  diff_table, status, edits_box, pdf_out]
    clear_edits_js = "() => { window.__ocrcClear && window.__ocrcClear(); }"
    file_in.change(on_upload, inputs=file_in, outputs=upload_out)
    sample_btn.click(load_sample, inputs=None, outputs=upload_out)
    # A new document starts with no edits (browser-side store).
    file_in.change(None, js=clear_edits_js)
    sample_btn.click(None, js=clear_edits_js)
    export_btn.click(export_pdf, inputs=[edits_box, src_state, pages_state],
                     outputs=[pdf_out, export_msg], concurrency_limit=4)

    def _page_changed(_):
        ready = placeholder_html("Click “Run both” to OCR this page.")
        return ready, ready, [], EMPTY_TABLE

    page_sel.change(_page_changed, inputs=page_sel,
                    outputs=[classic_panel, vl_panel, rows_state, diff_table])
    run_btn.click(
        run_both,
        inputs=[page_sel, pages_state, classic_size, vl_mode, use_vl, only_diff],
        outputs=[classic_panel, vl_panel, rows_state, diff_table, status],
    )
    only_diff.change(on_filter, inputs=[rows_state, only_diff], outputs=diff_table)

    # Display toggles are pure CSS classes on <body>; no server round trip.
    show_boxes.change(None, inputs=show_boxes, outputs=None,
                      js="(v) => { document.body.classList.toggle('ocrc-show-boxes', v); }")
    show_text.change(None, inputs=show_text, outputs=None,
                     js="(v) => { document.body.classList.toggle('ocrc-show-text', v); }")

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


if __name__ == "__main__":
    demo.launch()