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

# Pre-built causal_conv1d wheel (linear-attention dependency, can't compile in Space build).
# Pulled from Dao-AILab official GitHub release; matches torch==2.6.0 / cu12 / py3.10 / cxx11abi=FALSE.
_CAUSAL_CONV1D_WHEEL = (
    "https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.5.0.post8/"
    "causal_conv1d-1.5.0.post8+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
)
try:
    import causal_conv1d  # noqa: F401
except ImportError:
    subprocess.run(
        [sys.executable, "-m", "pip", "install", _CAUSAL_CONV1D_WHEEL, "--no-deps", "-q"],
        check=True,
    )
    import causal_conv1d  # noqa: F401

import gradio as gr
import spaces
from PIL import Image


MODEL_ID = os.getenv("MODEL_ID", "ebinan92/open-chandra-stage2-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(
    image: Image.Image,
    temperature: float,
) -> tuple[str, int, float]:
    """Inference path mirrors victor/chandra-ocr-2 exactly."""
    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)

    start = time.perf_counter()
    generated_ids = MODEL.generate(
        **inputs,
        max_new_tokens=MAX_NEW_TOKENS,
        eos_token_id=eos_token_id,
    )
    elapsed = time.perf_counter() - start

    new_tokens = generated_ids.shape[1] - inputs.input_ids.shape[1]
    generated_ids_trimmed = generated_ids[0][len(inputs.input_ids[0]):]
    text = PROCESSOR.decode(
        generated_ids_trimmed,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )
    return text, int(new_tokens), elapsed


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 _interactive_viewer_html(
    image: Image.Image,
    blocks: list[dict],
    title: str = "",
) -> str:
    header_inner = (
        f'<div class="hdr-title">{html.escape(title)}</div>' if title else ""
    )
    header_block = (
        f'<div class="hdr">{header_inner}</div>' if header_inner else ""
    )
    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}}
.hdr{{background:#16213e;padding:6px 12px;border-bottom:1px solid #333;display:flex;align-items:center;gap:12px;flex-shrink:0}}
.hdr-title{{font-size:13px;font-weight:600;color:#ddd;white-space:nowrap;overflow:hidden;text-overflow:ellipsis}}
.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}}
.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)}}
.bb-lbl{{position:absolute;top:-17px;left:0;font-size:10px;padding:1px 4px;color:#fff;border-radius:2px;white-space:nowrap;pointer-events:none}}
.zoom-ctl{{position:absolute;bottom:12px;left:12px;display:flex;gap:4px;background:rgba(15,52,96,.92);padding:4px 6px;border-radius:6px;z-index:10;align-items:center;font-size:11px;color:#bbb;border:1px solid #333;user-select:none}}
.zoom-ctl button{{padding:2px 8px;background:#0f3460;border:1px solid #333;color:#ccc;border-radius:3px;cursor:pointer;font-size:13px;line-height:1;min-width:24px;font-weight:600}}
.zoom-ctl button:hover{{background:#16498a;color:#fff}}
.zoom-ctl .pct{{min-width:42px;text-align:center;font-family:monospace;color:#bbb}}
.v-resize{{width:4px;cursor:col-resize;background:#333;flex-shrink:0}}
.v-resize:hover{{background:#555}}
.cmp-panel{{display:flex;flex:0 0 460px;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}}
.blk-col-head .tag{{font-size:10px;padding:2px 6px;border-radius:3px;background:#1565c0;color:#fff}}
.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-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}}
</style>
</head>
<body>
<div class="app">
  {header_block}
  <div class="main">
    <div class="img-panel" id="imgPanel">
      <div class="img-wrap" id="imgWrap">
        <img id="pgImg" src="{_image_data_url(image)}" alt="page" draggable="false">
        <div class="bbox-ov" id="bboxOv"></div>
      </div>
      <div class="zoom-ctl" title="Ctrl/Cmd + wheel to zoom · drag to pan · +/- /0 keys">
        <button id="zOut" title="Zoom out (-)">−</button>
        <span class="pct" id="zPct">—</span>
        <button id="zIn" title="Zoom in (+)">+</button>
        <button id="zFit" title="Fit (0)">Fit</button>
        <button id="zOne" title="Actual size">1:1</button>
      </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 class="tag">Pred</span></span><span style="font-size:10px;color:#888" id="blkCount"></span></div>
        <div class="blk-list" id="predList"></div>
      </div>
    </div>
  </div>
</div>
<script>
const blocks = {_json_for_script(blocks)};
const colors = {_json_for_script(LABEL_COLORS)};
let activeIdx = -1;
document.getElementById('blkCount').textContent = blocks.length + ' blocks';

function renderBBoxes() {{
  const ov = document.getElementById('bboxOv');
  if (!ov) return;
  ov.innerHTML = '';
  blocks.forEach((b, i) => addBBox(ov, b, i));
}}

function addBBox(ov, b, i) {{
  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;
  const l = document.createElement('div');
  l.className = 'bb-lbl';
  l.textContent = b.label;
  l.style.background = c;
  r.appendChild(l);
  r.onclick = e => {{ e.stopPropagation(); setActive(i); }};
  ov.appendChild(r);
}}

function renderBlocks() {{
  const list = document.getElementById('predList');
  list.innerHTML = '';
  if (blocks.length === 0) {{
    list.innerHTML = '<div class="blk-empty">(no blocks parsed)</div>';
    return;
  }}
  blocks.forEach((b, i) => {{
    const c = colors[b.label] || '#999';
    const card = document.createElement('div');
    card.className = 'blk-card' + (i === activeIdx ? ' on' : '');
    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);
    list.appendChild(card);
  }});
}}

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);
  }}
}});

(function setupResize(){{
  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);
  }});
}})();

function formatContent(html) {{
  if (!html) return '';
  let out = html.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 esc(s) {{
  const d = document.createElement('div');
  d.textContent = s == null ? '' : String(s);
  return d.innerHTML;
}}

renderBlocks();

// --- Zoom & pan ---
const _img = document.getElementById('pgImg');
const _panel = document.getElementById('imgPanel');
const _pctEl = document.getElementById('zPct');
let _scale = 1;
let _fitScale = 1;
let _ready = false;

function _applyScale() {{
  if (!_ready) return;
  _img.style.width = Math.round(_img.naturalWidth * _scale) + 'px';
  _img.style.height = 'auto';
  _pctEl.textContent = Math.round(_scale * 100) + '%';
}}

function _computeFit() {{
  const aw = Math.max(50, _panel.clientWidth - 24);
  return aw / _img.naturalWidth;  // fit by width — fills panel horizontally
}}

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;
  const 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 _initZoom() {{
  if (_img.naturalWidth === 0) {{
    _img.addEventListener('load', _initZoom, {{ once: true }});
    _img.addEventListener('error', () => {{ _ready = true; _applyScale(); renderBBoxes(); }}, {{ once: true }});
    return;
  }}
  _ready = true;
  _fitZoom();
  renderBBoxes();
}}

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

document.getElementById('zIn').onclick = () => {{
  const r = _panel.getBoundingClientRect();
  _zoomAt(1.2, r.left + r.width/2, r.top + r.height/2);
}};
document.getElementById('zOut').onclick = () => {{
  const r = _panel.getBoundingClientRect();
  _zoomAt(1/1.2, r.left + r.width/2, r.top + r.height/2);
}};
document.getElementById('zFit').onclick = _fitZoom;
document.getElementById('zOne').onclick = () => {{
  const r = _panel.getBoundingClientRect();
  const oldScale = _scale;
  const px = r.width/2, py = r.height/2;
  _scale = Math.max(_fitScale, 1);
  _applyScale();
  const ratio = _scale / oldScale;
  _panel.scrollLeft = (_panel.scrollLeft + px) * ratio - px;
  _panel.scrollTop = (_panel.scrollTop + py) * ratio - py;
}};

// --- Pan via drag ---
let _dragging = false, _dragX = 0, _dragY = 0, _startSL = 0, _startST = 0;
_panel.addEventListener('mousedown', e => {{
  if (e.target.closest('.bb') || e.target.closest('.zoom-ctl')) 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');
}});

// Re-fit on iframe resize while user has not zoomed away from fit
let _wasFit = true;
window.addEventListener('resize', () => {{
  if (Math.abs(_scale - _fitScale) < 1e-3) {{
    _fitZoom();
  }} else {{
    _fitScale = _computeFit();
  }}
}});

// Keyboard zoom shortcuts
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();
  }}
}});

_initZoom();
</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>'
    )


def _preview_html(image: Image.Image, content: str) -> str:
    blocks = parse_chandra_html(content)
    return _interactive_viewer_html(image, blocks, title="")


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

    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

    preview = _preview_html(image, "")
    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], "", preview, meta, status


def run_ocr(
    file_path: str | None,
    page_number: int,
    temperature: float,
    progress: gr.Progress = gr.Progress(track_tqdm=False),
):
    if not file_path:
        raise gr.Error("Please upload an image or PDF.")

    progress(0.0, desc="Preparing page")
    page_num, image = _load_image(file_path, int(page_number))

    progress(0.3, desc=f"Running OCR on page {page_num}")
    text, token_count, elapsed = _generate_page(
        image,
        temperature=temperature,
    )
    progress(1.0, desc="Done")

    metadata = [
        {
            "page": page_num,
            "tokens": token_count,
            "seconds": round(elapsed, 2),
            "image_size": f"{image.width}x{image.height}",
        }
    ]
    status = (
        f"Model: `{MODEL_ID}`  \n"
        f"Page: `{page_num}` / Tokens: `{token_count}` / Seconds: `{round(elapsed, 2)}`"
    )
    return (
        [image],
        text,
        _preview_html(image, text),
        metadata,
        status,
    )


CSS = """
.gradio-container { max-width: 100% !important; padding: 8px 16px !important; }
.app-title { margin: 0 0 4px !important; font-size: 18px !important; }
.app-sub { margin: 0 0 12px !important; font-size: 12px !important; color: #666 !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; }
"""


with gr.Blocks(title="Open Chandra OCR", css=CSS) as demo:
    gr.Markdown("# Open Chandra OCR", elem_classes=["app-title"])
    gr.Markdown(
        "Upload an image or PDF and run OCR. Predicted blocks are shown in the right pane.",
        elem_classes=["app-sub"],
    )

    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("Raw output / metadata / pages", 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("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(elem_classes=["viewer-host"])

    input_file.change(
        preview_uploaded,
        inputs=[input_file, page_number],
        outputs=[gallery, html_code, preview, metadata, status],
    )
    page_number.change(
        preview_uploaded,
        inputs=[input_file, page_number],
        outputs=[gallery, html_code, preview, metadata, status],
    )

    run_button.click(
        run_ocr,
        inputs=[input_file, page_number, temperature],
        outputs=[gallery, html_code, preview, metadata, status],
    )


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