from __future__ import annotations

import base64
import html
import io
import json
import os
import re
import subprocess
import sys
import time
from pathlib import Path
from threading import Thread

# Pre-built causal_conv1d wheel (linear-attention dependency, can't compile in Space build).
# Pulled from Dao-AILab official GitHub release. Auto-selects the cpXY wheel matching the
# running Python version (HF Space = cp310, local dev py3.12 = cp312, etc.).
# Soft-fails on unsupported combos — flash-linear-attention falls back to its torch impl.
_CAUSAL_CONV1D_PY = f"cp{sys.version_info.major}{sys.version_info.minor}"
_CAUSAL_CONV1D_WHEEL = (
    "https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.6.2.post1/"
    f"causal_conv1d-1.6.2.post1+cu12torch2.8cxx11abiTRUE-"
    f"{_CAUSAL_CONV1D_PY}-{_CAUSAL_CONV1D_PY}-linux_x86_64.whl"
)
try:
    import causal_conv1d  # noqa: F401
except ImportError:
    try:
        subprocess.run(
            [sys.executable, "-m", "pip", "install", _CAUSAL_CONV1D_WHEEL, "--no-deps", "-q"],
            check=True,
        )
        import causal_conv1d  # noqa: F401
    except Exception as e:
        print(f"[warn] causal_conv1d unavailable ({e}); slower linear-attention conv1d path will be used")

import gradio as gr
import spaces
from chandra.output import parse_markdown
from PIL import Image


MODEL_ID = os.getenv("MODEL_ID", "ebinan92/Qwen3.5-ocr-jp-2b")
PROMPT = os.getenv(
    "OCR_PROMPT", "OCR this image as HTML layout blocks with bbox and label."
)
MAX_NEW_TOKENS = 12288
# Area-based image cap (matches chandra-ocr-2 reference Space). 1536² px² ≈ 2300 visual tokens.
MAX_IMAGE_PIXELS = int(os.getenv("MAX_IMAGE_PIXELS", str(1536 * 1536)))
PDF_DPI = int(os.getenv("PDF_DPI", "180"))
HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")

LABEL_COLORS = {
    "Text": "#2196F3",
    "Section-Header": "#F44336",
    "Equation-Block": "#9C27B0",
    "Table": "#FF9800",
    "Image": "#4CAF50",
    "Figure": "#009688",
    "Caption": "#795548",
    "Footnote": "#607D8B",
    "List-Group": "#3F51B5",
    "Page-Header": "#CDDC39",
    "Page-Footer": "#FFC107",
    "Code-Block": "#00BCD4",
    "Bibliography": "#E91E63",
    "Complex-Block": "#8BC34A",
    "Form": "#FF5722",
    "Table-Of-Contents": "#673AB7",
    "Diagram": "#FFEB3B",
}

_BLOCK_RE = re.compile(
    r'<div\s+data-bbox="([^"]+)"\s+data-label="([^"]+)">(.*?)</div>',
    re.DOTALL,
)


def _load_model_eager():
    """Module-level eager load for ZeroGPU.

    Uses device_map="auto" + flash-linear-attention CUDA kernels (via the
    `fla` package pulled in by requirements). Without flash-linear-attention,
    the 18 linear_attention layers in Qwen3.5 fall back to a slow Python
    implementation — that's the 5–10x slowdown we hit before.
    """
    import torch
    from transformers import AutoModelForImageTextToText, AutoProcessor

    print(f"[startup] Loading model: {MODEL_ID}")
    kwargs = {"trust_remote_code": True, "device_map": "auto"}
    if HF_TOKEN:
        kwargs["token"] = HF_TOKEN
    try:
        model = AutoModelForImageTextToText.from_pretrained(
            MODEL_ID, dtype=torch.bfloat16, **kwargs,
        )
    except TypeError:
        model = AutoModelForImageTextToText.from_pretrained(
            MODEL_ID, torch_dtype=torch.bfloat16, **kwargs,
        )
    model.eval()

    processor = AutoProcessor.from_pretrained(
        MODEL_ID, trust_remote_code=True, token=HF_TOKEN,
    )
    if hasattr(processor, "tokenizer"):
        processor.tokenizer.padding_side = "left"
    print(f"[startup] Model ready")
    return model, processor


MODEL, PROCESSOR = _load_model_eager()


def _resize_if_needed(image: Image.Image) -> Image.Image:
    """Area-based scale-to-fit (matches chandra-ocr-2 Space)."""
    if image.mode != "RGB":
        image = image.convert("RGB")
    if MAX_IMAGE_PIXELS <= 0:
        return image
    w, h = image.size
    total = w * h
    if total <= MAX_IMAGE_PIXELS:
        return image
    scale = (MAX_IMAGE_PIXELS / total) ** 0.5
    return image.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)


def _render_pdf_page(path: Path, page_number: int) -> tuple[int, Image.Image]:
    import pypdfium2 as pdfium

    doc = pdfium.PdfDocument(str(path))
    try:
        if page_number < 1 or page_number > len(doc):
            raise gr.Error(
                f"Invalid page number ({page_number}). Must be between 1 and {len(doc)}."
            )
        page = doc[page_number - 1]
        image = page.render(scale=PDF_DPI / 72).to_pil()
        page.close()
        return page_number, _resize_if_needed(image)
    finally:
        doc.close()


def _load_image(file_path: str, page_number: int) -> tuple[int, Image.Image]:
    path = Path(file_path)
    if path.suffix.lower() == ".pdf":
        return _render_pdf_page(path, page_number)

    img = Image.open(path)
    img.load()  # eagerly read pixel data so the fp can be released
    return 1, _resize_if_needed(img)


@spaces.GPU(duration=180)
def _generate_page_stream(
    image: Image.Image,
    temperature: float,
):
    """Streaming inference. Yields (partial_text, token_count, elapsed, done)."""
    from transformers import TextIteratorStreamer

    conversation = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": PROMPT},
            ],
        }
    ]

    inputs = PROCESSOR.apply_chat_template(
        [conversation],
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
        padding=True,
    )
    inputs = inputs.to(MODEL.device)

    eos_token_id = MODEL.generation_config.eos_token_id
    im_end_id = PROCESSOR.tokenizer.convert_tokens_to_ids("<|im_end|>")
    if isinstance(eos_token_id, int):
        eos_token_id = [eos_token_id]
    elif eos_token_id is None:
        eos_token_id = []
    if im_end_id is not None and im_end_id not in eos_token_id:
        eos_token_id.append(im_end_id)

    streamer = TextIteratorStreamer(
        PROCESSOR.tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    gen_kwargs = dict(
        **inputs,
        max_new_tokens=MAX_NEW_TOKENS,
        eos_token_id=eos_token_id,
        streamer=streamer,
    )

    result_holder: dict = {}

    def _run():
        try:
            result_holder["ids"] = MODEL.generate(**gen_kwargs)
        except Exception as e:
            result_holder["error"] = e

    start = time.perf_counter()
    thread = Thread(target=_run)
    thread.start()

    text = ""
    chunk_count = 0
    try:
        for chunk in streamer:
            if not chunk:
                continue
            text += chunk
            chunk_count += 1
            yield text, chunk_count, time.perf_counter() - start, False
    finally:
        thread.join()

    if "error" in result_holder:
        raise result_holder["error"]

    elapsed = time.perf_counter() - start
    ids = result_holder.get("ids")
    if ids is not None:
        token_count = int(ids.shape[1] - inputs.input_ids.shape[1])
    else:
        token_count = chunk_count
    yield text, token_count, elapsed, True


def parse_chandra_html(html_text: str) -> list[dict]:
    blocks = []
    for match in _BLOCK_RE.finditer(html_text):
        bbox, label, content = match.groups()
        blocks.append({"bbox": bbox, "label": label, "content": content})
    return blocks


def _image_data_url(image: Image.Image) -> str:
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/png;base64,{b64}"


def _json_for_script(data) -> str:
    return json.dumps(data, ensure_ascii=False).replace("</", "<\\/")


VIEWER_HEIGHT_PX = int(os.getenv("VIEWER_HEIGHT_PX", "900"))


def _make_viewer_iframe() -> str:
    """Static viewer iframe, rendered ONCE at UI build time.

    Image and blocks are pushed in via ``postMessage`` from the parent page
    (see HEAD_SCRIPT and the change-handlers on image_signal / blocks_signal).
    Keeping the iframe stable across updates preserves zoom/scroll state and
    eliminates flicker.
    """
    srcdoc = f"""<!doctype html>
<html lang="ja">
<head>
<meta charset="utf-8">
<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.css">
<script defer src="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.js"></script>
<style>
*{{margin:0;padding:0;box-sizing:border-box}}
html,body{{height:100%}}
body{{font-family:-apple-system,"Helvetica Neue","Noto Sans JP",sans-serif;background:#1a1a2e;color:#e0e0e0;overflow:hidden}}
.app{{display:flex;flex-direction:column;height:100vh}}
.main{{display:flex;flex:1;overflow:hidden;min-height:0}}
.img-panel{{flex:1;overflow:auto;padding:12px;background:#0d1b36;min-width:280px;position:relative;cursor:grab}}
.img-panel.grabbing{{cursor:grabbing}}
.img-wrap{{position:relative;display:inline-block;line-height:0}}
.img-wrap img{{display:block;user-select:none;-webkit-user-drag:none;max-width:none}}
.bbox-ov{{position:absolute;top:0;left:0;width:100%;height:100%;pointer-events:none}}
.bb{{position:absolute;border:3px solid;opacity:.6;transition:opacity .15s,border-width .15s;pointer-events:auto;cursor:pointer}}
.bb:hover,.bb.on{{opacity:1}}
.bb.on{{border-width:5px;box-shadow:0 0 12px rgba(255,255,255,.3)}}
.v-resize{{width:4px;cursor:col-resize;background:#333;flex-shrink:0}}
.v-resize:hover{{background:#555}}
.cmp-panel{{display:flex;flex:0 0 360px;min-width:0}}
.blk-col{{flex:1;display:flex;flex-direction:column;min-width:0;overflow:hidden;background:#0f3460;border-left:1px solid #333}}
.blk-col-head{{padding:8px 12px;font-size:12px;font-weight:600;background:#16213e;border-bottom:1px solid #333;display:flex;justify-content:space-between;align-items:center;gap:8px}}
.blk-col-head .tag{{font-size:10px;padding:2px 6px;border-radius:3px;background:#1565c0;color:#fff}}
.blk-col-head .live{{font-size:10px;padding:2px 8px;border-radius:10px;display:inline-flex;align-items:center;gap:4px;background:#333;color:#888}}
.blk-col-head .live.on{{background:#0d4d2a;color:#7ee29a}}
.blk-col-head .live.on::before{{content:"";width:6px;height:6px;border-radius:50%;background:#7ee29a;animation:pulse 1s ease-in-out infinite}}
.blk-col-head .live.done{{background:#0a3d6e;color:#9ec9ff}}
.blk-col-head .live.done::before{{content:"✓";font-size:11px;line-height:1}}
@keyframes pulse{{0%,100%{{opacity:1}}50%{{opacity:.35}}}}
.blk-list{{flex:1;overflow-y:auto}}
.blk-card{{padding:8px 12px;border-bottom:1px solid rgba(255,255,255,.04);cursor:pointer;border-left:3px solid transparent}}
.blk-card:hover{{background:rgba(255,255,255,.03)}}
.blk-card.on{{background:rgba(33,150,243,.12)}}
.blk-card.fresh{{animation:flashIn .35s ease-out}}
@keyframes flashIn{{from{{background:rgba(126,226,154,.25)}}to{{background:transparent}}}}
.blk-hdr{{display:flex;align-items:center;gap:6px;margin-bottom:4px;flex-wrap:wrap}}
.blk-idx{{font-size:10px;color:#666;min-width:18px}}
.blk-tag{{font-size:10px;font-weight:600;padding:2px 7px;border-radius:3px;color:#fff}}
.blk-bb{{font-size:10px;color:#555;font-family:monospace}}
.blk-body{{font-size:12px;line-height:1.5;color:#bbb;max-height:220px;overflow-y:auto;word-break:break-word}}
.blk-body table{{border-collapse:collapse;width:100%;margin-top:4px;font-size:11px}}
.blk-body th,.blk-body td{{border:1px solid #444;padding:3px 6px;text-align:left}}
.blk-body th{{background:#1a2a4e;font-weight:600}}
.blk-body p{{margin:2px 0}}
.blk-body img{{max-width:100%;height:auto}}
.blk-empty{{padding:16px;color:#555;font-size:12px;text-align:center}}
.placeholder{{position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);color:#445;font-size:13px;pointer-events:none;text-align:center;line-height:1.6}}
</style>
</head>
<body>
<div class="app">
  <div class="main">
    <div class="img-panel" id="imgPanel">
      <div class="img-wrap" id="imgWrap">
        <img id="pgImg" alt="page" draggable="false" style="display:none">
        <div class="bbox-ov" id="bboxOv"></div>
      </div>
      <div class="placeholder" id="placeholder">📄<br>Upload an image or PDF</div>
    </div>
    <div class="v-resize" id="vResize"></div>
    <div class="cmp-panel">
      <div class="blk-col">
        <div class="blk-col-head">
          <span>Prediction</span>
          <span style="display:flex;align-items:center;gap:6px">
            <span class="live" id="liveBadge">idle</span>
            <span style="font-size:10px;color:#888" id="blkCount"></span>
          </span>
        </div>
        <div class="blk-list" id="predList"><div class="blk-empty">(no blocks parsed)</div></div>
      </div>
    </div>
  </div>
</div>
<script>
const colors = {_json_for_script(LABEL_COLORS)};
let blocks = [];
let activeIdx = -1;
const _img = document.getElementById('pgImg');
const _panel = document.getElementById('imgPanel');
const _placeholder = document.getElementById('placeholder');
const _badge = document.getElementById('liveBadge');
let _scale = 1, _fitScale = 1, _ready = false;

function setLive(state) {{
  _badge.classList.remove('on', 'done');
  if (state === 'streaming') {{ _badge.classList.add('on'); _badge.textContent = 'generating'; }}
  else if (state === 'done')  {{ _badge.classList.add('done'); _badge.textContent = 'done'; }}
  else                         {{ _badge.textContent = 'idle'; }}
}}

function esc(s) {{
  const d = document.createElement('div');
  d.textContent = s == null ? '' : String(s);
  return d.innerHTML;
}}

function formatContent(htmlStr) {{
  if (!htmlStr) return '';
  let out = htmlStr.replace(/<math([^>]*)>([\\s\\S]*?)<\\/math>/g, (_, attrs, tex) => {{
    const isBlock = attrs.includes('display="block"');
    if (typeof katex !== 'undefined') {{
      try {{ return katex.renderToString(tex, {{ displayMode: isBlock, throwOnError: false }}); }} catch(e) {{}}
    }}
    const d = document.createElement('div');
    d.textContent = tex;
    return '<code style="color:#ce93d8;background:#1a1a2e;padding:1px 4px;border-radius:2px">' + d.innerHTML + '</code>';
  }});
  if (out.match(/^\\s*<tbody[\\s>]/i) && !out.match(/<table[\\s>]/i)) {{
    out = '<table>' + out + '</table>';
  }}
  return out;
}}

function addBBox(b, i) {{
  const ov = document.getElementById('bboxOv');
  const parts = String(b.bbox || '').split(/\\s+/).map(Number);
  if (parts.length !== 4 || parts.some(v => Number.isNaN(v))) return;
  const [x0, y0, x1, y1] = parts;
  const c = colors[b.label] || '#999';
  const r = document.createElement('div');
  r.className = 'bb' + (i === activeIdx ? ' on' : '');
  r.dataset.i = i;
  r.style.left = (x0/10)+'%'; r.style.top = (y0/10)+'%';
  r.style.width = ((x1-x0)/10)+'%'; r.style.height = ((y1-y0)/10)+'%';
  r.style.borderColor = c;
  r.onclick = e => {{ e.stopPropagation(); setActive(i); }};
  ov.appendChild(r);
}}

function appendBlockCard(b, i, fresh) {{
  const c = colors[b.label] || '#999';
  const card = document.createElement('div');
  card.className = 'blk-card' + (i === activeIdx ? ' on' : '') + (fresh ? ' fresh' : '');
  card.dataset.i = i;
  if (i === activeIdx) card.style.borderLeftColor = c;
  card.innerHTML = `
    <div class="blk-hdr">
      <span class="blk-idx">#${{i}}</span>
      <span class="blk-tag" style="background:${{c}}">${{esc(b.label)}}</span>
      <span class="blk-bb">${{esc(b.bbox)}}</span>
    </div>
    <div class="blk-body">${{formatContent(b.content) || '<em style="color:#444">(empty)</em>'}}</div>`;
  card.onclick = () => setActive(i);
  document.getElementById('predList').appendChild(card);
}}

function clearAllBlocks() {{
  blocks = [];
  activeIdx = -1;
  document.getElementById('bboxOv').innerHTML = '';
  document.getElementById('predList').innerHTML = '<div class="blk-empty">(no blocks parsed)</div>';
  document.getElementById('blkCount').textContent = '';
}}

function applyBlocksUpdate(newBlocks) {{
  if (!Array.isArray(newBlocks)) return;
  if (newBlocks.length < blocks.length) {{
    clearAllBlocks();
  }}
  if (newBlocks.length > 0) {{
    const empty = document.querySelector('#predList .blk-empty');
    if (empty) empty.remove();
  }}
  for (let i = blocks.length; i < newBlocks.length; i++) {{
    const b = newBlocks[i];
    blocks.push(b);
    addBBox(b, i);
    appendBlockCard(b, i, true);
  }}
  document.getElementById('blkCount').textContent = blocks.length ? blocks.length + ' blocks' : '';
}}

function setActive(idx) {{
  const same = activeIdx === idx;
  activeIdx = same ? -1 : idx;
  document.querySelectorAll('.bb').forEach(el => {{
    el.classList.toggle('on', +el.dataset.i === activeIdx);
  }});
  document.querySelectorAll('.blk-card').forEach(el => {{
    const on = +el.dataset.i === activeIdx;
    el.classList.toggle('on', on);
    if (on) {{
      el.style.borderLeftColor = colors[blocks[+el.dataset.i].label] || '#999';
      el.scrollIntoView({{ block: 'nearest' }});
    }} else {{
      el.style.borderLeftColor = 'transparent';
    }}
  }});
}}

document.addEventListener('keydown', e => {{
  if (blocks.length === 0) return;
  if (e.key === 'j' || e.key === 'ArrowDown') {{
    e.preventDefault();
    const nxt = activeIdx < 0 ? 0 : Math.min(activeIdx + 1, blocks.length - 1);
    activeIdx = -1; setActive(nxt);
  }} else if (e.key === 'k' || e.key === 'ArrowUp') {{
    e.preventDefault();
    const prv = activeIdx < 0 ? 0 : Math.max(activeIdx - 1, 0);
    activeIdx = -1; setActive(prv);
  }} else if (e.key === 'Escape') {{
    if (activeIdx >= 0) setActive(activeIdx);
  }}
}});

document.addEventListener('mousedown', e => {{
  const h = e.target.closest('#vResize');
  if (!h) return;
  const cmp = document.querySelector('.cmp-panel');
  if (!cmp) return;
  document.body.style.cursor = 'col-resize';
  const onMove = ev => {{
    const w = Math.max(280, Math.min(document.body.clientWidth - ev.clientX, 1400));
    cmp.style.flex = '0 0 ' + w + 'px';
  }};
  const onUp = () => {{
    document.body.style.cursor = '';
    document.removeEventListener('mousemove', onMove);
    document.removeEventListener('mouseup', onUp);
  }};
  document.addEventListener('mousemove', onMove);
  document.addEventListener('mouseup', onUp);
}});

// --- Zoom & pan ---
function _applyScale() {{
  if (!_ready) return;
  _img.style.width = Math.round(_img.naturalWidth * _scale) + 'px';
  _img.style.height = 'auto';
}}
function _computeFit() {{
  const aw = Math.max(50, _panel.clientWidth - 24);
  return aw / Math.max(1, _img.naturalWidth);
}}
function _fitZoom() {{
  _fitScale = _computeFit();
  _scale = _fitScale;
  _applyScale();
  _panel.scrollLeft = 0; _panel.scrollTop = 0;
}}
function _zoomAt(factor, cx, cy) {{
  if (!_ready) return;
  const rect = _panel.getBoundingClientRect();
  const px = cx - rect.left, py = cy - rect.top;
  const oldScale = _scale;
  const maxScale = Math.max(10, _fitScale * 2);
  _scale = Math.max(_fitScale, Math.min(maxScale, _scale * factor));
  if (_scale === oldScale) return;
  _applyScale();
  const r = _scale / oldScale;
  _panel.scrollLeft = (_panel.scrollLeft + px) * r - px;
  _panel.scrollTop  = (_panel.scrollTop  + py) * r - py;
}}

function applySetImage(url) {{
  if (!url) return;
  _ready = false;
  clearAllBlocks();
  _img.style.display = '';
  _placeholder.style.display = 'none';
  const onLoad = () => {{ _ready = true; _fitZoom(); }};
  _img.addEventListener('load', onLoad, {{ once: true }});
  _img.addEventListener('error', () => {{ _ready = true; }}, {{ once: true }});
  _img.src = url;
}}

window.addEventListener('message', (e) => {{
  const d = e.data || {{}};
  if (d.type === 'oc-set-image') applySetImage(d.url);
  else if (d.type === 'oc-update') applyBlocksUpdate(d.blocks || []);
  else if (d.type === 'oc-state')  setLive(d.state);
}});

_panel.addEventListener('wheel', e => {{
  if (!(e.ctrlKey || e.metaKey)) return;
  e.preventDefault();
  const factor = e.deltaY < 0 ? 1.15 : 1/1.15;
  _zoomAt(factor, e.clientX, e.clientY);
}}, {{ passive: false }});

let _dragging = false, _dragX = 0, _dragY = 0, _startSL = 0, _startST = 0;
_panel.addEventListener('mousedown', e => {{
  if (e.target.closest('.bb')) return;
  _dragging = true;
  _dragX = e.clientX; _dragY = e.clientY;
  _startSL = _panel.scrollLeft; _startST = _panel.scrollTop;
  _panel.classList.add('grabbing');
  e.preventDefault();
}});
window.addEventListener('mousemove', e => {{
  if (!_dragging) return;
  _panel.scrollLeft = _startSL - (e.clientX - _dragX);
  _panel.scrollTop  = _startST - (e.clientY - _dragY);
}});
window.addEventListener('mouseup', () => {{
  if (!_dragging) return;
  _dragging = false;
  _panel.classList.remove('grabbing');
}});

window.addEventListener('resize', () => {{
  if (Math.abs(_scale - _fitScale) < 1e-3) _fitZoom();
  else _fitScale = _computeFit();
}});

document.addEventListener('keydown', e => {{
  if (e.key === '+' || e.key === '=') {{
    e.preventDefault();
    const r = _panel.getBoundingClientRect();
    _zoomAt(1.2, r.left + r.width/2, r.top + r.height/2);
  }} else if (e.key === '-' || e.key === '_') {{
    e.preventDefault();
    const r = _panel.getBoundingClientRect();
    _zoomAt(1/1.2, r.left + r.width/2, r.top + r.height/2);
  }} else if (e.key === '0') {{
    e.preventDefault();
    _fitZoom();
  }}
}});

// Tell the parent we're ready for state replays.
try {{ window.parent.postMessage({{type: 'oc-ready'}}, '*'); }} catch(_) {{}}
</script>
</body>
</html>"""
    return (
        f'<iframe class="oc-viewer" style="width:100%;height:{VIEWER_HEIGHT_PX}px;'
        'border:0;border-radius:8px;background:#1a1a2e;display:block" '
        f'sandbox="allow-scripts" srcdoc="{html.escape(srcdoc, quote=True)}"></iframe>'
    )


_EMPTY_BLOCKS_JSON = "[]"


def preview_uploaded(file_path: str | None, page_number: int):
    """Render the uploaded image (or selected PDF page) without running the model.

    Outputs:
        gallery, html_code, markdown_code, image_signal, blocks_signal,
        state_signal, metadata, status
    """
    if not file_path:
        return [], "", "", "", _EMPTY_BLOCKS_JSON, "idle", [], ""

    try:
        page_num, image = _load_image(file_path, int(page_number or 1))
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Failed to load file: {e}") from e

    status = (
        f"Loaded `{Path(file_path).name}` · page `{page_num}` · "
        f"`{image.width}x{image.height}` — click **Run OCR** to predict."
    )
    meta = [
        {
            "file": Path(file_path).name,
            "page": page_num,
            "image_size": f"{image.width}x{image.height}",
        }
    ]
    return (
        [image],
        "",
        "",
        _image_data_url(image),
        _EMPTY_BLOCKS_JSON,
        "idle",
        meta,
        status,
    )


def run_ocr(
    file_path: str | None,
    page_number: int,
    temperature: float,
):
    """Streaming OCR. Yields:
        gallery, html_code, markdown_code, image_signal, blocks_signal,
        state_signal, metadata, status
    """
    if not file_path:
        raise gr.Error("Please upload an image or PDF.")

    page_num, image = _load_image(file_path, int(page_number))
    image_url = _image_data_url(image)

    # First yield: clear any previous blocks, mark "streaming", ensure image set.
    yield (
        [image],
        "",
        gr.skip(),
        image_url,
        _EMPTY_BLOCKS_JSON,
        "streaming",
        gr.skip(),
        f"⏳ Generating… page `{page_num}` · `{image.width}x{image.height}`",
    )

    last_block_count = 0
    text = ""
    token_count = 0
    elapsed = 0.0

    for text, token_count, elapsed, done in _generate_page_stream(
        image, temperature
    ):
        if done:
            break
        blocks = parse_chandra_html(text)
        cur_count = len(blocks)
        if cur_count != last_block_count:
            blocks_out = json.dumps(blocks, ensure_ascii=False)
            last_block_count = cur_count
        else:
            blocks_out = gr.skip()
        status = (
            f"⏳ Generating… `{cur_count}` blocks · `{token_count}` chunks · "
            f"`{round(elapsed, 1)}s`"
        )
        yield (
            [image],
            text,
            gr.skip(),
            gr.skip(),
            blocks_out,
            gr.skip(),
            gr.skip(),
            status,
        )

    try:
        markdown = parse_markdown(text)
    except Exception as e:
        markdown = f"<!-- markdown conversion failed: {e} -->"

    blocks = parse_chandra_html(text)
    blocks_json = json.dumps(blocks, ensure_ascii=False)
    metadata = [
        {
            "page": page_num,
            "tokens": token_count,
            "seconds": round(elapsed, 2),
            "image_size": f"{image.width}x{image.height}",
            "blocks": len(blocks),
        }
    ]
    status = (
        f"### ✅ Done — `{len(blocks)}` blocks · `{token_count}` tokens · "
        f"`{round(elapsed, 2)}s`\n"
        f"Model: `{MODEL_ID}` · Page: `{page_num}`"
    )
    yield (
        [image],
        text,
        markdown,
        gr.skip(),
        blocks_json,
        "done",
        metadata,
        status,
    )


CSS = """
.gradio-container { max-width: 100% !important; padding: 8px 16px !important; }
.viewer-host > .label-wrap { display: none !important; }
.viewer-host .gradio-html { padding: 0 !important; }
.input-pane { border-right: 1px solid #d0d7de; padding-right: 12px !important; }
.section-h { margin: 4px 0 4px !important; font-size: 13px !important; font-weight: 600 !important; color: #444 !important; }
.model-link { margin: 0 0 8px !important; font-size: 12px !important; color: #555 !important; }
"""


HEAD_SCRIPT = """
<script>
(function() {
  if (window._ocBus) return;
  window._ocBus = { latestImage: null, latestBlocks: null, latestState: null };
  function _iframes() { return document.querySelectorAll('iframe.oc-viewer'); }
  function _post(msg, target) {
    const targets = target ? [target] : Array.from(_iframes()).map(f => f.contentWindow);
    targets.forEach(w => { if (!w) return; try { w.postMessage(msg, '*'); } catch(_) {} });
  }
  window._ocSetImage = function(url) {
    if (!url) return;
    window._ocBus.latestImage = url;
    // New image: reset block cache too — sender is expected to follow up with []
    window._ocBus.latestBlocks = [];
    _post({type:'oc-set-image', url});
  };
  window._ocSetBlocks = function(blocks) {
    if (!Array.isArray(blocks)) return;
    window._ocBus.latestBlocks = blocks;
    _post({type:'oc-update', blocks});
  };
  window._ocSetState = function(state) {
    window._ocBus.latestState = state;
    _post({type:'oc-state', state});
  };
  window.addEventListener('message', (e) => {
    if (!e.data || e.data.type !== 'oc-ready') return;
    const b = window._ocBus;
    if (b.latestImage)   _post({type:'oc-set-image', url: b.latestImage}, e.source);
    if (b.latestBlocks)  _post({type:'oc-update', blocks: b.latestBlocks},  e.source);
    if (b.latestState)   _post({type:'oc-state', state: b.latestState},     e.source);
  });
})();
</script>
"""


with gr.Blocks(title="Qwen3.5-OCR-JP-2B", css=CSS, head=HEAD_SCRIPT) as demo:
    gr.Markdown(
        "Model: [`ebinan92/Qwen3.5-ocr-jp-2b`](https://huggingface.co/ebinan92/Qwen3.5-ocr-jp-2b)",
        elem_classes=["model-link"],
    )

    with gr.Row():
        with gr.Column(scale=1, min_width=300, elem_classes=["input-pane"]):
            input_file = gr.File(
                label="Image / PDF",
                file_types=["image", ".pdf"],
                type="filepath",
            )
            page_number = gr.Number(
                label="PDF page number",
                value=1,
                precision=0,
                minimum=1,
                info="Used only for PDF input (ignored for images).",
            )
            temperature = gr.Slider(
                0.0, 1.0, value=0.0, step=0.05, label="Temperature"
            )
            run_button = gr.Button("Run OCR", variant="primary")
            status = gr.Markdown()

            with gr.Accordion("Generated HTML / Markdown", open=False):
                with gr.Tabs():
                    with gr.Tab("Generated HTML"):
                        html_code = gr.Code(
                            label="Generated Chandra HTML",
                            language="html",
                            lines=18,
                        )
                    with gr.Tab("Markdown"):
                        markdown_code = gr.Code(
                            label="Markdown (chandra parse_markdown)",
                            language="markdown",
                            lines=18,
                        )
                    with gr.Tab("Input pages"):
                        gallery = gr.Gallery(
                            label="Input pages",
                            columns=1,
                            height=260,
                            object_fit="contain",
                        )
                    with gr.Tab("Metadata"):
                        metadata = gr.JSON(label="Run metadata")

        with gr.Column(scale=4):
            preview = gr.HTML(
                value=_make_viewer_iframe(), elem_classes=["viewer-host"]
            )

    # Hidden signals: changes are forwarded to the iframe via postMessage so the
    # iframe is never re-rendered (preserves zoom/scroll, kills flicker).
    image_signal = gr.Textbox(visible=False)
    blocks_signal = gr.Textbox(visible=False, value=_EMPTY_BLOCKS_JSON)
    state_signal = gr.Textbox(visible=False, value="idle")

    image_signal.change(
        None,
        inputs=[image_signal],
        outputs=[],
        js="(url) => { if (url && window._ocSetImage) window._ocSetImage(url); }",
    )
    blocks_signal.change(
        None,
        inputs=[blocks_signal],
        outputs=[],
        js=(
            "(s) => { if (s == null) return;"
            " let b; try { b = JSON.parse(s); } catch(_) { return; }"
            " if (window._ocSetBlocks) window._ocSetBlocks(b); }"
        ),
    )
    state_signal.change(
        None,
        inputs=[state_signal],
        outputs=[],
        js="(s) => { if (window._ocSetState) window._ocSetState(s); }",
    )

    common_outputs = [
        gallery,
        html_code,
        markdown_code,
        image_signal,
        blocks_signal,
        state_signal,
        metadata,
        status,
    ]

    input_file.change(
        preview_uploaded,
        inputs=[input_file, page_number],
        outputs=common_outputs,
    )
    page_number.change(
        preview_uploaded,
        inputs=[input_file, page_number],
        outputs=common_outputs,
    )

    run_button.click(
        run_ocr,
        inputs=[input_file, page_number, temperature],
        outputs=common_outputs,
        show_progress="full",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=int(os.getenv("CONCURRENCY_LIMIT", "1")))
    demo.launch()