"""
Small OCR Comparison — gr.Workflow on ZeroGPU
=============================================

    [image] ─┬─▶ (fn) lighton_ocr ──▶ 📝 LightOnOCR-3-0.8B markdown
             ├─▶ (fn) typhoon_ocr ──▶ 📝 typhoon-ocr1.5-2b markdown
             ├─▶ (fn) nanonets_ocr ─▶ 📝 Nanonets-OCR2-3B markdown
             └─▶ (fn) paddle_ocr ───▶ 📝 PaddleOCR-VL-1.6 markdown

Each operator is a `fn` node bound to a Python function. Models are lazy-loaded on
CPU at first use (cached for the rest of the process) and moved to CUDA inside the
`@spaces.GPU` call, so the ZeroGPU slice only pays for inference, not downloads.
Each model runs exactly the inference recipe from its model card.
"""

import spaces  # must be imported before torch touches CUDA

import base64
import io
import os

import gradio as gr
import torch
from PIL import Image
from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    Qwen3_5ForConditionalGeneration,
)

LIGHTON_ID = "lightonai/LightOnOCR-3-0.8B"
TYPHOON_ID = "typhoon-ai/typhoon-ocr1.5-2b"
NANONETS_ID = "nanonets/Nanonets-OCR2-3B"
PADDLE_ID = "PaddlePaddle/PaddleOCR-VL-1.6"

MODEL_IDS = {
    "lighton": LIGHTON_ID,
    "typhoon": TYPHOON_ID,
    "nanonets": NANONETS_ID,
    "paddle": PADDLE_ID,
}

# Prompts are copied verbatim from each model card — all four are prompt-sensitive.
TYPHOON_PROMPT = """Extract all text from the image.

Instructions:
- Only return the clean Markdown.
- Do not include any explanation or extra text.
- You must include all information on the page.

Formatting Rules:
- Tables: Render tables using <table>...</table> in clean HTML format.
- Equations: Render equations using LaTeX syntax with inline ($...$) and block ($$...$$).
- Images/Charts/Diagrams: Wrap any clearly defined visual areas (e.g. charts, diagrams, pictures) in:

<figure>
Describe the image's main elements (people, objects, text), note any contextual clues (place, event, culture), mention visible text and its meaning, provide deeper analysis when relevant (especially for financial charts, graphs, or documents), comment on style or architecture if relevant, then give a concise overall summary. Describe in Thai.
</figure>

- Page Numbers: Wrap page numbers in <page_number>...</page_number> (e.g., <page_number>14</page_number>).
- Checkboxes: Use ☐ for unchecked and ☑ for checked boxes."""

NANONETS_PROMPT = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""

PADDLE_PROMPT = "OCR:"

_MODELS = {}  # key -> (model on cuda, processor); built at module level below


# ── helpers ────────────────────────────────────────────────────────────────────

def _num(value, default, lo=None, hi=None, integer=False):
    if isinstance(value, bool):
        value = int(value)
    try:
        n = float(str(value).strip())
        if n != n or n in (float("inf"), float("-inf")):
            raise ValueError
    except (TypeError, ValueError, AttributeError):
        n = float(default)
    if lo is not None:
        n = max(lo, n)
    if hi is not None:
        n = min(hi, n)
    return int(round(n)) if integer else float(n)


def _load_image(value):
    """Accept every shape an image port can carry and return a PIL RGB image."""
    if value is None or value == "":
        raise gr.Error("No image reached this node — upload an image in the input node.")
    if isinstance(value, Image.Image):
        return value.convert("RGB")
    src = value
    if isinstance(value, dict):
        src = value.get("path") or value.get("url") or value.get("name") or value.get("data")
    if not isinstance(src, str) or not src:
        raise gr.Error(f"Unrecognised image value: {type(value).__name__}")
    if src.startswith("data:"):
        payload = src.split(",", 1)[1]
        return Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
    if src.startswith("http://") or src.startswith("https://"):
        import urllib.request
        buf = io.BytesIO(urllib.request.urlopen(src, timeout=30).read())
        return Image.open(buf).convert("RGB")
    return Image.open(src).convert("RGB")


def _resize_max(img, max_size):
    """Typhoon-OCR is trained at a fixed 1800 px image dimension (per its card)."""
    w, h = img.size
    if w > max_size or h > max_size:
        if w >= h:
            img = img.resize((max_size, int(h * max_size / float(w))), Image.Resampling.LANCZOS)
        else:
            img = img.resize((int(w * max_size / float(h)), max_size), Image.Resampling.LANCZOS)
    return img


def _build(key):
    model_id = MODEL_IDS[key]
    print(f"[load] {key}: {model_id}", flush=True)
    if key == "lighton":
        model = Qwen3_5ForConditionalGeneration.from_pretrained(model_id, dtype=torch.bfloat16)
    else:
        model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch.bfloat16)
    processor = AutoProcessor.from_pretrained(model_id)
    # ZeroGPU docs: place models on cuda at module level — the CUDA-emulation mode
    # packs them at startup and streams them into the real GPU when a @spaces.GPU
    # call lands. Loading inside the request handler is what expired the ZeroGPU
    # proxy token: first-run downloads made the handler outlive the token TTL.
    _MODELS[key] = (model.eval().to("cuda"), processor)


for _key in MODEL_IDS:
    try:
        _build(_key)
    except Exception as _e:  # one bad repo must not take down the whole Space
        print(f"[load] FAILED {_key}: {_e}", flush=True)


@spaces.GPU(duration=240)
def _run(key, img, max_new_tokens):
    """Move the cached model to CUDA and run the card's inference recipe."""
    if key not in _MODELS:
        raise gr.Error(f"Model {MODEL_IDS.get(key, key)} failed to load at startup — check the Space logs.")
    model, processor = _MODELS[key]

    if key == "lighton":
        # Card: empty prompt = transcription-only; thinking disabled.
        messages = [{"role": "user", "content": [{"type": "image", "image": img}]}]
        inputs = processor.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=True,
            return_dict=True, return_tensors="pt", enable_thinking=False,
        ).to(model.device)
        out = model.generate(**inputs, max_new_tokens=int(max_new_tokens),
                             do_sample=True, temperature=0.2)
    elif key == "typhoon":
        img = _resize_max(img, 1800)
        messages = [{"role": "user", "content": [
            {"type": "image", "image": img},
            {"type": "text", "text": TYPHOON_PROMPT},
        ]}]
        inputs = processor.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt",
        ).to(model.device)
        out = model.generate(**inputs, max_new_tokens=int(max_new_tokens))
    elif key == "nanonets":
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": NANONETS_PROMPT},
            ]},
        ]
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[text], images=[img], padding=True, return_tensors="pt")
        inputs = inputs.to(model.device)
        out = model.generate(**inputs, max_new_tokens=int(max_new_tokens), do_sample=False)
    elif key == "paddle":
        messages = [{"role": "user", "content": [
            {"type": "image", "image": img},
            {"type": "text", "text": PADDLE_PROMPT},
        ]}]
        # Card snippet used images_kwargs={"size": {"shortest_edge":
        # processor.image_processor.min_pixels, ...}}, but PaddleOCRVLImageProcessor
        # has no min_pixels attribute. For the non-spotting "OCR:" task the card
        # specifies default sizing (~1M pixels), so omit the override entirely.
        inputs = processor.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=True,
            return_dict=True, return_tensors="pt",
        ).to(model.device)
        out = model.generate(**inputs, max_new_tokens=int(max_new_tokens))
    else:
        raise gr.Error(f"Unknown model key: {key}")

    gen_ids = out[0][inputs["input_ids"].shape[-1]:]
    text_out = processor.decode(gen_ids, skip_special_tokens=True)
    print(f"[run] {key}: {len(text_out)} chars, {int(max_new_tokens)} max_new_tokens", flush=True)
    return text_out if text_out.strip() else "(empty output)"


# ── the four bound nodes ───────────────────────────────────────────────────────

def lighton_ocr(image, max_new_tokens=4096):
    img = _load_image(image)
    return _run("lighton", img, _num(max_new_tokens, 4096, 256, 16384, integer=True))


def typhoon_ocr(image, max_new_tokens=8192):
    img = _load_image(image)
    return _run("typhoon", img, _num(max_new_tokens, 8192, 256, 16384, integer=True))


def nanonets_ocr(image, max_new_tokens=8192):
    img = _load_image(image)
    return _run("nanonets", img, _num(max_new_tokens, 8192, 256, 16384, integer=True))


def paddle_ocr(image, max_new_tokens=4096):
    img = _load_image(image)
    return _run("paddle", img, _num(max_new_tokens, 4096, 256, 16384, integer=True))


BIND = {
    "lighton_ocr": lighton_ocr,
    "typhoon_ocr": typhoon_ocr,
    "nanonets_ocr": nanonets_ocr,
    "paddle_ocr": paddle_ocr,
}

WORKFLOW = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflow.json")
demo = gr.Workflow(WORKFLOW, bind=BIND)

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