app.py

import io
import os
import tempfile
from collections import Counter

import multiprocessing
import numpy as np
import pandas as pd
import streamlit as st
from PIL import Image

import job_store

# With multiprocessing "spawn", every worker process re-imports __main__ (this
# file).  We must not execute any Streamlit calls inside a worker or they will
# produce ScriptRunContext warnings and can cause the worker to hang.
# All helper functions below are safe to define in both contexts; only the UI
# execution block at the bottom is gated on _IS_MAIN.
_IS_MAIN = multiprocessing.current_process().name == "MainProcess"

if _IS_MAIN:
    # Diagnostic marker: fires on every single script execution, including
    # Streamlit reruns and fresh sessions from new WebSocket connections.
    # Comparing session_id across consecutive prints tells us whether the
    # analysis loop is being killed by an in-session rerun (same session_id)
    # or the browser/proxy dropping and re-opening the WebSocket connection
    # (different session_id each time) — see restart-loop investigation.
    import time as _diag_time
    from streamlit.runtime.scriptrunner import get_script_run_ctx as _get_script_run_ctx
    _diag_ctx = _get_script_run_ctx()
    _diag_session_id = _diag_ctx.session_id if _diag_ctx else "NO_CTX"
    _diag_run_count = st.session_state.get("_diag_run_count", 0) + 1
    st.session_state["_diag_run_count"] = _diag_run_count
    print(
        f"[app][{_diag_time.strftime('%H:%M:%S')}] SCRIPT START #{_diag_run_count} "
        f"pid={os.getpid()} session_id={_diag_session_id}",
        flush=True,
    )


# ── Helpers ───────────────────────────────────────────────────────────────────

def _bgr_pages_from_bytes(pdf_bytes: bytes, dpi: int = 200) -> list[np.ndarray]:
    """Render each PDF page to a BGR numpy array using PyMuPDF (no poppler needed)."""
    import fitz  # PyMuPDF
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    mat = fitz.Matrix(dpi / 72, dpi / 72)
    pages = []
    for page in doc:
        pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB)
        arr = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)
        pages.append(arr[:, :, ::-1].copy())  # RGB → BGR
    return pages


def _pages_to_pdf_bytes(bgr_images: list[np.ndarray]) -> bytes:
    """Convert BGR numpy arrays back to a PDF via Pillow (no cv2 needed)."""
    pils = [Image.fromarray(img[:, :, ::-1]) for img in bgr_images]  # BGR → RGB
    buf = io.BytesIO()
    pils[0].save(buf, format="PDF", save_all=True, append_images=pils[1:])
    return buf.getvalue()


def _pages_matching_text(pdf_bytes: bytes, needle: str, crop_dpi: int = 150) -> list[dict]:
    """Search each page's bottom-right quadrant for `needle` (case-insensitive).
    Architectural title blocks live there, so restricting the search to that
    corner skips boilerplate notes and callouts scattered across the rest of
    the sheet that happen to mention the same words.

    Matching works at word granularity (via PyMuPDF's "words" extraction)
    rather than a plain substring check, since title-block sheet names often
    wrap across two lines (e.g. "1st Story Floor" / "Plan Code Study") --
    concatenating word-by-word with single spaces (instead of the newline
    PyMuPDF would join lines with) lets "floor plan" match across that wrap,
    and keeping each word's own rect lets the match be highlighted precisely
    even when it spans lines.

    Returns one dict per matching page -- {"page": 1-indexed page number,
    "snippet": surrounding text, "crop": RGB uint8 array of the matched
    region with a box drawn around the match} -- so a human can see what was
    actually matched before trusting it."""
    import fitz  # PyMuPDF
    from PIL import ImageDraw

    CONTEXT      = 50  # characters of snippet context on each side of the match
    CROP_PAD_PT  = 40  # points of padding around the match, for the preview crop
    BOX_COLOR    = (220, 30, 30)

    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    needle_norm = " ".join(needle.lower().split())
    scale = crop_dpi / 72
    matches = []
    for i, page in enumerate(doc):
        r = page.rect
        quadrant = fitz.Rect(r.width / 2, r.height / 2, r.width, r.height)
        words = page.get_text("words", clip=quadrant)
        if not words:
            continue

        # Concatenate words (PyMuPDF's own reading-order) with single spaces,
        # tracking each word's character span so a match can be mapped back
        # to the word rect(s) it came from.
        concat = ""
        spans = []  # (start_char, end_char, word_rect)
        for w in words:
            if concat:
                concat += " "
            start = len(concat)
            concat += w[4]
            spans.append((start, len(concat), fitz.Rect(w[0], w[1], w[2], w[3])))

        pos = concat.lower().find(needle_norm)
        if pos == -1:
            continue
        match_end = pos + len(needle_norm)

        involved = [rect for s, e, rect in spans if e > pos and s < match_end]
        if not involved:
            continue
        match_rect = involved[0]
        for rect in involved[1:]:
            match_rect |= rect

        # Grow the box 50% (25% per side) around its own center so
        # descenders/ascenders and characters that PyMuPDF's word rects clip
        # a little tight on are still fully inside the drawn box.
        _grow_x = match_rect.width * 0.25
        _grow_y = match_rect.height * 0.25
        box_rect = fitz.Rect(
            match_rect.x0 - _grow_x, match_rect.y0 - _grow_y,
            match_rect.x1 + _grow_x, match_rect.y1 + _grow_y,
        )

        snip_start = max(0, pos - CONTEXT)
        snip_end   = min(len(concat), match_end + CONTEXT)
        snippet = concat[snip_start:snip_end]
        if snip_start > 0:
            snippet = "…" + snippet
        if snip_end < len(concat):
            snippet = snippet + "…"

        crop_rect = fitz.Rect(
            max(0, box_rect.x0 - CROP_PAD_PT), max(0, box_rect.y0 - CROP_PAD_PT),
            min(r.width, box_rect.x1 + CROP_PAD_PT), min(r.height, box_rect.y1 + CROP_PAD_PT),
        )
        pix = page.get_pixmap(matrix=fitz.Matrix(scale, scale), clip=crop_rect)
        crop_img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
        ImageDraw.Draw(crop_img).rectangle(
            (
                (box_rect.x0 - crop_rect.x0) * scale, (box_rect.y0 - crop_rect.y0) * scale,
                (box_rect.x1 - crop_rect.x0) * scale, (box_rect.y1 - crop_rect.y0) * scale,
            ),
            outline=BOX_COLOR, width=3,
        )

        matches.append({
            "page":    i + 1,
            "snippet": snippet,
            "crop":    np.array(crop_img),
        })
    doc.close()
    return matches


def _extract_pdf_pages(pdf_bytes: bytes, page_indices: list[int]) -> bytes:
    """Return a new PDF (as bytes) containing only the given 0-indexed pages,
    in the order given."""
    import fitz  # PyMuPDF
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    doc.select(page_indices)
    out = doc.tobytes()
    doc.close()
    return out


def _annotate_pages(results: list[dict], valid) -> list[np.ndarray]:
    import classical_cv_detector as ccv
    return [ccv.annotate(r["bgr"], r["dets"], valid=valid) for r in results]


def _find_text_matches(
    pdf_bytes: bytes, needle: str, dpi: int = 200
) -> list[list[tuple[int, int, int, int]]]:
    """
    Search every page's embedded text layer for `needle` (case-insensitive)
    and return, per page, the pixel-space bbox of each occurrence -- scaled
    by dpi/72 to line up with the page images from _bgr_pages_from_bytes,
    which uses the same default dpi. Scanned/rasterized pages with no text
    layer will simply yield zero matches.
    """
    import fitz  # PyMuPDF
    scale = dpi / 72
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    matches = [
        [
            (int(r.x0 * scale), int(r.y0 * scale), int(r.x1 * scale), int(r.y1 * scale))
            for r in page.search_for(needle)
        ]
        for page in doc
    ]
    doc.close()
    return matches


def _mark_text_matches(
    pages: list[np.ndarray], matches_per_page: list[list[tuple[int, int, int, int]]]
) -> list[np.ndarray]:
    """Draw a box around every text match, on a copy of each page."""
    import cv2

    CLR_TEXT_MATCH = (255, 255, 0)  # cyan (BGR) -- distinct from annotate()'s
                                     # green/red/blue/magenta detection colours
    out = []
    for page, boxes in zip(pages, matches_per_page):
        img = page.copy()
        for x1, y1, x2, y2 in boxes:
            cv2.rectangle(img, (x1, y1), (x2, y2), CLR_TEXT_MATCH, 3)
        out.append(img)
    return out


def _build_clean_bytes(results: list[dict], annotated_pages: list[np.ndarray]) -> bytes:
    """
    One PDF page per input page: all detection crops arranged in a left-to-right,
    top-to-bottom grid.  Crops are taken from the annotated source so the
    green/red/blue colour coding is preserved.
    """
    import cv2

    CELL_MARGIN = 12   # px gap around each crop
    LABEL_H     = 22   # px reserved below each crop for the code label
    COLS        = 6    # columns per row

    clean_pages = []
    for r, ann in zip(results, annotated_pages):
        ph, pw = ann.shape[:2]

        crops: list[tuple[np.ndarray, str]] = []
        for d in r["dets"]:
            x, y, bw, bh = d["bbox"]
            pad = 15
            x1 = max(0, x - pad)
            y1 = max(0, y - pad)
            x2 = min(pw, x + bw + pad)
            y2 = min(ph, y + bh + pad)
            if x2 > x1 and y2 > y1:
                crops.append((ann[y1:y2, x1:x2].copy(), d["code"]))

        if not crops:
            clean_pages.append(np.full((200, 800, 3), 255, dtype=np.uint8))
            continue

        max_cw = max(c.shape[1] for c, _ in crops)
        max_ch = max(c.shape[0] for c, _ in crops)

        cols   = min(COLS, len(crops))
        cell_w = max_cw + 2 * CELL_MARGIN
        cell_h = max_ch + LABEL_H + 2 * CELL_MARGIN
        rows   = (len(crops) + cols - 1) // cols

        canvas_w = cols * cell_w + CELL_MARGIN
        canvas_h = rows * cell_h + CELL_MARGIN
        canvas   = np.full((canvas_h, canvas_w, 3), 255, dtype=np.uint8)

        for i, (crop, code) in enumerate(crops):
            row_i, col_i = divmod(i, cols)
            cell_x = CELL_MARGIN + col_i * cell_w
            cell_y = CELL_MARGIN + row_i * cell_h

            ch, cw = crop.shape[:2]
            x_off = (max_cw - cw) // 2
            y_off = (max_ch - ch) // 2
            canvas[cell_y + y_off : cell_y + y_off + ch,
                   cell_x + x_off : cell_x + x_off + cw] = crop

            cv2.rectangle(
                canvas,
                (cell_x + x_off - 1, cell_y + y_off - 1),
                (cell_x + x_off + cw, cell_y + y_off + ch),
                (200, 200, 200), 1,
            )

            (tw, _), _ = cv2.getTextSize(code, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            tx = cell_x + (max_cw - tw) // 2
            ty = cell_y + max_ch + CELL_MARGIN + 14
            cv2.putText(canvas, code, (tx, ty), cv2.FONT_HERSHEY_SIMPLEX,
                        0.5, (60, 60, 60), 1, cv2.LINE_AA)

        clean_pages.append(canvas)

    return _pages_to_pdf_bytes(clean_pages)


def _build_csv_bytes(results: list[dict], legend_df=None) -> bytes:
    rows = []
    for r in results:
        codes = [d["code"] for d in r["dets"] if d["code"] != "?"]
        for code, count in Counter(codes).items():
            rows.append({"Page": r["page"], "Code": code, "Count": count})
    if not rows:
        return b"Code,Total\n"
    df = pd.DataFrame(rows)
    pivot = df.pivot_table(
        index="Code", columns="Page", values="Count",
        fill_value=0, aggfunc="sum",
    )
    pivot.columns = [f"Page {c}" for c in pivot.columns]
    pivot["Total"] = pivot.sum(axis=1)
    pivot = pivot.reset_index()

    if legend_df is not None and "TYPE MARK" in legend_df.columns:
        legend_cols = [c for c in legend_df.columns if c != "TYPE MARK"]
        legend_trim = legend_df[["TYPE MARK"] + legend_cols].rename(
            columns={"TYPE MARK": "Code"}
        )
        pivot = legend_trim.merge(pivot, on="Code", how="right")

    return pivot.to_csv(index=False).encode()


def _summary_pivot(results: list[dict], valid) -> pd.DataFrame | None:
    rows = []
    for r in results:
        codes   = [d["code"] for d in r["dets"] if d["code"] != "?"]
        unknown = sum(1 for d in r["dets"] if d["code"] == "?")
        for code, cnt in Counter(codes).items():
            rows.append({"Page": r["page"], "Code": code, "Count": cnt})
        if unknown:
            rows.append({"Page": r["page"], "Code": "?", "Count": unknown})
    if not rows:
        return None
    df = pd.DataFrame(rows)
    pivot = df.pivot_table(
        index="Code", columns="Page", values="Count",
        fill_value=0, aggfunc="sum",
    )
    pivot.columns = [f"Page {c}" for c in pivot.columns]
    pivot["Total"] = pivot.sum(axis=1)
    if valid is not None:
        pivot.index = pd.Index(
            [
                f"{c}  [valid]"   if c in valid
                else (c           if c == "?"
                else f"{c}  [flagged]")
                for c in pivot.index
            ],
            name="Code",
        )
    total_row = pivot.sum(axis=0).rename("Total").to_frame().T
    total_row.index.name = "Code"
    return pd.concat([pivot, total_row])


def _get_validated_results(results: list[dict]) -> list[dict]:
    """
    Return a copy of results reflecting the current state of the Validate tab
    widgets.  Unchecked detections are dropped; edited codes are substituted.
    Falls back to the original detection when a widget key hasn't been created
    yet (i.e. the Validate tab has never been opened).
    """
    validated = []
    for r in results:
        validated_dets = []
        for j, d in enumerate(r["dets"]):
            key_check = f"det_check_{r['page']}_{j}"
            key_code  = f"det_code_{r['page']}_{j}"
            included  = st.session_state.get(key_check, True)
            code      = st.session_state.get(key_code,  d["code"])
            if included:
                validated_dets.append({**d, "code": code})
        validated.append({**r, "dets": validated_dets})
    return validated


# ── Streamlit UI — only runs in the main process ──────────────────────────────
if _IS_MAIN:
    st.set_page_config(
        page_title="Blueprint Window Shape Detector",
        layout="wide",
        initial_sidebar_state="collapsed",
    )

    # ── Session state ─────────────────────────────────────────────────────────
    _STATE_DEFAULTS = {
        "running":     False,
        "pending":     None,   # {pdf, legend_bytes, legend_name} written before rerun
        "results":     None,
        "valid":       None,
        "legend_df":   None,
        "ann_bytes":   None,
        "csv_bytes":   None,
        "clean_bytes": None,
        "ann_pages":   None,   # list[np.ndarray] — kept for Validate tab crops
        "last_run_summary": None,  # {total_min, pg_done, n, failed} — survives past running=False
        "floor_plan_source_id":     None,  # blueprint_file.file_id the search below was run against
        "floor_plan_matches":       None,  # [{"page": 1-idx, "snippet": str, "crop": ndarray}, ...] from the last search
        "floor_plan_pdf_bytes":     None,  # filtered PDF bytes built from the applied selection, or None
        "floor_plan_applied_pages": None,  # 1-indexed page numbers baked into floor_plan_pdf_bytes
        "mode":              "Run Detection Analysis",  # or "Search Blueprint Text"
        "search_source_id":  None,  # blueprint_file.file_id the search results below were computed from
        "search_results":    {},    # {search text: [count on page 1, count on page 2, ...]}
        "search_marked_pdf": None,  # bytes — all searched terms highlighted, built on demand
    }
    for _key, _val in _STATE_DEFAULTS.items():
        if _key not in st.session_state:
            st.session_state[_key] = _val

    if not st.session_state.running and not st.session_state.pending:
        _orphan_key = job_store.any_active_key()
        if _orphan_key is not None:
            # A job is still running server-side in job_store from a
            # previous session this browser tab has no memory of (a brand
            # new session gets fresh, empty st.session_state, but job_store
            # is a plain module-level dict that survives session loss -- see
            # job_store.py). Resume showing progress for it instead of
            # silently landing on the idle upload screen while the job
            # keeps running unseen.
            st.session_state.running = True
            st.session_state.pending = {"_orphan_job_key": _orphan_key}

    # ── Tabs ──────────────────────────────────────────────────────────────────
    _ready = bool(st.session_state.results)
    tab_input, tab_search, tab_progress, tab_output, tab_validate = st.tabs([
        "Input",
        "Search",
        "Progress" if st.session_state.running else "Progress (idle)",
        "Output"   if _ready else "Output (locked)",
        "Validate" if _ready else "Validate (locked)",
    ])

    # ── PROGRESS TAB ──────────────────────────────────────────────────────────
    # Widgets live inside the tab (not above it) so the main page stays clean
    # while an analysis runs -- previously this whole block rendered above
    # every tab, on screen no matter which tab was open. st.empty()
    # placeholders update in place regardless of which tab is currently
    # selected in the browser, so the polling loop further down still works
    # unchanged; the user just needs this tab open to see the updates.
    _progress_bar          = None
    _status_text           = None
    _page_status_container = None
    with tab_progress:
        if st.session_state.running:
            st.info("Analysis in progress. Please wait — this may take several minutes per page.")
            _progress_bar          = st.progress(0)
            _status_text           = st.empty()
            _page_status_container = st.empty()
        elif st.session_state.last_run_summary:
            # Show the completed/failed run's final numbers instead of just
            # going back to a bare idle message -- previously this
            # information (total runtime, per-page breakdown) was computed
            # and then immediately discarded on the st.rerun() that follows
            # completion, since running flips to False right before it.
            _summary = st.session_state.last_run_summary
            if _summary.get("failed"):
                st.error(f"Last analysis failed after {_summary['total_min']} minutes.")
            else:
                st.success(f"Last analysis completed in {_summary['total_min']} minutes.")
            _lines = ["**Page Breakdown**"]
            for _pi in range(_summary["n"]):
                if _pi in _summary["pg_done"]:
                    _lines.append(
                        f"&nbsp;&nbsp;&nbsp;&nbsp;Page {_pi + 1} finished processing "
                        f"in {_summary['pg_done'][_pi]} minutes"
                    )
            st.markdown("  \n".join(_lines))
            st.caption("Start a new analysis from the **Input** tab.")
        else:
            st.caption("No analysis is currently running. Start one from the **Input** tab.")

    # ══ INPUT TAB ════════════════════════════════════════════════════════════
    with tab_input:
        st.title("Blueprint Diamond Detector")
        st.caption(
            "Detects diamond-shaped unit symbols (e.g. A1, B2) in construction blueprints "
            "using classical computer vision and OCR."
        )
        st.divider()

        blueprint_file = st.file_uploader(
            "Blueprint PDF  *(required)*",
            type=["pdf"],
            help="The construction blueprint PDF containing diamond unit symbols.",
            disabled=bool(st.session_state.running),
        )

        if blueprint_file:
            import fitz as _fitz
            _doc = _fitz.open(stream=blueprint_file.read(), filetype="pdf")
            _n_pages = len(_doc)
            _doc.close()
            blueprint_file.seek(0)

            _parallel_batches = (_n_pages + 1) // 2
            _est_lo = _parallel_batches * 5
            _est_hi = _parallel_batches * 8
            _time_note = (
                "approximately 5–8 minutes" if _n_pages == 1
                else f"approximately {_est_lo}–{_est_hi} minutes"
            )
            st.caption(
                f"**{_n_pages} page(s) detected.** "
                f"Pages are analysed 2 at a time in parallel (using both vCPUs), "
                f"so the estimated runtime is {_time_note}."
            )

            st.radio(
                "What would you like to do?",
                ["Run Detection Analysis", "Search Blueprint Text"],
                key="mode",
                horizontal=True,
                disabled=bool(st.session_state.running),
            )
            if st.session_state.mode == "Search Blueprint Text":
                st.info(
                    "Detection analysis is disabled while Search is selected. "
                    "Switch to the **Search** tab above to search this blueprint's text."
                )
            else:
                if st.session_state.floor_plan_source_id != blueprint_file.file_id:
                    # A different file than the one the current search (if any)
                    # was run against -- drop the stale results rather than
                    # silently applying them to unrelated pages.
                    st.session_state.floor_plan_source_id     = None
                    st.session_state.floor_plan_matches        = None
                    st.session_state.floor_plan_pdf_bytes       = None
                    st.session_state.floor_plan_applied_pages   = None

                if st.button(
                    "Search for 'Floor Plan' Pages",
                    disabled=bool(st.session_state.running),
                ):
                    st.session_state.floor_plan_source_id     = blueprint_file.file_id
                    st.session_state.floor_plan_matches        = _pages_matching_text(
                        blueprint_file.getvalue(), "Floor Plan"
                    )
                    st.session_state.floor_plan_pdf_bytes       = None
                    st.session_state.floor_plan_applied_pages   = None
                    for _k in list(st.session_state.keys()):
                        if _k.startswith("floor_plan_check_"):
                            del st.session_state[_k]
                    st.rerun()

                if st.session_state.floor_plan_source_id == blueprint_file.file_id:
                    _fp_matches = st.session_state.floor_plan_matches
                    if not _fp_matches:
                        st.warning("No pages containing the text 'Floor Plan' were found.")
                    else:
                        _fp_included = sum(
                            st.session_state.get(f"floor_plan_check_{m['page']}", True)
                            for m in _fp_matches
                        )
                        with st.expander(
                            f"Floor Plan Matches — {len(_fp_matches)} of {_n_pages} page(s), "
                            f"{_fp_included} included",
                            expanded=False,
                        ):
                            st.caption(
                                "Each page below matched 'Floor Plan' in its title-block "
                                "corner. Uncheck any that aren't real floor plan pages, "
                                "then press **Apply Selection** to commit."
                            )
                            with st.form(key="floor_plan_form", enter_to_submit=False):
                                for _m in _fp_matches:
                                    _key_check = f"floor_plan_check_{_m['page']}"
                                    _fp_c1, _fp_c2 = st.columns([1, 4])
                                    with _fp_c1:
                                        st.checkbox(
                                            f"Page {_m['page']}",
                                            value=st.session_state.get(_key_check, True),
                                            key=_key_check,
                                        )
                                    with _fp_c2:
                                        st.image(_m["crop"], caption=_m["snippet"], width=280)
                                _fp_submitted = st.form_submit_button(
                                    "Apply Selection", use_container_width=True
                                )

                            if _fp_submitted:
                                _fp_selected = [
                                    m["page"] for m in _fp_matches
                                    if st.session_state.get(f"floor_plan_check_{m['page']}", True)
                                ]
                                if _fp_selected:
                                    st.session_state.floor_plan_pdf_bytes = _extract_pdf_pages(
                                        blueprint_file.getvalue(),
                                        [p - 1 for p in _fp_selected],
                                    )
                                    st.session_state.floor_plan_applied_pages = _fp_selected
                                else:
                                    st.session_state.floor_plan_pdf_bytes     = None
                                    st.session_state.floor_plan_applied_pages = None
                                    st.warning("No pages selected — all pages will be used instead.")
                                st.rerun()

                        if st.session_state.floor_plan_applied_pages:
                            _fp_applied = st.session_state.floor_plan_applied_pages
                            st.success(
                                f"Using {len(_fp_applied)} of {_n_pages} page(s): "
                                f"{', '.join(str(n) for n in _fp_applied)}."
                            )
                            if st.button("Use All Pages Instead", disabled=bool(st.session_state.running)):
                                st.session_state.floor_plan_pdf_bytes     = None
                                st.session_state.floor_plan_applied_pages = None
                                st.rerun()

        legend_file = st.file_uploader(
            "Legend Image  *(optional)*",
            type=["png", "jpg", "jpeg", "bmp", "tiff", "webp"],
            help="A screenshot or scan of the window schedule / legend. "
                 "Used to validate detected codes as confirmed or flagged.",
            disabled=bool(st.session_state.running),
        )

        exemplar_file = st.file_uploader(
            "Example Target Shape  *(optional)*",
            type=["png", "jpg", "jpeg", "bmp", "tiff", "webp"],
            help="A tight crop of a single instance of the target symbol, including "
                 "its interior code (e.g. A1). When provided, detection uses this "
                 "shape instead of the default diamond detector.",
            disabled=bool(st.session_state.running),
        )

        if exemplar_file:
            import cv2 as _cv2
            import classical_cv_detector as _ccv

            _exemplar_bytes = exemplar_file.read()
            exemplar_file.seek(0)
            _exemplar_bgr = _cv2.imdecode(
                np.frombuffer(_exemplar_bytes, np.uint8), _cv2.IMREAD_COLOR
            )
            _exemplar_data = _ccv.extract_shape_template(_exemplar_bgr) if _exemplar_bgr is not None else None

            if _exemplar_data is None:
                st.warning(
                    "No clear shape outline could be found in this image. "
                    "Try a tighter crop with the symbol's outline clearly visible."
                )
            else:
                _prev_col1, _prev_col2 = st.columns(2)
                with _prev_col1:
                    st.caption(f"Detected shape type: **{_exemplar_data['shape_type']}**")
                    st.image(_exemplar_bgr[:, :, ::-1], caption="Example", width=120)
                with _prev_col2:
                    st.caption("Extracted mask")
                    st.image(_exemplar_data["mask"], caption="Mask", width=120)

        if st.button(
            "Run Analysis",
            type="primary",
            disabled=(
                bool(st.session_state.running)
                or blueprint_file is None
                or st.session_state.mode == "Search Blueprint Text"
            ),
        ):
            _use_floor_plan_filter = (
                st.session_state.floor_plan_source_id == blueprint_file.file_id
                and st.session_state.floor_plan_pdf_bytes is not None
            )
            st.session_state.pending = {
                "pdf": (
                    st.session_state.floor_plan_pdf_bytes if _use_floor_plan_filter
                    else blueprint_file.read()
                ),
                "legend_bytes":   legend_file.read()   if legend_file   else None,
                "legend_name":    legend_file.name     if legend_file   else None,
                "exemplar_bytes": exemplar_file.read() if exemplar_file else None,
            }
            # Clear previous results and all validation widget state
            for _k in ("results", "valid", "legend_df", "ann_bytes", "csv_bytes",
                       "clean_bytes", "ann_pages"):
                st.session_state[_k] = None
            for _k in list(st.session_state.keys()):
                if _k.startswith("det_check_") or _k.startswith("det_code_"):
                    del st.session_state[_k]
            st.session_state.running = True
            st.rerun()

    # ══ SEARCH TAB ═══════════════════════════════════════════════════════════
    with tab_search:
        st.subheader("Search Blueprint Text")
        if blueprint_file is None:
            st.warning("Upload a blueprint PDF on the **Input** tab first.")
        else:
            if st.session_state.search_source_id != blueprint_file.file_id:
                # A different file than the one the search results below were
                # computed from -- drop stale results rather than showing
                # counts for pages that no longer correspond to this file.
                st.session_state.search_source_id  = blueprint_file.file_id
                st.session_state.search_results     = {}
                st.session_state.search_marked_pdf  = None

            st.caption(
                "Search the blueprint's embedded text layer for a phrase and see how "
                "many times it appears on each page. Only text in the PDF's text layer "
                "is searched -- scanned/rasterized pages won't have matches. Enter one "
                "phrase at a time; each search is added to the table below."
            )

            with st.form(key="search_form", clear_on_submit=True):
                _search_col, _button_col = st.columns([4, 1])
                with _search_col:
                    _search_needle = st.text_input(
                        "Text to search for",
                        placeholder="e.g. Floor Plan",
                        label_visibility="collapsed",
                    )
                with _button_col:
                    _search_submitted = st.form_submit_button(
                        "Search", use_container_width=True
                    )

            if _search_submitted and _search_needle.strip():
                _needle = _search_needle.strip()
                _matches = _find_text_matches(blueprint_file.getvalue(), _needle)
                st.session_state.search_results[_needle] = [len(m) for m in _matches]
                st.session_state.search_marked_pdf = None  # stale — rebuild on demand

            if st.session_state.search_results:
                import fitz as _fitz_search
                _search_doc = _fitz_search.open(
                    stream=blueprint_file.getvalue(), filetype="pdf"
                )
                _search_n_pages = len(_search_doc)
                _search_doc.close()

                _rows = []
                for _term, _counts in st.session_state.search_results.items():
                    _row = {"Search Text": _term}
                    for _pi in range(_search_n_pages):
                        _row[f"Page {_pi + 1}"] = _counts[_pi]
                    _row["Total"] = sum(_counts)
                    _rows.append(_row)
                _search_df = pd.DataFrame(_rows).set_index("Search Text")
                st.dataframe(_search_df, use_container_width=True)

                if st.button("Generate Highlighted PDF (all searched terms)"):
                    _marked_pages = _bgr_pages_from_bytes(blueprint_file.getvalue())
                    for _term in st.session_state.search_results:
                        _term_matches = _find_text_matches(blueprint_file.getvalue(), _term)
                        _marked_pages = _mark_text_matches(_marked_pages, _term_matches)
                    st.session_state.search_marked_pdf = _pages_to_pdf_bytes(_marked_pages)

                if st.session_state.search_marked_pdf:
                    st.download_button(
                        "Download Highlighted PDF",
                        data=st.session_state.search_marked_pdf,
                        file_name="search_highlighted.pdf",
                        mime="application/pdf",
                    )
            else:
                st.caption("No searches yet — enter text above and press **Search**.")

    # ══ OUTPUT TAB ═══════════════════════════════════════════════════════════
    with tab_output:
        if not st.session_state.results:
            st.warning("Results are not yet available. Upload a blueprint and run the analysis on the **Input** tab first.")
        else:
            results = st.session_state.results
            valid   = st.session_state.valid

            # Use validated results (reflects Validate tab edits) for summary + CSV
            _v_results = _get_validated_results(results)

            st.subheader("Detection Summary")
            st.caption("Updates live as you make changes in the Validate tab.")
            pivot = _summary_pivot(_v_results, valid)
            if pivot is not None:
                st.dataframe(pivot, use_container_width=True)
                if valid is not None:
                    st.caption(
                        "[valid] = code found in legend   "
                        "[flagged] = code not in legend   "
                        "? = shape detected but OCR failed"
                    )
            else:
                st.write("No detections found.")

            st.divider()
            st.subheader("Downloads")

            col1, col2, col3 = st.columns(3)

            with col1:
                st.download_button(
                    label="Download Annotated PDF",
                    data=st.session_state.ann_bytes,
                    file_name="annotated_blueprint.pdf",
                    mime="application/pdf",
                    use_container_width=True,
                )
                st.caption(
                    "Original pages with bounding boxes. "
                    "Green = confirmed, Red = flagged, Blue = shape-only. "
                    "Use Regenerate PDFs in the Validate tab to reflect edits."
                )

            with col2:
                st.download_button(
                    label="Download CSV Report",
                    data=_build_csv_bytes(_v_results, st.session_state.legend_df),
                    file_name="detection_report.csv",
                    mime="text/csv",
                    use_container_width=True,
                )
                st.caption(
                    "Pivot table of code counts by page. Reflects Validate tab edits live."
                )

            with col3:
                st.download_button(
                    label="Download Detections-Only PDF",
                    data=st.session_state.clean_bytes,
                    file_name="detections_only.pdf",
                    mime="application/pdf",
                    use_container_width=True,
                )
                st.caption(
                    "Grid of detection crops per page. "
                    "Use Regenerate PDFs in the Validate tab to reflect edits."
                )

    # ══ VALIDATE TAB ═════════════════════════════════════════════════════════
    with tab_validate:
        if not st.session_state.results:
            st.warning("Results are not yet available. Upload a blueprint and run the analysis on the **Input** tab first.")
        else:
            results   = st.session_state.results
            valid     = st.session_state.valid
            ann_pages = st.session_state.ann_pages  # annotated BGR arrays

            st.subheader("Validate Detections")
            st.caption(
                "Each detection is shown with its annotated crop. "
                "Uncheck false positives or edit codes, then press **Apply Changes** "
                "for that page to commit. The summary table and CSV update immediately. "
                "Use **Regenerate PDFs** at the bottom to rebuild the PDF exports."
            )

            _VALIDATE_COLS = 6
            # Crop pad: enough to show the text label drawn above each bbox.
            # annotate() draws text at y-6 with ~15px height, so we need >21px
            # above the box top.  25px gives comfortable margin on all sides.
            _CROP_PAD = 25

            for _ri, _r in enumerate(results):
                _page_num  = _r["page"]
                _dets      = _r["dets"]
                _ann_page  = ann_pages[_ri] if ann_pages else None
                _ph, _pw   = (_ann_page.shape[:2] if _ann_page is not None
                              else _r["bgr"].shape[:2])
                _src_page  = _ann_page if _ann_page is not None else _r["bgr"]

                _n_included = sum(
                    st.session_state.get(f"det_check_{_page_num}_{j}", True)
                    for j in range(len(_dets))
                )

                with st.expander(
                    f"Page {_page_num} — {len(_dets)} detection(s), "
                    f"{_n_included} included",
                    expanded=False,
                ):
                    if not _dets:
                        st.write("No detections on this page.")
                        continue

                    # Sort: green (confirmed) first, red (flagged) second,
                    # blue (shape-only / "?") last.  Original index is preserved
                    # alongside each detection so widget keys remain stable.
                    def _sort_priority(code):
                        if code == "?":
                            return 2
                        if valid is None or code in valid:
                            return 0
                        return 1

                    _dets_sorted = sorted(
                        enumerate(_dets),
                        key=lambda x: _sort_priority(x[1]["code"]),
                    )

                    # enter_to_submit=False prevents pressing Enter in a code
                    # field from triggering form submission (Streamlit >= 1.39).
                    with st.form(key=f"validate_form_{_page_num}", enter_to_submit=False):
                        # Fixed-height scrollable container (Streamlit >= 1.30) so
                        # opening an expander doesn't push the rest of the page down.
                        # The caption and submit button sit below it, always visible.
                        with st.container(height=1000):
                            for _row_start in range(0, len(_dets_sorted), _VALIDATE_COLS):
                                _row_dets = _dets_sorted[_row_start : _row_start + _VALIDATE_COLS]
                                _cols = st.columns(_VALIDATE_COLS)  # always full width
                                for _ci, (_col, (_det_idx, _d)) in enumerate(zip(_cols, _row_dets)):
                                    _x, _y, _bw, _bh = _d["bbox"]
                                    _x1 = max(0, _x - _CROP_PAD)
                                    _y1 = max(0, _y - _CROP_PAD)
                                    _x2 = min(_pw, _x + _bw + _CROP_PAD)
                                    _y2 = min(_ph, _y + _bh + _CROP_PAD)
                                    _crop_rgb = _src_page[_y1:_y2, _x1:_x2, ::-1]

                                    _key_check = f"det_check_{_page_num}_{_det_idx}"
                                    _key_code  = f"det_code_{_page_num}_{_det_idx}"

                                    with _col:
                                        st.image(_crop_rgb, width="stretch")
                                        st.checkbox(
                                            "Include",
                                            value=st.session_state.get(_key_check, True),
                                            key=_key_check,
                                        )
                                        st.text_input(
                                            "Code",
                                            value=st.session_state.get(_key_code, _d["code"]),
                                            key=_key_code,
                                            label_visibility="collapsed",
                                            placeholder="e.g. A1",
                                        )

                        st.caption(
                            "Uncheck any false positives and correct any misread codes above, "
                            "then press this button to save your changes. The summary table "
                            "and CSV in the Output tab will update immediately."
                        )
                        st.form_submit_button(
                            f"Apply Changes — Page {_page_num}",
                            use_container_width=True,
                        )

            st.divider()
            st.caption(
                "Once you have applied changes across all pages, press this button to rebuild "
                "the annotated and detections-only PDFs so the downloads in the Output tab "
                "reflect your edits."
            )
            if st.button("Regenerate PDFs", type="primary"):
                _v = _get_validated_results(st.session_state.results)
                _new_ann = _annotate_pages(_v, valid)
                st.session_state.ann_pages   = _new_ann
                st.session_state.ann_bytes   = _pages_to_pdf_bytes(_new_ann)
                st.session_state.clean_bytes = _build_clean_bytes(_v, _new_ann)
                st.success("PDFs regenerated — download updated files from the Output tab.")

    # ══ ANALYSIS (runs after the UI re-renders with all widgets disabled) ═════
    if st.session_state.running and st.session_state.pending:
        import time as _time

        def _ts() -> str:
            # Wall-clock stamp on our own log lines -- Streamlit/Tornado's own
            # log lines already carry a timestamp, so this lets the two be
            # correlated directly when diagnosing a restart's exact timing.
            return _time.strftime("%H:%M:%S")

        _pending = st.session_state.pending
        _is_orphan_attach = "_orphan_job_key" in _pending
        if _is_orphan_attach:
            _job_key = _pending["_orphan_job_key"]
        else:
            _job_key = job_store.job_key(
                _pending["pdf"], _pending["legend_bytes"], _pending.get("exemplar_bytes")
            )

        if not _is_orphan_attach and job_store.claim(_job_key):
            # First time seeing this exact job (same PDF+legend+exemplar
            # bytes): do the one-time setup (legend parse, PDF render, spawn
            # the Pool) and stash everything in job_store, keyed by a hash
            # of the input bytes rather than by Streamlit session -- see
            # job_store.py for why session-scoped storage isn't enough
            # (a session can be destroyed outright, not just reconnected).
            # claim() (not a plain get()-is-None check) guards this so two
            # reruns racing on the same job can't both create a Pool.
            #
            # Deliberately NOT using `with pool:` -- that would close the
            # pool the moment THIS script execution ends, but if a rerun or
            # a brand new session picks this job back up mid-poll, we want
            # it to find the pool still alive and keep polling it rather
            # than closing it out from under the still-running workers.
            _valid = None
            if _pending["legend_bytes"]:
                # Dispatch by the exemplar's classified shape_type, mirroring
                # classical_cv_detector.run()'s own dispatch (see its comment
                # at the equivalent point): hexagon schedules use a 'WINDOW
                # LETTER' column (single letter, or letter+0-2 digits -- e.g.
                # Colorado Grand Oaks' W1..W11), diamond/default schedules use
                # 'TYPE MARK' (letter+digit, e.g. A1). This dispatch was
                # already fixed in the CLI script but never ported here --
                # app.py always called the diamond-only parser regardless of
                # exemplar shape, which raises on a hexagon-format legend
                # (no 'TYPE MARK' column) and was being silently swallowed
                # by the except below, leaving _valid=None. With no legend
                # to check against, annotate() then marks every OCR'd code
                # "confirmed" (green) -- including random page text that
                # happens to OCR into something letter+digit-shaped -- with
                # no visible error at all.
                import cv2 as _cv2_legend
                import classical_cv_detector as _ccv_legend
                _is_hexagon_legend = False
                _exemplar_bytes_for_legend = _pending.get("exemplar_bytes")
                if _exemplar_bytes_for_legend:
                    _ex_bgr = _cv2_legend.imdecode(
                        np.frombuffer(_exemplar_bytes_for_legend, np.uint8), _cv2_legend.IMREAD_COLOR
                    )
                    _ex_data = _ccv_legend.extract_shape_template(_ex_bgr) if _ex_bgr is not None else None
                    _is_hexagon_legend = _ex_data is not None and _ex_data["shape_type"] == "hexagon"

                _suffix = "." + _pending["legend_name"].rsplit(".", 1)[-1]
                with tempfile.NamedTemporaryFile(suffix=_suffix, delete=False) as _tmp:
                    _tmp.write(_pending["legend_bytes"])
                    _tmp_path = _tmp.name
                try:
                    if _is_hexagon_legend:
                        from legend_parser import parse_hexagon_legend_image, valid_codes_hexagon
                        _legend_df = parse_hexagon_legend_image(_tmp_path)
                        _valid = valid_codes_hexagon(_legend_df)
                    else:
                        from legend_parser import parse_legend_image, valid_codes
                        _legend_df = parse_legend_image(_tmp_path)
                        _valid = valid_codes(_legend_df)
                    st.session_state.legend_df = _legend_df
                except Exception:
                    _valid = None
                finally:
                    try:
                        os.unlink(_tmp_path)
                    except OSError:
                        pass

            _status_text.text("Rendering PDF pages…")
            _pages = _bgr_pages_from_bytes(_pending["pdf"])
            _n = len(_pages)

            import classical_cv_detector as ccv

            _ctx = multiprocessing.get_context("spawn")

            # Cross-process progress channel: stdout print()s from inside a
            # worker are only visible via HF's log-stream endpoint, which has
            # proven unreliable for this (a live capture across two full
            # 8-page runs captured zero of the tile-level progress lines
            # added for the pages-6/7 stall, despite ~90+ expected). A
            # multiprocessing.Manager().dict() proxy was tried next and ALSO
            # showed nothing -- and because that failure mode was silently
            # swallowed too, there was no way to tell whether writes were
            # failing or simply never happening. Plain files on the
            # container's local disk are about as hard to silently break as
            # cross-process communication gets: the main process and every
            # worker subprocess share the same filesystem unconditionally,
            # no proxy/socket lifecycle involved. One file per page index.
            _progress_dir = f"/tmp/tile_progress_{_job_key}"
            os.makedirs(_progress_dir, exist_ok=True)

            _exemplar_bytes = _pending.get("exemplar_bytes")
            _worker_args = [
                (p.tobytes(), p.shape, p.dtype.str, ccv.TILE_OVERLAP, ccv.OCR_UPSCALE,
                 _exemplar_bytes, _progress_dir, _pi)
                for _pi, p in enumerate(_pages)
            ]

            _pg_start: dict[int, float] = {}   # page index -> start timestamp
            _pg_done:  dict[int, int]   = {}   # page index -> elapsed minutes once finished
            _t0 = _time.time()
            for _pi in range(min(2, _n)):
                _pg_start[_pi] = _t0

            print(f"[app][{_ts()}] creating Pool(processes=2) for {_n} page(s)", flush=True)
            _t_pool = _time.time()
            _pool = _ctx.Pool(processes=2, initializer=ccv.lower_worker_priority)
            print(f"[app][{_ts()}] pool created in {_time.time()-_t_pool:.1f}s, submitting tasks", flush=True)
            _async_results = [
                _pool.apply_async(ccv.detect_page_worker, (arg,))
                for arg in _worker_args
            ]

            job_store.put(_job_key, {
                "pool":           _pool,
                "progress_dir":   _progress_dir,
                "async_results":  _async_results,
                "pages":          _pages,
                "n":              _n,
                "all_dets":       [None] * _n,
                "pg_done":        _pg_done,
                "pg_start":       _pg_start,
                "t_pool":         _t_pool,
                "valid":          _valid,
                "legend_df":      st.session_state.legend_df,
                "done_count":     0,
            })
        else:
            # Either the job already fully exists, another thread just
            # claimed it and is still building it (spawning a Pool + parsing
            # the legend can take a few seconds), or this is a fresh session
            # auto-adopting an orphaned job -- wait for it to appear rather
            # than racing ahead with job_store.get() returning None.
            _wait_start = _time.time()
            while job_store.get(_job_key) is None and _time.time() - _wait_start < 30:
                _time.sleep(0.2)
            print(f"[app][{_ts()}] rerun/reconnect detected -- reattaching to in-flight "
                  f"analysis (job={_job_key[:8]})", flush=True)

        _state = job_store.get(_job_key)
        if _state is None:
            # The job vanished (finished and was popped by another session)
            # in the narrow window between us discovering it and getting
            # here -- nothing left to attach to, so drop back to idle rather
            # than crash on a missing state dict.
            st.session_state.running = False
            st.session_state.pending = None
            st.rerun()

        st.session_state.legend_df = _state["legend_df"]
        _pool           = _state["pool"]
        _async_results  = _state["async_results"]
        _pages          = _state["pages"]
        _n              = _state["n"]
        _all_dets       = _state["all_dets"]
        _pg_done        = _state["pg_done"]
        _pg_start       = _state["pg_start"]
        _t_pool         = _state["t_pool"]
        _valid          = _state["valid"]
        # .get() with a fallback: a job already in flight when this deploy
        # landed was created by prior code and has no "progress_dir" key --
        # reattaching to it must not KeyError.
        _progress_dir   = _state.get("progress_dir")

        def _render_page_status():
            # Computed fresh on every call (rather than reading a static
            # cached string) so an in-progress page's elapsed time ticks up
            # in real time as this function is re-invoked on every poll
            # cycle below, not just when a page transitions to done.
            lines = ["**Processing Pages**"]
            _now = _time.time()
            for _pi in range(_n):
                if _pi in _pg_done:
                    lines.append(
                        f"&nbsp;&nbsp;&nbsp;&nbsp;Page {_pi + 1} finished processing "
                        f"in {_pg_done[_pi]} minutes"
                    )
                elif _pi in _pg_start:
                    _live_min = round((_now - _pg_start[_pi]) / 60)
                    lines.append(
                        f"&nbsp;&nbsp;&nbsp;&nbsp;Page {_pi + 1} has been processing "
                        f"for {_live_min} minutes"
                    )
                    _detail = None
                    if _progress_dir:
                        try:
                            with open(os.path.join(_progress_dir, f"page_{_pi}.txt")) as _f:
                                _detail = _f.read().strip()
                        except Exception:
                            # File may not exist yet (worker hasn't written
                            # its first checkpoint) -- never let a
                            # diagnostic read break the actual progress UI.
                            _detail = None
                    if _detail:
                        lines.append(f"&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;→ {_detail}")
            _page_status_container.markdown("  \n".join(lines))

        _progress_bar.progress(_state["done_count"] / _n)
        _render_page_status()

        # Poll for results instead of blocking on pool.imap(): a fully blocking
        # wait sends zero WebSocket traffic to the browser for however long a
        # page takes (several minutes), and the HF Spaces reverse proxy treats
        # that as a dead connection and reconnects, which reruns this script
        # from the top (occasionally landing on a brand new session if the
        # old one was destroyed) -- which is why the pool/async-results are
        # persisted in job_store above rather than kept as local variables
        # that a rerun would silently discard.
        _POLL_INTERVAL      = 2.0   # seconds between readiness checks
        _HEARTBEAT_INTERVAL = 15.0  # seconds between keepalive UI pushes

        _done_count     = _state["done_count"]
        _last_heartbeat = _time.time()
        try:
            while _done_count < _n:
                for _i, _ar in enumerate(_async_results):
                    if _all_dets[_i] is not None or not _ar.ready():
                        continue
                    _dets = _ar.get()
                    _all_dets[_i] = _dets
                    _done_count += 1
                    _state["done_count"] = _done_count

                    _elapsed_min = round((_time.time() - _pg_start.get(_i, _t_pool)) / 60)
                    print(f"[app][{_ts()}] received page {_i + 1}/{_n} result in {_elapsed_min}min", flush=True)

                    _progress_bar.progress(_done_count / _n)
                    _pg_done[_i] = _elapsed_min

                    _next = _i + 2
                    if _next < _n and _next not in _pg_start:
                        _pg_start[_next] = _time.time()

                    _render_page_status()
                    _last_heartbeat = _time.time()

                if _done_count >= _n:
                    break

                # Re-render on every poll cycle (not just on page-completion
                # transitions or the coarser heartbeat below) so in-progress
                # pages' elapsed-minutes text updates in real time.
                _render_page_status()

                _now = _time.time()
                if _now - _last_heartbeat >= _HEARTBEAT_INTERVAL:
                    _status_text.text(
                        f"Still processing… ({_done_count}/{_n} pages complete, "
                        f"{round((_now - _t_pool) / 60)} min elapsed)"
                    )
                    _last_heartbeat = _now

                _time.sleep(_POLL_INTERVAL)
        except Exception:
            # Don't leak the worker Pool on an unexpected error -- an orphaned
            # Pool would keep burning CPU indefinitely, which is exactly the
            # kind of contention this whole persistence scheme is trying to
            # avoid. Also clear the stored state so a retry starts clean
            # instead of reattaching to a dead job.
            print(f"[app][{_ts()}] analysis failed, terminating pool", flush=True)
            _pool.terminate()
            _pool.join()
            if _progress_dir:
                import shutil as _shutil
                _shutil.rmtree(_progress_dir, ignore_errors=True)
            job_store.pop(_job_key)
            st.session_state.last_run_summary = {
                "total_min": round((_time.time() - _t_pool) / 60),
                "pg_done":   dict(_pg_done),
                "n":         _n,
                "failed":    True,
            }
            st.session_state.running = False
            st.session_state.pending = None
            raise

        _pool.close()
        _pool.join()
        if _progress_dir:
            import shutil as _shutil
            _shutil.rmtree(_progress_dir, ignore_errors=True)
        job_store.pop(_job_key)

        print(f"[app][{_ts()}] all pages done", flush=True)
        _status_text.empty()

        _results = [
            {"page": _i + 1, "bgr": _pages[_i], "dets": _all_dets[_i]}
            for _i in range(_n)
        ]

        # ── Build output files ────────────────────────────────────────────────
        _status_text.text("Building output files…")
        _ann_pages = _annotate_pages(_results, _valid)
        st.session_state.ann_pages   = _ann_pages
        st.session_state.ann_bytes   = _pages_to_pdf_bytes(_ann_pages)
        st.session_state.csv_bytes   = _build_csv_bytes(_results, st.session_state.legend_df)
        st.session_state.clean_bytes = _build_clean_bytes(_results, _ann_pages)
        st.session_state.results     = _results
        st.session_state.valid       = _valid
        st.session_state.last_run_summary = {
            "total_min": round((_time.time() - _t_pool) / 60),
            "pg_done":   dict(_pg_done),
            "n":         _n,
            "failed":    False,
        }
        st.session_state.pending     = None
        st.session_state.running     = False

        st.rerun()