"""
Chandra OCR 2 — Hugging Face Space demo.

Model:  datalab-to/chandra-ocr-2  (~10B, bf16)
Docs:   https://huggingface.co/datalab-to/chandra-ocr-2

Two pipeline stages:
  1. Text & layout extraction — all text in reading order, with tables, math,
     forms and headings preserved as markdown (chandra's parse_markdown).
  2. Visual understanding — locate figures/charts/diagrams/photos/stamps in the
     layout output, crop them at full resolution, and surface the model's
     caption + any structured data it read from them.

Hardware: needs ZeroGPU (H200 slice) or a paid A100/L40S/RTX PRO 6000.
The model is ~20 GB in bf16 and will NOT run on the free CPU tier.

API facts verified against the installed `chandra` package (not guessed):
  - generate_hf(batch, model, max_output_tokens=...) -> list[GenerationResult]
      GenerationResult has .raw, .token_count, .error
  - PROMPT_MAPPING contains exactly: "ocr_layout", "ocr"
  - parse_markdown(raw, include_headers_footers=False, include_images=True)
  - parse_chunks(raw, image, bbox_scale=1000) -> list of
      {"bbox": [x0, y0, x1, y1] (pixels), "label": str, "content": str}
  - settings.BBOX_SCALE == 1000 (bbox is normalised per-axis, so it is
      resolution-independent and can be mapped onto the full-res source image)
"""

import inspect
import json
import os
import time
import zipfile
from pathlib import Path

import gradio as gr
import torch
from PIL import Image

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

#MODEL_ID = "datalab-to/chandra-ocr-2"
MODEL_ID = "datalab-to/surya-ocr-2"

PAGES_PER_GPU_CALL = 3      # keep each ZeroGPU allocation inside its duration budget
GPU_DURATION = 180          # seconds requested per allocation
MAX_PAGES = 20              # guard against someone uploading a 500-page PDF

OUT_DIR = Path(os.environ.get("CHANDRA_OUT_DIR", "/tmp/chandra_out"))
OUT_DIR.mkdir(parents=True, exist_ok=True)

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}
PDF_EXTS = {".pdf"}

# Labels the layout model can attach to blocks that count as "figures" for
# stage 2 (visual understanding). Kept lowercase for matching.
FIGURE_LABELS = {
    "figure", "chart", "diagram", "photo", "image", "picture", "stamp",
    "graph", "plot", "illustration", "logo", "icon", "drawing", "map",
}
CAPTION_LABELS = {"caption", "figure-caption", "figcaption", "figure_caption"}

# ---------------------------------------------------------------------------
# ZeroGPU shim — lets the same file run locally without the `spaces` package
# ---------------------------------------------------------------------------

ON_ZERO = os.environ.get("SPACES_ZERO_GPU") == "true"

try:
    import spaces

    gpu = spaces.GPU
except ImportError:  # local / non-ZeroGPU deploy

    def gpu(*args, **kwargs):
        if args and callable(args[0]):
            return args[0]

        def deco(fn):
            return fn

        return deco


# ---------------------------------------------------------------------------
# Chandra package — preferred path. Falls back to plain transformers.
# ---------------------------------------------------------------------------

CHANDRA_ERR = None
USE_CHANDRA = False
HAS_PARSE_CHUNKS = False
PROMPT_TYPES = ["ocr_layout"]  # safe default; replaced below if package present

try:
    from chandra.model import generate_hf
    from chandra.model.schema import BatchInputItem
    from chandra.model import PROMPT_MAPPING  # enumerate real prompt types

    try:
        from chandra.output import parse_markdown, parse_chunks
    except ImportError:
        from chandra.model.output import parse_markdown, parse_chunks

    USE_CHANDRA = True
    HAS_PARSE_CHUNKS = True
    PROMPT_TYPES = list(PROMPT_MAPPING.keys())
except Exception as e:  # noqa: BLE001
    CHANDRA_ERR = f"{type(e).__name__}: {e}"
    USE_CHANDRA = False
    HAS_PARSE_CHUNKS = False

DEFAULT_PROMPT_TYPE = "ocr_layout" if "ocr_layout" in PROMPT_TYPES else PROMPT_TYPES[0]


# ---------------------------------------------------------------------------
# Model — loaded at module scope, guarded so a failure never kills the Space
# ---------------------------------------------------------------------------

MODEL = None
PROCESSOR = None
MODEL_ERROR = None


def _load():
    from transformers import AutoProcessor

    try:
        from transformers import AutoModelForImageTextToText as VLM
    except ImportError:
        from transformers import AutoModelForVision2Seq as VLM

    # On ZeroGPU there is no GPU visible at import time, so accelerate's
    # device_map="auto" would strand the model on CPU. ZeroGPU instead
    # intercepts .to("cuda") at global scope. Elsewhere, device_map is fine.
    # Blackwell (sm_120) -> SDPA, never flash-attn (no prebuilt wheels).
    kw = dict(low_cpu_mem_usage=True, attn_implementation="sdpa")
    if not ON_ZERO:
        kw["device_map"] = "auto"

    try:
        m = VLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, **kw)
    except TypeError:  # transformers < 4.56 spelled it torch_dtype
        kw.pop("attn_implementation", None)
        m = VLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, **kw)

    if ON_ZERO:
        m = m.to("cuda")

    m.eval()
    p = AutoProcessor.from_pretrained(MODEL_ID)
    p.tokenizer.padding_side = "left"   # required by chandra's batched generate
    m.processor = p
    return m, p


try:
    print(f"Loading {MODEL_ID} ...")
    _t0 = time.time()
    MODEL, PROCESSOR = _load()
    print(f"Loaded in {time.time() - _t0:.0f}s | chandra pkg: {USE_CHANDRA} "
          f"({CHANDRA_ERR or 'ok'})")
except Exception as e:  # noqa: BLE001
    MODEL_ERROR = f"{type(e).__name__}: {e}"
    print(f"MODEL LOAD FAILED: {MODEL_ERROR}")

MODEL_LOADED = MODEL is not None


# ---------------------------------------------------------------------------
# Diagnostics — surfaced in the UI, not just logs
# ---------------------------------------------------------------------------

def _transformers_version():
    try:
        import transformers
        return transformers.__version__
    except Exception:  # noqa: BLE001
        return "n/a"


def diagnostics_md() -> str:
    cuda = torch.cuda.is_available()
    gpu_name = torch.cuda.get_device_name(0) if cuda else "none visible"
    rows = [
        ("torch", torch.__version__),
        ("transformers", _transformers_version()),
        ("gradio", gr.__version__),
        ("CUDA visible", str(cuda)),
        ("GPU", gpu_name),
        ("chandra package", "imported" if USE_CHANDRA else f"FAILED — {CHANDRA_ERR}"),
        ("prompt types", ", ".join(PROMPT_TYPES)),
        ("model loaded", "yes" if MODEL_LOADED else "NO"),
        ("model error", MODEL_ERROR or "—"),
    ]
    return "\n".join(f"- **{k}**: `{v}`" for k, v in rows)


# ---------------------------------------------------------------------------
# Page extraction
# ---------------------------------------------------------------------------

def as_path(f) -> Path:
    """Normalise str / os.PathLike / Gradio file object into a real Path.

    pathlib.Path also has a .name attribute, but there it is the *basename* —
    a hasattr(f, "name") check would silently drop the directory. Handle the
    types distinctly instead.
    """
    if isinstance(f, (str, os.PathLike)):
        return Path(f)
    # Gradio file object (or tempfile.NamedTemporaryFile) exposes .name as a path
    return Path(getattr(f, "name", str(f)))


def pdf_to_images(path: Path, dpi: int):
    try:
        import pymupdf
    except ImportError:
        import fitz as pymupdf

    doc = pymupdf.open(str(path))
    pages = []
    for i, page in enumerate(doc):
        pix = page.get_pixmap(dpi=dpi)
        pages.append((f"{path.stem}_p{i + 1:03d}",
                      Image.frombytes("RGB", (pix.width, pix.height), pix.samples)))
    doc.close()
    return pages


def collect_pages(files, dpi: int):
    """Return (pages, rejected) where pages is [(name, full_res_PIL)] and
    rejected is [(name, reason)]. One bad file never fails the run."""
    pages, rejected = [], []
    for f in files:
        p = as_path(f)
        ext = p.suffix.lower()
        try:
            if ext in PDF_EXTS:
                pages.extend(pdf_to_images(p, dpi))
            elif ext in IMAGE_EXTS:
                pages.append((p.stem, Image.open(p).convert("RGB")))
            else:
                rejected.append((p.name, f"unsupported type '{ext or 'none'}'"))
        except Exception as e:  # noqa: BLE001
            rejected.append((p.name, f"{type(e).__name__}: {e}"))
    return pages, rejected


def fit(img: Image.Image, max_side: int) -> Image.Image:
    """Cap the long edge — visual token count scales with area, so this is the
    single biggest lever on VRAM and latency."""
    img = img.convert("RGB")
    if max(img.size) > max_side:
        s = max_side / max(img.size)
        img = img.resize((max(1, int(img.width * s)), max(1, int(img.height * s))),
                         Image.LANCZOS)
    return img


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

@gpu(duration=GPU_DURATION)
@torch.inference_mode()
def _infer_chunk(images, prompt_type: str, max_new_tokens: int):
    """OCR a small batch of PIL images. Returns a list of raw model strings."""
    if MODEL is None:
        raise RuntimeError("model not loaded")

    if USE_CHANDRA:
        batch = [BatchInputItem(image=im, prompt_type=prompt_type) for im in images]
        results = generate_hf(batch, MODEL, max_output_tokens=max_new_tokens)
        return [getattr(r, "raw", None) or getattr(r, "markdown", "") or str(r)
                for r in results]

    # ---- fallback: drive the chat template directly (degraded) ----
    outs = []
    for im in images:
        msgs = [{"role": "user", "content": [
            {"type": "image", "image": im},
            {"type": "text", "text": prompt_type},
        ]}]
        inputs = PROCESSOR.apply_chat_template(
            msgs, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt",
        ).to(MODEL.device)
        if "pixel_values" in inputs:
            inputs["pixel_values"] = inputs["pixel_values"].to(MODEL.dtype)
        gen = MODEL.generate(**inputs, max_new_tokens=int(max_new_tokens),
                             do_sample=False)
        trimmed = [o[len(i):] for i, o in zip(inputs["input_ids"], gen)]
        outs.append(PROCESSOR.batch_decode(trimmed, skip_special_tokens=True)[0].strip())
    return outs


def to_markdown(raw: str) -> str:
    if USE_CHANDRA:
        try:
            return parse_markdown(raw)
        except Exception:  # noqa: BLE001
            pass
    return raw


# ---------------------------------------------------------------------------
# Stage 2 — figure extraction (visual understanding)
# ---------------------------------------------------------------------------

def _find_caption(chunks, fig_idx: int):
    """Look for a caption chunk immediately after (or before) the figure."""
    for j in (fig_idx + 1, fig_idx - 1):
        if 0 <= j < len(chunks):
            lbl = (chunks[j].get("label") or "").lower()
            if lbl in CAPTION_LABELS:
                return chunks[j].get("content") or ""
    return ""


def extract_figures(raw: str, full_res: Image.Image):
    """Locate figure-like blocks in the layout output and crop them from the
    FULL-RESOLUTION source image (bbox is per-axis normalised, so it maps
    cleanly onto any resolution with the same aspect ratio).

    Returns (figures, degraded) where figures is a list of dicts and degraded
    is True when we could not use the real chandra path (so captions/data are
    NOT from the model's figure understanding).
    """
    if not USE_CHANDRA or not HAS_PARSE_CHUNKS:
        return [], True

    try:
        chunks = parse_chunks(raw, full_res, bbox_scale=1000)
    except Exception:  # noqa: BLE001
        return [], True

    figures = []
    for i, ch in enumerate(chunks):
        label = (ch.get("label") or "").lower()
        if label not in FIGURE_LABELS:
            continue
        bbox = ch.get("bbox")
        if not bbox or len(bbox) != 4:
            continue
        x0, y0, x1, y1 = (int(v) for v in bbox)
        w, h = full_res.size
        x0, y0 = max(0, x0), max(0, y0)
        x1, y1 = min(w, x1), min(h, y1)
        if x1 <= x0 or y1 <= y0:
            continue
        crop = full_res.crop((x0, y0, x1, y1))
        figures.append({
            "page": None,          # filled by caller
            "label": label,
            "bbox": [x0, y0, x1, y1],
            "crop": crop,
            "caption": _find_caption(chunks, i),
            "structured": (ch.get("content") or "").strip(),
        })
    return figures, False


# ---------------------------------------------------------------------------
# Orchestration — streaming generator
# ---------------------------------------------------------------------------

FIG_HEADERS = ["Page", "Label", "BBox (x0,y0,x1,y1)", "Caption", "Structured data"]


def run(files, prompt_type, dpi, max_side, max_new_tokens, progress=gr.Progress()):

    def empty(status):
        return "", "", "", status, None, [], None, None

    if not MODEL_LOADED:
        yield empty(f"Model failed to load — see diagnostics.\n{MODEL_ERROR}")
        return

    if not files:
        yield empty("Upload a PDF or some images first.")
        return

    try:
        pages, rejected = collect_pages(files, int(dpi))
    except Exception as e:  # noqa: BLE001
        yield empty(f"Could not read those files:\n{type(e).__name__}: {e}")
        return

    if not pages:
        yield empty("No readable PDF or image files found in that upload.")
        return

    skip_note = ""
    if rejected:
        skip_note = "\nSkipped:\n" + "\n".join(f"  - {n}: {r}" for n, r in rejected)

    truncated = ""
    if len(pages) > MAX_PAGES:
        truncated = f"  (truncated from {len(pages)})"
        pages = pages[:MAX_PAGES]

    md_parts, raw_parts = [], []
    all_figs = []          # flattened figure records for the gallery/table
    per_page = []          # (name, md, raw, figs)
    t_start = time.time()

    for start in range(0, len(pages), PAGES_PER_GPU_CALL):
        chunk = pages[start:start + PAGES_PER_GPU_CALL]
        names = [n for n, _ in chunk]
        full_res = [im for _, im in chunk]
        imgs = [fit(im, int(max_side)) for im in full_res]

        progress(start / len(pages),
                 desc=f"{names[0]} … ({start + 1}-{start + len(chunk)}/{len(pages)})")

        try:
            raws = _infer_chunk(imgs, prompt_type, int(max_new_tokens))
        except torch.cuda.OutOfMemoryError:
            torch.cuda.empty_cache()
            raws = ["[OUT OF MEMORY — lower 'Max image side']"] * len(chunk)
        except Exception as e:  # noqa: BLE001
            raws = [f"[FAILED: {type(e).__name__}: {e}]"] * len(chunk)

        for name, fimg, raw in zip(names, full_res, raws):
            md = to_markdown(raw)
            figs, degraded = extract_figures(raw, fimg)
            for fg in figs:
                fg["page"] = name
                fg["degraded"] = degraded
                all_figs.append(fg)

            md_parts.append(f"\n\n---\n\n## {name}\n\n{md}")
            raw_parts.append(f"===== {name} =====\n{raw}")
            per_page.append((name, md, raw, figs))

        elapsed = time.time() - t_start
        joined = "\n".join(md_parts)
        yield (joined, joined,
               "\n\n".join(raw_parts),
               f"{len(per_page)}/{len(pages)} pages{truncated} · {elapsed:.0f}s "
               f"({elapsed / max(1, len(per_page)):.1f}s/page){skip_note}",
               _gallery(all_figs), _fig_rows(all_figs),
               None, None)

    # ---- artefacts ----
    stamp = time.strftime("%Y%m%d_%H%M%S")
    md_path = OUT_DIR / f"chandra_{stamp}.md"
    md_path.write_text("\n".join(md_parts), encoding="utf-8")

    zip_path = OUT_DIR / f"chandra_{stamp}.zip"
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
        for name, md, raw, figs in per_page:
            safe = "".join(c if c.isalnum() or c in "-_." else "_" for c in name)
            z.writestr(f"markdown/{safe}.md", md)
            z.writestr(f"raw/{safe}.txt", raw)
            for k, fg in enumerate(figs):
                buf = _to_png_bytes(fg["crop"])
                if buf:
                    z.writestr(f"figures/{safe}_{k:02d}_{fg['label']}.png", buf)
        z.writestr("manifest.json", json.dumps({
            "model": MODEL_ID,
            "prompt_type": prompt_type,
            "dpi": dpi,
            "max_side": max_side,
            "pages": [n for n, _, _, _ in per_page],
            "figures": [
                {"page": f["page"], "label": f["label"], "bbox": f["bbox"],
                 "caption": f["caption"], "structured": f["structured"]}
                for f in all_figs
            ],
        }, indent=2))

    total = time.time() - t_start
    joined = "\n".join(md_parts)
    yield (joined, joined,
           "\n\n".join(raw_parts),
           f"Done — {len(per_page)} page(s){truncated} in {total:.0f}s "
           f"({total / max(1, len(per_page)):.1f}s/page){skip_note}",
           _gallery(all_figs), _fig_rows(all_figs),
           str(md_path), str(zip_path))


def _to_png_bytes(img: Image.Image):
    import io
    buf = io.BytesIO()
    try:
        img.save(buf, format="PNG")
        return buf.getvalue()
    except Exception:  # noqa: BLE001
        return None


def _gallery(figs):
    """Gallery entries: (crop, caption). Caption is clearly marked degraded when
    it did not come from the model's figure-understanding path."""
    out = []
    for f in figs:
        cap = f["caption"] or "(no caption from model)"
        if f.get("degraded"):
            cap = f"[degraded — not from model] {cap}"
        out.append((f["crop"], f"{f['page']} · {f['label']} · {cap}"))
    return out or None


def _fig_rows(figs):
    rows = []
    for f in figs:
        cap = f["caption"] or ""
        if f.get("degraded"):
            cap = f"[degraded] {cap}".strip()
        rows.append([f["page"], f["label"],
                     f"({f['bbox'][0]},{f['bbox'][1]},{f['bbox'][2]},{f['bbox'][3]})",
                     cap, f["structured"]])
    return rows


# ---------------------------------------------------------------------------
# UI — Gradio 4/5/6 tolerant
# ---------------------------------------------------------------------------

GR_MAJOR = int(gr.__version__.split(".")[0])


def C(cls, **kw):
    """Build a component, dropping kwargs this Gradio version rejects.

    Gradio 6 removed Textbox.show_copy_button and moved theme/css from
    Blocks() to launch(). This keeps one file working across 4/5/6.
    """
    try:
        allowed = set(inspect.signature(cls.__init__).parameters)
        if "kwargs" not in allowed:
            kw = {k: v for k, v in kw.items() if k in allowed}
    except (TypeError, ValueError):
        pass
    return cls(**kw)


CSS = """
#raw_out textarea { font-family: ui-monospace, monospace; font-size: 12px; }
#status textarea  { font-family: ui-monospace, monospace; font-size: 12px; }
.md_pane { max-height: 640px; overflow-y: auto; }
"""

_STYLE = dict(theme=gr.themes.Soft(), css=CSS)
_BLOCKS_KW = {} if GR_MAJOR >= 6 else _STYLE
_LAUNCH_KW = _STYLE if GR_MAJOR >= 6 else {}

with gr.Blocks(title="Chandra OCR 2", **_BLOCKS_KW) as demo:
    gr.Markdown(
        f"""
# Chandra OCR 2 — document → markdown / HTML / JSON + figures

Layout-aware OCR from [Datalab](https://datalab.to). Handles tables, math,
forms, handwriting and 90+ languages, preserving reading order and structure.

Two outputs per page:
1. **Text & layout** — markdown with tables/math/forms/headings preserved.
2. **Figures** — charts, diagrams, photos and stamps cropped at full resolution,
   with the model's caption and any structured data it read from them.

Upload **PDFs and/or images**; each page is processed separately and results
stream in below. Capped at **{MAX_PAGES} pages** per run in this demo.

Model: [`{MODEL_ID}`](https://huggingface.co/{MODEL_ID}) · weights are under a
modified OpenRAIL-M licence (free for research, personal use, and companies
under $2M funding/revenue — **not** for building a competitor to Datalab's API).
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            # NOTE: no file_types filter — it silently rejects valid files on
            # some browser/version combinations. Validation happens in Python.
            files = C(gr.Files, label="PDFs / images", file_count="multiple")
            go = C(gr.Button, value="Run OCR", variant="primary")

            prompt_type = C(gr.Dropdown, label="Prompt type", choices=PROMPT_TYPES,
                            value=DEFAULT_PROMPT_TYPE, allow_custom_value=True,
                            info="Enumerated from the chandra package; editable "
                                 "so you can try undocumented values.")

            with gr.Accordion("Advanced", open=False):
                dpi = C(gr.Slider, minimum=100, maximum=400, value=200, step=25,
                        label="PDF render DPI", info="200-300 suits most scans.")
                max_side = C(gr.Slider, minimum=768, maximum=2560, value=1540, step=64,
                             label="Max image side (px)",
                             info="Biggest lever on speed and VRAM.")
                max_new = C(gr.Slider, minimum=512, maximum=8192, value=4096, step=256,
                            label="Max new tokens",
                            info="Layout output is verbose — keep this high.")

            status = C(gr.Textbox, label="Status", lines=3, elem_id="status")

            # Diagnostics — auto-open when the model failed to load.
            diag = gr.Accordion("Diagnostics", open=not MODEL_LOADED)
            with diag:
                gr.Markdown(diagnostics_md())

        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("Rendered"):
                    md_view = C(gr.Markdown, value="", elem_classes=["md_pane"])
                with gr.Tab("Markdown source"):
                    md_src = C(gr.Textbox, label=None, lines=24, show_copy_button=True)
                with gr.Tab("Raw model output"):
                    raw_view = C(gr.Textbox, label=None, lines=24,
                                 elem_id="raw_out", show_copy_button=True)
            with gr.Tab("Figures"):
                fig_gallery = C(gr.Gallery, label="Extracted figures",
                                columns=3, height="auto",
                                object_fit="contain")
                fig_table = C(gr.Dataframe, headers=FIG_HEADERS,
                              label="Figure metadata", interactive=False)
            with gr.Row():
                md_file = C(gr.File, label="Combined .md")
                zip_file = C(gr.File, label="All pages .zip")

    go.click(
        run,
        inputs=[files, prompt_type, dpi, max_side, max_new],
        outputs=[md_view, md_src, raw_view, status, fig_gallery, fig_table,
                 md_file, zip_file],
    )


if __name__ == "__main__":
    demo.queue(max_size=12).launch(show_error=True, **_LAUNCH_KW)