"""Progressive document workspace: choose a PDF, run a decision, inspect the pages."""

from __future__ import annotations

import html
import logging
from functools import lru_cache

import blink
import documents
import gradio as gr
import ui

e = html.escape
LOG = logging.getLogger(__name__)
LABELS = {"text": "Selectable text", "ocr": "English OCR", "blank": "Nearly blank page",
          "partial": "Partial text / review", "unreadable": "No extracted text / review"}
REVISION = r"""(...a) => {
  a[0] = String(window.__blinkDocRev = (window.__blinkDocRev || 0) + 1);
  return a;
}"""
READ_REVISION = r"""(...a) => {
  a[0] = String(window.__blinkDocRev || 0);
  return a;
}"""
GUARD_JS = r"""
(() => {
  if (window.__blinkDocGuard) return;
  window.__blinkDocGuard = true;
  document.addEventListener('change', (event) => {
    const input = event.target;
    if (!(input instanceof HTMLInputElement) || input.type !== 'file' || !input.closest('#doc-upload')) return;
    const file = input.files?.[0];
    if (!file || !file.name.toLowerCase().endsWith('.pdf') || file.size > 12 * 1024 * 1024) return;
    window.__blinkDocRev = (window.__blinkDocRev || 0) + 1;
    const node = document.getElementById('doc-run');
    const button = node instanceof HTMLButtonElement ? node : node?.querySelector('button');
    if (button) button.disabled = true;
  }, true);
  setInterval(() => {
    const panel = document.querySelector('#doc-result [data-doc-rev]');
    const stale = panel && panel.dataset.docRev !== String(window.__blinkDocRev || 0);
    const outputs = document.getElementById('doc-outputs');
    if (outputs) outputs.classList.toggle('doc-stale', Boolean(stale));
  }, 150);
})();
"""


@lru_cache(maxsize=3)
def sample_document(key: str) -> documents.Document:
    return documents.parse_document(documents.sample_path(key))


FILE_ICON = ('<svg viewBox="0 0 24 24" aria-hidden="true"><path d="M14 3H7a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h10a2 2 0 0 0 '
             '2-2V8z"/><path d="M14 3v5h5M9 13h6M9 17h6"/></svg>')
SOURCES = {"saved": "Saved model run", "live": "Live model run", "mock": "Mock: fabricated scores"}
DOT = " \u00b7 "
FLAG = ' class="flag"'
COMPARE_HEAD = '<p class="doc-compare-head">Same PDF and rules, other model</p>'


def intro() -> str:
    return (
        '<div class="doc-intro"><h2>PDF in. Documents out.</h2>'
        '<p class="blk-note">Label a document or separate a packet. See the original pages, scores, and boundaries.</p>'
        '</div>'
    )


def step(number: int, title: str) -> str:
    return f'<h3 class="doc-step"><span>{number}</span>{e(title)}</h3>'


def _count(n: int, word: str) -> str:
    return f"{n} {word}" if n == 1 else f"{n} {word}s"


def page_choices(document: documents.Document) -> list[tuple[str, int]]:
    return [(f"Page {p.number} / {len(document.pages)}", p.number) for p in document.pages]


def preview_html(document: documents.Document | None, number: int = 1) -> str:
    if document is None:
        return '<div class="doc-empty"><p>Choose a sample or upload a PDF to see its pages.</p></div>'
    page = document.pages[int(number) - 1]
    return (
        '<figure class="doc-preview"><div class="doc-sheet">'
        f'<img src="data:image/jpeg;base64,{page.preview}" '
        f'alt="Original page {page.number} of {len(document.pages)}"></div>'
        f'<figcaption><span>Page {page.number} / {len(document.pages)}</span>'
        f'<span class="doc-how">{e(LABELS[page.method])}</span></figcaption></figure>'
    )


def source_html(document: documents.Document, sample: str | None) -> str:
    scans = sum(p.method == "ocr" for p in document.pages)
    review = sum(p.method in ("partial", "unreadable") for p in document.pages)
    caption = "Original fictional sample. Not an accuracy benchmark." if sample else "Your uploaded PDF."
    facts = [_count(len(document.pages), "page")]
    if scans:
        facts.append(f"{scans} OCR")
    elif not review:
        facts.append("selectable text")
    if review:
        facts.append(f"{_count(review, 'page')} need text review")
    return (
        f'<div class="doc-source"><span class="doc-file">{FILE_ICON}</span>'
        f'<p><b>{e(document.name)}</b><span>{DOT.join(facts)}</span></p>'
        f'<p class="doc-caption">{caption}</p></div>'
    )


def _range(numbers: list[int]) -> str:
    return f"Page {numbers[0]}" if len(numbers) == 1 else f"Pages {numbers[0]}-{numbers[-1]}"


def _bars(segment: dict, labels: dict) -> str:
    return "".join(
        f'<li class="blk-opt{" win" if key == segment["category"] else ""}" style="--i:{i}">'
        f'<span class="k">{e(labels.get(key, key))}</span>'
        f'<span class="track"><i style="width:{max(p * 100, 0.6):.1f}%"></i></span>'
        f'<span class="v">{p:.1%}</span></li>'
        for i, (key, p) in enumerate(sorted(segment["probabilities"].items(), key=lambda item: -item[1]))
    )


def _reasons(segment: dict) -> str:
    items = "".join(f'<li>{e(r["message"])} ({_range(r["pages"])})</li>' for r in segment["review_reasons"])
    return f'<ul class="doc-review-reasons">{items}</ul>' if items else ""


def result_html(result: dict | None, revision: str = "0") -> str:
    if result is None:
        return ('<div class="doc-empty"><p>Ready to run. Blink scores your categories and page boundaries '
                'without writing text.</p></div>')
    segments = result["segments"]
    model = str(result["model"]).split("/")[-1]
    kind = result["source"]
    count = len(segments)
    review = sum(bool(s["review_reasons"]) for s in segments)
    pages = result["document"]["page_count"]
    labels = {c["id"]: c["label"] for c in result["rules"]["categories"]}
    split = result["mode"] != "classify"
    title = f"{count} documents" if split else segments[0]["label"]
    needs = "document needs" if review == 1 else "documents need"
    note = _count(pages, "original page") + (f"{DOT}{review} {needs} review" if review else f"{DOT}no review flags")
    if split:
        timeline = "".join(
            f'<li style="--pages:{len(s["pages"])}"{FLAG if s["review_reasons"] else ""} '
            f'title="{e(s["label"])}{DOT}{_range(s["pages"])}">'
            f'<b>{e(s["label"])}</b><span>{_range(s["pages"])}</span></li>'
            for s in segments
        )
        rows = "".join(
            '<li class="doc-segment"><details class="doc-seg"><summary>'
            f'<span class="doc-seg-name"><b>{e(s["label"])}</b><span>{_range(s["pages"])}</span>'
            + ('<em class="doc-review">Review</em>' if s["review_reasons"] else "")
            + '</span><span class="doc-seg-p"><span class="doc-sr">Option probabilities, top score </span>'
            f'{s["probabilities"][s["category"]]:.1%}</span></summary>'
            f'<ul class="blk-opts">{_bars(s, labels)}</ul></details>{_reasons(s)}</li>'
            for s in segments
        )
        body = (f'<ol class="doc-timeline" aria-label="Documents in page order">{timeline}</ol>'
                f'<p class="doc-list-head"><span>Document</span><span>Top score</span></p>'
                f'<ol class="doc-segments">{rows}</ol>')
    else:
        segment = segments[0]
        body = (f'<div class="doc-scores"><p class="doc-list-head"><span>Option probabilities</span>'
                f'<span>{_range(segment["pages"])}</span></p>'
                f'<ul class="blk-opts">{_bars(segment, labels)}</ul>{_reasons(segment)}</div>')
    model_ms = result["metrics"].get("model_ms")
    call_ms = result["metrics"].get("latency_ms")
    timing = (f'<span><b>{model_ms:,.0f} ms</b> model only</span>' if model_ms is not None else
              f'<span><b>{call_ms:,.0f} ms</b> decision call</span>' if call_ms is not None else "")
    return (
        f'<div class="doc-result {e(kind)}" data-doc-rev="{e(str(revision))}">'
        f'<div class="doc-result-head"><p class="doc-provenance"><span class="doc-dot"></span>'
        f'{SOURCES[kind]} / {e(model)}</p>'
        f'<h3>{e(title)}</h3><p class="doc-verdict-note">{note}</p></div>'
        + body
        + '<div class="doc-metrics">'
        + (f'<span><b>{result["document"]["preparation_ms"]:,.0f} ms</b> PDF preparation</span>'
           if kind != "saved" else "")
        + timing
        + '<span><b>0</b> generated tokens</span></div>'
        '<details class="doc-method"><summary>What the scores and timings mean</summary>'
        f'<p>{e(result["probability_note"])}</p>'
        '<p>Review flags use fixed demo thresholds, not a fitted confidence guarantee. '
        'Nearly blank pages stay with the preceding document; leading blank pages stay with the first. '
        'Unreadable pages are attached provisionally and always need review.</p>'
        '<p>Model-only time excludes PDF preparation and the wait for a device. '
        'Preparation includes text extraction, scan OCR when needed, and preview rendering. '
        'Saved model timings describe the recorded Space call, not this page load. '
        'The saved JSON also retains separately measured preparation from the recording client; '
        'it is not a hosted parse measurement and is omitted from this timing strip.</p>'
        + (f'<p>Decision call: {call_ms:,.0f} ms. This includes the device wait inside the call.</p>'
           if call_ms is not None else "")
        + '</details></div>'
    )


def _cache_exports(document: documents.Document, result: dict, component: gr.File) -> list[str]:
    paths = []
    for name, data in documents.export_bytes(document, result).items():
        path = gr.processing_utils.save_bytes_to_cache(data, name, component.GRADIO_CACHE)
        component.temp_files.add(path)
        paths.append(path)
    return paths


def document_api(file: str, mode: str, category_rules: str, model: str) -> dict:
    """File-component preprocessing permits only uploaded files, not arbitrary local paths."""
    try:
        document = documents.parse_document(file)
        return documents.process(document, category_rules, mode, model)
    except (documents.DocumentError, blink.BlinkError) as exc:
        raise gr.Error(str(exc)) from None


def build() -> dict:
    sample = sample_document("packet")
    model = documents.default_model()
    saved = documents.saved_result(sample, documents.DEFAULT_RULES, "split", model)
    held = gr.State(sample)
    result = gr.State(saved)
    source_seen = gr.State({"source": 0, "selection": ("packet", None)})
    rev = gr.Textbox("0", visible=False)
    gr.HTML(intro())
    with gr.Row(equal_height=False, elem_classes="doc-workspace"):
        with gr.Column(scale=5, min_width=280, elem_classes="doc-col doc-col-source"):
            gr.HTML(step(1, "Choose a PDF"), elem_classes="doc-stephead")
            preset = gr.Radio(choices=[(s["title"], s["id"]) for s in documents.sample_info()],
                              value="packet", show_label=False, elem_id="doc-samples")
            with gr.Accordion("Use your own PDF", open=False, elem_id="doc-own"):
                upload = gr.File(label="PDF \u00b7 up to 12 MiB and 12 pages", file_types=[".pdf"], type="filepath",
                                 elem_id="doc-upload")
                gr.HTML('<p class="doc-caption">Processing runs in this Space. Temporary uploads and downloads '
                        'are cleared within an hour. Do not upload confidential documents.</p>')
            source = gr.HTML(source_html(sample, "packet"), elem_id="doc-source")
            page = gr.Dropdown(choices=page_choices(sample), value=1, label="Original page", elem_id="doc-page")
            preview = gr.HTML(preview_html(sample), elem_id="doc-preview")
            with gr.Accordion("Extracted page text", open=False, elem_id="doc-text-fold"):
                text = gr.Textbox(sample.pages[0].text, label="Text that Blink reads", lines=12,
                                  interactive=False, elem_id="doc-text")
        with gr.Column(scale=6, min_width=280, elem_classes="doc-col doc-col-task"):
            gr.HTML(step(2, "Choose a task"), elem_classes="doc-stephead")
            mode = gr.Radio([("Classify", "classify"), ("Split packet", "split")], value="split",
                            show_label=False, elem_id="doc-mode")
            gr.HTML('<p class="doc-caption">Classify gives one label. Split packet finds separate documents, '
                    'including adjacent documents with the same label.</p>')
            with gr.Accordion("Categories and rules", open=False, elem_id="doc-rules-fold"):
                rule = gr.Code(documents.rules_text(), language="json", label="Your categories and boundary rules",
                               lines=18, wrap_lines=True, elem_id="doc-rules")
            with gr.Accordion("Run options", open=False, elem_id="doc-options"):
                models = gr.State(model)
                compare = gr.Checkbox(False, label="Compare with Blink 4B on the same page text",
                                      visible=len(blink.models()) == 2, elem_id="doc-compare")
                gr.HTML('<p class="doc-caption">The PDF is prepared once. Each model gets identical text and rules. '
                        'This is a demonstration, not a matched accuracy study.</p>')
            run = gr.Button("Run Blink", variant="primary", elem_id="doc-run")
            gr.HTML('<p class="doc-caption doc-runnote">This button runs the live model. The initial result, '
                    'when present, is a saved run on the fictional sample.</p>')
            gr.HTML(step(3, "Inspect the result"), elem_classes="doc-stephead")
            with gr.Column(elem_id="doc-outputs", elem_classes="doc-panel"):
                answer = gr.HTML(result_html(saved), elem_id="doc-result")
                inspect = gr.Dropdown(
                    choices=[(f'{s["label"]} / {_range(s["pages"])}', s["pages"][0]) for s in saved["segments"]]
                    if saved else [], value=None, label="Jump to a document", visible=saved is not None,
                    info="Shows the document's first page in the preview.", elem_id="doc-inspect",
                )
                comparison = gr.HTML(visible=False, elem_id="doc-comparison")
                files = gr.File(label="Download result JSON and split PDFs", file_count="multiple",
                                interactive=False, visible=saved is not None, elem_id="doc-downloads")
                if saved is not None:
                    initial = files.postprocess(_cache_exports(sample, saved, files))
                    files.value = gr.processing_utils.move_files_to_cache(
                        initial.model_dump(), files, postprocess=True, keep_in_cache=True,
                    )
                with gr.Accordion("Full result JSON", open=False, elem_id="doc-json-fold"):
                    raw = gr.JSON(saved, label="Page decisions, review flags, and provenance", elem_id="doc-json")
            gr.HTML('<p class="doc-stale-notice">Inputs changed. Run Blink again to see a matching result.</p>')
    def clear(revision):
        return (None, result_html(None), gr.update(value=None, visible=False),
                gr.update(value=None, visible=False), None, gr.update(value="", visible=False), revision)

    clear_out = [result, answer, files, inspect, raw, comparison, rev]
    for control in (mode, rule, compare):
        control.input(clear, [rev], clear_out, js=REVISION, queue=False,
                      show_progress="hidden", api_visibility="private")

    def load_source(revision, key, uploaded, current_mode, current_rules, current_model):
        try:
            document = documents.parse_document(uploaded) if uploaded else sample_document(key)
        except (documents.DocumentError, blink.BlinkError) as exc:
            return (None, gr.update(choices=[], value=None), preview_html(None), "", "",
                    None, f'<div class="blk-warn">{e(str(exc))}</div>', gr.update(value=None, visible=False),
                    gr.update(choices=[], value=None, visible=False), None,
                    gr.update(value="", visible=False), revision)
        warning = ""
        try:
            recorded = documents.saved_result(document, current_rules, current_mode, current_model)
        except documents.DocumentError as exc:
            recorded = None
            warning = f'<div class="blk-warn" data-doc-rev="{e(str(revision))}">{e(str(exc))} The PDF remains loaded.</div>'
        return (document, gr.update(choices=page_choices(document), value=1), preview_html(document),
                document.pages[0].text, source_html(document, None if uploaded else key),
                recorded, warning or result_html(recorded, revision),
                gr.update(value=_cache_exports(document, recorded, files) if recorded else None, visible=bool(recorded)),
                gr.update(choices=[(f'{s["label"]} / {_range(s["pages"])}', s["pages"][0])
                                   for s in recorded["segments"]] if recorded else [],
                          value=None, visible=bool(recorded)),
                recorded, gr.update(value="", visible=False), revision)

    load_out = [held, page, preview, text, source, result, answer, files, inspect, raw, comparison, rev]

    def preparing(revision, key, uploaded, seen):
        if not ui.newest(seen, "source", revision):
            return [gr.skip() for _ in preparing_out]
        seen["selection"] = (key, str(uploaded) if uploaded else None)
        return [None, '<div class="doc-empty"><p class="blk-run">Preparing the selected PDF</p></div>',
                "", *clear(revision), gr.update(interactive=False)]

    preparing_out = [held, preview, source, *clear_out, run]

    def load_ready(revision, key, uploaded, current_mode, current_rules, current_model, seen):
        selection = (key, str(uploaded) if uploaded else None)
        if ui.as_rev(revision) < ui.applied(seen, "source") or seen.get("selection") != selection:
            return [gr.skip() for _ in [*load_out, run]]
        source_revision = ui.applied(seen, "source")
        loaded = load_source(revision, key, uploaded, current_mode, current_rules, current_model)
        if source_revision != ui.applied(seen, "source") or seen.get("selection") != selection:
            return [gr.skip() for _ in [*load_out, run]]
        return [*loaded, gr.update(interactive=loaded[0] is not None)]

    for event in (upload.upload, upload.clear):
        event(preparing, [rev, preset, upload, source_seen], preparing_out, js=REVISION, queue=False,
              show_progress="hidden", api_visibility="private").then(
            load_ready, [rev, preset, upload, mode, rule, models, source_seen], [*load_out, run],
            js=READ_REVISION, show_progress="minimal", api_visibility="private",
        )

    def pick_sample(revision, key, seen):
        if ui.as_rev(revision) < ui.applied(seen, "source"):
            return [gr.skip() for _ in [upload, mode, *preparing_out]]
        return [None, next(s["mode"] for s in documents.sample_info() if s["id"] == key),
                *preparing(revision, key, None, seen)]

    preset.input(pick_sample, [rev, preset, source_seen], [upload, mode, *preparing_out],
                 js=REVISION, queue=False, api_visibility="private").then(
        load_ready, [rev, preset, upload, mode, rule, models, source_seen], [*load_out, run],
        js=READ_REVISION, show_progress="minimal", api_visibility="private",
    )

    def show_page(document, number):
        if document is None or number is None:
            return preview_html(None), ""
        return preview_html(document, int(number)), document.pages[int(number) - 1].text

    page.change(show_page, [held, page], [preview, text], queue=False,
                show_progress="hidden", api_visibility="private")
    inspect.input(lambda number: gr.update(value=number), [inspect], [page], queue=False,
                  show_progress="hidden", api_visibility="private")

    def run_document(revision, document, category_rules, task, selected, both, uploaded, selected_sample):
        if document is None:
            return (None, '<div class="blk-warn">Choose a readable PDF before running Blink.</div>',
                    gr.update(value=None, visible=False), gr.update(value=None, visible=False),
                    None, gr.update(value="", visible=False))
        try:
            source_path = uploaded or documents.sample_path(selected_sample)
            if not documents.matches_source(document, source_path):
                raise documents.DocumentError("The new PDF is still being prepared. Wait for its page preview before running Blink.")
            outcome = documents.process(document, category_rules, task, selected)
            other = None
            if both and len(blink.models()) == 2:
                other_model = next(m for m in blink.models() if m != selected)
                other = documents.process(document, category_rules, task, other_model)
            paths = _cache_exports(document, outcome, files)
        except (documents.DocumentError, blink.BlinkError) as exc:
            return (None, f'<div class="blk-warn">{e(str(exc))}</div>',
                    gr.update(value=None, visible=False), gr.update(value=None, visible=False),
                    None, gr.update(value="", visible=False))
        except (RuntimeError, gr.Error):
            LOG.exception("Live document processing failed")
            return (None, ('<div class="blk-warn">The live document run could not finish. '
                          'Check your remaining Space quota or try again later. No saved result was substituted.</div>'),
                    gr.update(value=None, visible=False), gr.update(value=None, visible=False),
                    None, gr.update(value="", visible=False))
        return (outcome, result_html(outcome, revision), gr.update(value=paths, visible=True),
                gr.update(choices=[(f'{s["label"]} / {_range(s["pages"])}', s["pages"][0])
                                   for s in outcome["segments"]], value=None, visible=True),
                {"primary": outcome, "comparison": other} if other else outcome,
                gr.update(value=(COMPARE_HEAD + result_html(other, revision)) if other else "",
                          visible=other is not None))

    run.click(run_document, [rev, held, rule, mode, models, compare, upload, preset],
              [result, answer, files, inspect, raw, comparison], js=READ_REVISION,
              show_progress="minimal", concurrency_id="blink-document-read", concurrency_limit=1,
              api_visibility="private")
    return {"model": models, "api": document_api}


def mount_api() -> None:
    with gr.Row(visible=False):
        file = gr.File(file_types=[".pdf"], type="filepath")
        mode = gr.Dropdown(choices=list(documents.MODES), value="classify")
        rule = gr.Textbox(documents.rules_text())
        model = gr.Dropdown(choices=blink.models(), value=documents.default_model())
        output = gr.JSON()
        gr.Button().click(document_api, [file, mode, rule, model], output,
                          api_name="v1_document", concurrency_id="blink-document-read", concurrency_limit=1)