"""Surya2 streamlit app — exercise layout, recognition, table_rec via the
inference manager. Detection + OCR-error stay in their own torch paths."""

from __future__ import annotations

import io
import re
import tempfile
import time
from typing import List

import pypdfium2
import streamlit as st
import streamlit.components.v1 as components
from PIL import Image, ImageDraw

from surya.debug.draw import draw_polys_on_image, draw_bboxes_on_image
from surya.detection import TextDetectionResult
from surya.inference import SuryaInferenceManager
from surya.layout import LayoutPredictor
from surya.layout.schema import LayoutResult
from surya.recognition import RecognitionPredictor
from surya.recognition.schema import PageOCRResult
from surya.settings import settings
from surya.table_rec import TableRecPredictor
from surya.table_rec.schema import TableResult


# KaTeX-enabled HTML wrapper. The OCR HTML wraps math in <math>...</math>
# (KaTeX-compatible LaTeX inside), which a browser would otherwise show as
# raw text. We convert those tags to \( \) / \[ \] delimiters and let KaTeX
# auto-render typeset them inside an iframe component.
_KATEX_HEAD = r"""<!doctype html><html><head>
<meta charset="utf-8">
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.css">
<script src="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.js"></script>
<script src="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/contrib/auto-render.min.js"></script>
<style>
/* White "paper" card so the text stays readable in both light and dark
   Streamlit themes (the iframe is otherwise transparent and our text is dark). */
html,body{background:#ffffff;}
body{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",sans-serif;font-size:15px;line-height:1.55;color:#111111;margin:0;padding:14px;}
table{border-collapse:collapse;margin:6px 0;} td,th{border:1px solid #bbb;padding:3px 6px;color:#111111;}
[data-label="SectionHeader"],[data-label="PageHeader"]{font-weight:600;}
</style></head><body>
"""

_KATEX_TAIL = r"""
<script>
renderMathInElement(document.body, {
  delimiters: [
    {left: "\\[", right: "\\]", display: true},
    {left: "\\(", right: "\\)", display: false}
  ],
  throwOnError: false
});
</script></body></html>
"""

_MATH_RE = re.compile(r"<math\b([^>]*)>(.*?)</math>", re.DOTALL | re.IGNORECASE)


def _math_to_katex(html_str: str) -> str:
    """Rewrite <math>...</math> tags into KaTeX \\( \\) / \\[ \\] delimiters."""

    def repl(m: "re.Match") -> str:
        attrs, inner = m.group(1), m.group(2)
        if re.search(r"""display\s*=\s*["']block["']""", attrs):
            return "\\[" + inner + "\\]"
        return "\\(" + inner + "\\)"

    return _MATH_RE.sub(repl, html_str or "")


def render_ocr_html(html_str: str, height: int = 400) -> None:
    """Render OCR HTML with math typeset by KaTeX (iframe component)."""
    components.html(
        _KATEX_HEAD + _math_to_katex(html_str) + _KATEX_TAIL,
        height=height,
        scrolling=True,
    )


def _assemble_page_html(page: PageOCRResult) -> str:
    """Reconstruct a div-block whole-page HTML from a PageOCRResult."""
    parts: List[str] = []
    for blk in page.blocks:
        if blk.skipped:
            continue
        x0, y0, x1, y1 = (int(c) for c in blk.bbox)
        body = blk.html or ""
        parts.append(
            f'<div data-bbox="{x0} {y0} {x1} {y1}" data-label="{blk.label}">{body}</div>'
        )
    return "\n".join(parts)


def _show_timing(label: str, elapsed_s: float, extra: str = "") -> None:
    """Render a small caption with wall-clock + optional extra detail."""
    detail = f" — {extra}" if extra else ""
    st.caption(f"⏱ {label}: {elapsed_s * 1000:.0f} ms ({elapsed_s:.2f}s){detail}")


@st.cache_resource()
def load_predictors_cached():
    manager = SuryaInferenceManager()
    layout_predictor = LayoutPredictor(manager)
    rec_predictor = RecognitionPredictor(manager)
    table_rec_predictor = TableRecPredictor(manager)

    # Lazy-import detection / ocr_error to keep startup snappy when the user
    # only wants VLM modes
    from surya.detection import DetectionPredictor
    from surya.ocr_error import OCRErrorPredictor

    return {
        "manager": manager,
        "layout": layout_predictor,
        "recognition": rec_predictor,
        "table_rec": table_rec_predictor,
        "detection": DetectionPredictor(),
        "ocr_error": OCRErrorPredictor(),
    }


def text_detection(img) -> tuple[Image.Image, TextDetectionResult, float]:
    t = time.perf_counter()
    text_pred = predictors["detection"]([img])[0]
    elapsed = time.perf_counter() - t
    text_polygons = [p.polygon for p in text_pred.bboxes]
    det_img = draw_polys_on_image(text_polygons, img.copy())
    return det_img, text_pred, elapsed


def layout_detection(img) -> tuple[Image.Image, LayoutResult, float]:
    t = time.perf_counter()
    pred = predictors["layout"]([img])[0]
    elapsed = time.perf_counter() - t
    polygons = [p.polygon for p in pred.bboxes]
    labels = [
        f"{p.label}-{p.position}-c{p.count}-{round(p.confidence or 0, 2)}"
        for p in pred.bboxes
    ]
    annotated = draw_polys_on_image(
        polygons, img.copy(), labels=labels, label_font_size=14
    )
    return annotated, pred, elapsed


def block_ocr(img) -> tuple[Image.Image, PageOCRResult, LayoutResult, float, float]:
    """Layout → block crops → BLOCK_PROMPT. Returns layout + block-OCR timings."""
    t_layout = time.perf_counter()
    layout = predictors["layout"]([img])[0]
    layout_elapsed = time.perf_counter() - t_layout

    t_blocks = time.perf_counter()
    page_results = predictors["recognition"]([img], [layout])
    blocks_elapsed = time.perf_counter() - t_blocks
    page = page_results[0]

    annotated = img.copy()
    draw = ImageDraw.Draw(annotated)
    for blk in page.blocks:
        x0, y0, x1, y1 = blk.bbox
        color = "red" if blk.error else ("orange" if blk.skipped else "green")
        draw.rectangle((x0, y0, x1, y1), outline=color, width=3)
        draw.text((x0 + 4, y0 + 4), f"{blk.reading_order} {blk.label}", fill=color)
    return annotated, page, layout, layout_elapsed, blocks_elapsed


def full_page_ocr(img) -> tuple[Image.Image, PageOCRResult, float]:
    """Single HIGH_ACCURACY_BBOX_PROMPT call on the whole page."""
    t = time.perf_counter()
    page_results = predictors["recognition"]([img], full_page=True)
    elapsed = time.perf_counter() - t
    page = page_results[0]
    annotated = img.copy()
    draw = ImageDraw.Draw(annotated)
    for blk in page.blocks:
        x0, y0, x1, y1 = blk.bbox
        color = "red" if blk.error else ("orange" if blk.skipped else "green")
        draw.rectangle((x0, y0, x1, y1), outline=color, width=3)
        draw.text((x0 + 4, y0 + 4), f"{blk.reading_order} {blk.label}", fill=color)
    return annotated, page, elapsed


def table_recognition(
    img: Image.Image,
    mode: str,
    skip_table_detection: bool,
) -> tuple[Image.Image, List[TableResult], float, float]:
    """Returns (annotated_img, table_preds, layout_elapsed, table_rec_elapsed)."""
    layout_elapsed = 0.0
    if skip_table_detection:
        table_imgs = [img]
        table_counts = [0]
        table_bboxes = [(0, 0, img.size[0], img.size[1])]
    else:
        t = time.perf_counter()
        layout = predictors["layout"]([img])[0]
        layout_elapsed = time.perf_counter() - t
        tables = [b for b in layout.bboxes if b.label in ("Table", "TableOfContents")]
        if not tables:
            return img.copy(), [], layout_elapsed, 0.0
        table_bboxes = [tuple(int(c) for c in b.bbox) for b in tables]
        table_imgs = [img.crop(b) for b in table_bboxes]
        table_counts = [b.count for b in tables]

    t = time.perf_counter()
    if mode == "full":
        table_preds = predictors["table_rec"].predict_full(
            table_imgs, counts=table_counts
        )
    else:
        table_preds = predictors["table_rec"].predict_simple(table_imgs)
    table_rec_elapsed = time.perf_counter() - t

    out_img = img.copy()
    for pred, table_img, tbbox in zip(table_preds, table_imgs, table_bboxes):
        if pred.error or pred.mode != "simple" or not pred.rows:
            continue
        row_bboxes = [r.bbox for r in pred.rows]
        col_bboxes = [c.bbox for c in pred.cols]
        row_labels = [r.label for r in pred.rows]
        col_labels = [c.label for c in pred.cols]
        annot = table_img.copy()
        annot = draw_bboxes_on_image(
            row_bboxes, annot, labels=row_labels, label_font_size=14, color="blue"
        )
        annot = draw_bboxes_on_image(
            col_bboxes, annot, labels=col_labels, label_font_size=14, color="red"
        )
        # Paste annotated crop back at the table's position in the page.
        out_img.paste(annot, (tbbox[0], tbbox[1]))
    return out_img, table_preds, layout_elapsed, table_rec_elapsed


def ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15):
    from pdftext.extraction import plain_text_output

    with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
        f.write(pdf_file.getvalue())
        f.seek(0)

        page_middle = page_count // 2
        page_range = range(
            max(page_middle - max_pages, 0), min(page_middle + max_pages, page_count)
        )
        text = plain_text_output(f.name, page_range=page_range)

    sample_gap = len(text) // max_samples
    if len(text) == 0 or sample_gap == 0:
        return "This PDF has no text or very little text", ["no text"]

    if sample_gap < sample_len:
        sample_gap = sample_len

    samples = []
    for i in range(0, len(text), sample_gap):
        samples.append(text[i : i + sample_len])

    results = predictors["ocr_error"](samples)
    label = "This PDF has good text."
    if results.labels.count("bad") / len(results.labels) > 0.2:
        label = "This PDF may have garbled or bad OCR text."
    return label, results.labels


def open_pdf(pdf_file):
    stream = io.BytesIO(pdf_file.getvalue())
    return pypdfium2.PdfDocument(stream)


@st.cache_data()
def get_page_image(pdf_file, page_num, dpi=settings.IMAGE_DPI):
    doc = open_pdf(pdf_file)
    renderer = doc.render(
        pypdfium2.PdfBitmap.to_pil,
        page_indices=[page_num - 1],
        scale=dpi / 72,
    )
    png = list(renderer)[0]
    png_image = png.convert("RGB")
    doc.close()
    return png_image


@st.cache_data()
def page_counter(pdf_file):
    doc = open_pdf(pdf_file)
    doc_len = len(doc)
    doc.close()
    return doc_len


st.set_page_config(layout="wide")
col1, col2 = st.columns([0.55, 0.45])

predictors = load_predictors_cached()

st.markdown(
    """
# Surya 2 Demo

VLM-backed layout, OCR, and table recognition. The model runs in a local
`llama-server` (or vllm) process, started on first use.

Modes:
- **Layout**: page → list of blocks with label + bbox + token count
- **Block OCR**: layout + per-block HTML
- **Table Rec (simple)**: row + column bboxes only
- **Table Rec (full)**: full HTML for each detected table
"""
)

in_file = st.sidebar.file_uploader(
    "PDF file or image:", type=["pdf", "png", "jpg", "jpeg", "gif", "webp"]
)

if in_file is None:
    st.stop()

filetype = in_file.type
page_count = None
if "pdf" in filetype:
    page_count = page_counter(in_file)
    page_number = st.sidebar.number_input(
        f"Page number out of {page_count}:", min_value=1, value=1, max_value=page_count
    )
    # Render at high DPI so the OCR / table-rec demos see fine glyphs.
    # Layout + detection internally downsample (or accept the small perf hit
    # at demo scale); we always render and display the high-DPI page here.
    pil_image = get_page_image(in_file, page_number, settings.IMAGE_DPI_HIGHRES)
else:
    pil_image = Image.open(in_file).convert("RGB")
    page_number = None

run_full_page_ocr = st.sidebar.button("Run Full-Page OCR")
run_text_det = st.sidebar.button("Run Text Detection")
run_layout = st.sidebar.button("Run Layout Analysis")
run_table_rec = st.sidebar.button("Run Table Rec")
run_block_ocr = st.sidebar.button("Run Block OCR")
run_ocr_errors = st.sidebar.button("Run bad-PDF-text detection")

table_mode = st.sidebar.radio(
    "Table mode",
    options=["simple", "full"],
    index=0,
    help="simple: rows+cols only. full: full HTML.",
)
skip_table_detection = st.sidebar.checkbox(
    "Skip table detection",
    value=False,
    help="Treat the entire page/image as a single table.",
)

if pil_image is None:
    st.stop()


if run_text_det:
    det_img, text_pred, elapsed = text_detection(pil_image)
    with col1:
        _show_timing("Text detection", elapsed, f"{len(text_pred.bboxes)} polys")
        st.image(det_img, caption="Detected Text", use_container_width=True)
        st.json(
            text_pred.model_dump(exclude=["heatmap", "affinity_map"]), expanded=False
        )


if run_layout:
    annotated, pred, elapsed = layout_detection(pil_image)
    with col1:
        _show_timing("Layout", elapsed, f"{len(pred.bboxes)} blocks")
        st.image(annotated, caption="Detected Layout", use_container_width=True)
        st.json(pred.model_dump(), expanded=False)


if run_block_ocr:
    annotated, page, layout, t_layout, t_blocks = block_ocr(pil_image)
    with col1:
        n_blocks = len(page.blocks)
        n_ok = sum(1 for b in page.blocks if not b.skipped and not b.error)
        _show_timing("Block OCR — layout", t_layout, f"{n_blocks} blocks")
        _show_timing("Block OCR — per-block OCR", t_blocks, f"{n_ok} OCR'd")
        _show_timing("Block OCR — total", t_layout + t_blocks)
        st.image(
            annotated,
            caption="Block OCR (green=ok, orange=skipped, red=error)",
            use_container_width=True,
        )
        full_html = _assemble_page_html(page)
        with st.expander("Full page HTML (rendered)", expanded=False):
            render_ocr_html(full_html, height=600)
        with st.expander("Full page HTML (source)", expanded=False):
            st.code(full_html, language="html")
        for blk in page.blocks:
            with st.expander(
                f"#{blk.reading_order} {blk.label} (conf {blk.confidence:.2f})"
            ):
                # Diagnostics: show numeric bbox + polygon + a thumbnail with the
                # drawn rectangle highlighted, then the actual crop fed to OCR.
                xs = [p[0] for p in blk.polygon]
                ys = [p[1] for p in blk.polygon]
                bbox_drawn = [int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))]
                cx0 = max(0, int(min(xs)) - 4)
                cy0 = max(0, int(min(ys)) - 4)
                cx1 = min(pil_image.size[0], int(max(xs)) + 4)
                cy1 = min(pil_image.size[1], int(max(ys)) + 4)
                st.text(
                    f"bbox(drawn) = {bbox_drawn}\n"
                    f"crop(ocr)  = {(cx0, cy0, cx1, cy1)}  (= bbox ± 4px pad)"
                )
                # Thumbnail with this block's rectangle highlighted in red.
                thumb = pil_image.copy()
                ImageDraw.Draw(thumb).rectangle(bbox_drawn, outline="red", width=4)
                st.image(thumb, caption="this block's drawn rect (red)", width=300)
                # The actual crop fed to OCR
                if cx1 > cx0 and cy1 > cy0:
                    st.image(pil_image.crop((cx0, cy0, cx1, cy1)), caption="OCR crop")
                if blk.skipped:
                    st.info("Block skipped (visual label)")
                elif blk.error:
                    st.error("Block OCR errored")
                else:
                    render_ocr_html(blk.html, height=160)
                    st.code(blk.html, language="html")


if run_full_page_ocr:
    annotated, page, elapsed = full_page_ocr(pil_image)
    with col1:
        n_blocks = len(page.blocks)
        n_ok = sum(1 for b in page.blocks if not b.skipped and not b.error)
        _show_timing("Full-Page OCR", elapsed, f"{n_blocks} blocks parsed, {n_ok} OK")
        st.image(
            annotated,
            caption="Full-Page OCR (green=ok, orange=skipped, red=error)",
            use_container_width=True,
        )
        full_html = _assemble_page_html(page)
        with st.expander("Full page HTML (rendered)", expanded=False):
            render_ocr_html(full_html, height=600)
        with st.expander("Full page HTML (source)", expanded=False):
            st.code(full_html, language="html")
        for blk in page.blocks:
            with st.expander(
                f"#{blk.reading_order} {blk.label} (conf {blk.confidence:.2f})"
            ):
                if blk.skipped:
                    st.info("Block skipped (visual label)")
                elif blk.error:
                    st.error("Block OCR errored")
                else:
                    render_ocr_html(blk.html, height=160)
                    st.code(blk.html, language="html")


if run_table_rec:
    table_img, preds, t_layout, t_table = table_recognition(
        pil_image, table_mode, skip_table_detection
    )
    with col1:
        if not skip_table_detection:
            _show_timing("Table Rec — layout", t_layout, f"{len(preds)} tables found")
        _show_timing(f"Table Rec — {table_mode}", t_table)
        if not skip_table_detection:
            _show_timing("Table Rec — total", t_layout + t_table)
        st.image(table_img, caption="Table Recognition", use_container_width=True)
        for pred in preds:
            if pred.mode == "full" and pred.html:
                with st.expander("Table HTML"):
                    render_ocr_html(pred.html, height=400)
                    st.code(pred.html, language="html")
            else:
                st.json(pred.model_dump(), expanded=False)


if run_ocr_errors:
    if "pdf" not in filetype:
        st.error("This feature only works with PDFs.")
    else:
        label, results = ocr_errors(in_file, page_count)
        with col1:
            st.write(label)
            st.json(results)


with col2:
    st.image(pil_image, caption="Uploaded Image", use_container_width=True)
