app.py

"""Streamlit UI — layout, rendering, and user interaction only."""

import base64
import hashlib
import logging
import logging.handlers
import time
import json
import os
import tempfile
from pathlib import Path

import pandas as pd
import streamlit as st

# ── File logging (mirrors Django's LOGGING config so both write the same file) ─
_LOG_DIR = Path(__file__).parent / 'logs'
_LOG_DIR.mkdir(exist_ok=True)
_root = logging.getLogger()
if not any(isinstance(h, logging.handlers.RotatingFileHandler) for h in _root.handlers):
    _fh = logging.handlers.RotatingFileHandler(
        _LOG_DIR / 'docprocessor.log',
        maxBytes=10 * 1024 * 1024,
        backupCount=5,
        encoding='utf-8',
    )
    _fh.setFormatter(logging.Formatter(
        '%(asctime)s %(levelname)-8s %(name)s  %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S',
    ))
    _fh.setLevel(logging.DEBUG)
    _root.addHandler(_fh)
    _root.setLevel(logging.DEBUG)
for _mod in ('llm', 'extractor'):
    logging.getLogger(_mod).setLevel(logging.DEBUG)

logger = logging.getLogger(__name__)

from extractor import (
    IMAGE_EXTS, DEVICE_AUTO, DEVICE_CPU, DEVICE_GPU,
    count_pdf_pages, detect_cuda, get_page_text,
    process_file, process_file_high,
)
from llm import (
    PRIMARY_MODEL, FALLBACK_MODEL, QA_SUGGESTIONS, set_num_gpu,
    extract_fields, extract_header, extract_totals, extract_line_items_page,
    extract_line_items_fallback, stream_qa,
)

MODES = ["Standard", "High Accuracy"]

# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="DocProcessor",
    page_icon="",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Load Font Awesome + shared styles
st.markdown("""
<link rel="stylesheet"
      href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.1/css/all.min.css"
      crossorigin="anonymous">
<style>
[data-testid="metric-container"] {
    background: rgba(255,255,255,0.04);
    border: 1px solid rgba(255,255,255,0.09);
    border-radius: 10px;
    padding: 10px 14px;
}
.info-card {
    background: rgba(255,255,255,0.04);
    border-left: 3px solid #4c9be8;
    border-radius: 6px;
    padding: 10px 14px;
    margin-bottom: 8px;
    font-size: 0.87rem;
    line-height: 1.7;
}
.info-card .card-title { font-size: 0.93rem; font-weight: 700; }
.info-card .fa-solid { opacity: 0.55; width: 16px; }
.grand-total {
    font-size: 1.55rem;
    font-weight: 800;
    color: #4ade80;
    margin-top: 4px;
}
.section-label {
    font-size: 0.72rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.07em;
    opacity: 0.50;
    margin-bottom: 6px;
}
.st-emotion-cache-vl2mil{font-size:15px;}
.section-label .fa-solid { margin-right: 5px; }
.reconciliation-row td { color: #f87171 !important; font-style: italic; }
div[data-testid="stDataFrame"] { border-radius: 8px; overflow: hidden; }
.high-badge {
    display: inline-block;
    background: #1d4ed8;
    color: #fff;
    font-size: 0.68rem;
    font-weight: 700;
    padding: 2px 7px;
    border-radius: 99px;
    letter-spacing: 0.04em;
    vertical-align: middle;
    margin-left: 6px;
}
.job-row {
    display: flex;
    align-items: flex-start;
    gap: 7px;
    padding: 6px 0;
    border-bottom: 1px solid rgba(255,255,255,0.06);
    font-size: 0.80rem;
    line-height: 1.45;
}
.job-row:last-child { border-bottom: none; }
.job-icon { flex-shrink: 0; margin-top: 1px; }
.job-name { font-weight: 600; word-break: break-all; }
.job-meta { opacity: 0.55; font-size: 0.74rem; }
.job-step { color: #60a5fa; font-size: 0.74rem; }
</style>
""", unsafe_allow_html=True)


# ── UI helpers ────────────────────────────────────────────────────────────────

def _file_hash(b: bytes) -> str:
    return hashlib.md5(b).hexdigest()


def _fmt(val, currency: str = "") -> str:
    if val is None:
        return "—"
    try:
        return f"{currency}{float(val):,.2f}"
    except (TypeError, ValueError):
        return str(val)


def _pdf_iframe(pdf_bytes: bytes, height: int = 540) -> None:
    """Embed PDF in an iframe — browser handles page navigation natively."""
    b64 = base64.b64encode(pdf_bytes).decode()
    st.markdown(
        f'<iframe src="data:application/pdf;base64,{b64}" '
        f'width="100%" height="{height}px" '
        f'style="border:none;border-radius:8px;background:#1a1a2e"></iframe>',
        unsafe_allow_html=True,
    )


def _save_tmp(raw: bytes, ext: str) -> str:
    f = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
    f.write(raw)
    f.close()
    return f.name


def _fa(icon: str, label: str = "") -> str:
    """Render a Font Awesome icon with optional label as HTML."""
    txt = f" {label}" if label else ""
    return f'<i class="fa-solid {icon}"></i>{txt}'


def _section(icon: str, label: str):
    st.markdown(
        f"<div class='section-label'>{_fa(icon)} {label}</div>",
        unsafe_allow_html=True,
    )


def _info_card(fa_icon: str, title: str, rows: list[tuple[str, str, str]]):
    """
    Render a bordered info card.
    rows = [(fa_icon, label, value), ...]  — rows with empty value are skipped.
    """
    body = f"<div class='card-title'>{_fa(fa_icon)} {title}</div>"
    for row_icon, lbl, val in rows:
        if val:
            label_html = f"<span style='opacity:.55'>{lbl}</span> " if lbl else ""
            body += f"<div>{_fa(row_icon)} {label_html}{val}</div>"
    st.markdown(f"<div class='info-card'>{body}</div>", unsafe_allow_html=True)


def _build_line_items_df(items: list, standardized: bool) -> pd.DataFrame:
    if not items:
        return pd.DataFrame()
    if standardized:
        return pd.DataFrame([
            {
                "Description": it.get("description") or "—",
                "Qty":         it.get("qty"),
                "Unit Price":  it.get("unit_price"),
                "Total":       it.get("total"),
            }
            for it in items
        ])
    return pd.DataFrame(items)


def _render_right(
    data: dict,
    items: list,
    extraction_mode: str,
    elapsed: float = None,
    mode_used: str = "",
    model_used: str = "",
    show_totals: bool = True,
):
    """
    Render the full right column: header details + line items table below.
    Called from streaming (per page) and from static session-state display.
    """
    cur = data.get("currency") or ""
    badge = (
        "<span class='high-badge'>HIGH ACCURACY</span>"
        if mode_used == "High Accuracy" else ""
    )
    time_html = (
        f"<span style='font-weight:400;opacity:.5;margin-left:8px'>{elapsed:.1f}s</span>"
        if elapsed else ""
    )
    model_html = (
        f"<span style='font-weight:400;opacity:.45;margin-left:6px;font-size:0.68rem'>"
        f"<i class='fa-solid fa-microchip'></i> {model_used}</span>"
        if model_used else ""
    )
    st.markdown(
        f"<div class='section-label'>"
        f"<i class='fa-solid fa-wand-magic-sparkles'></i> Extracted Details{badge}{time_html}{model_html}"
        f"</div>",
        unsafe_allow_html=True,
    )
    c1, c2, c3 = st.columns(3)
    with c1:
        st.metric("Invoice / Payment No.", data.get("payment_no") or "—")
    with c2:
        st.metric("Date", data.get("date") or "—")
    with c3:
        st.metric("Due Date", data.get("due_date") or "—")
    st.markdown("")
    vc, bc = st.columns(2)
    with vc:
        v = data.get("vendor") or {}
        _info_card("fa-building", "Vendor / Seller", [
            ("fa-signature",    "",         v.get("name")),
            ("fa-location-dot", "Address:", v.get("address")),
            ("fa-phone",        "Phone:",   v.get("phone")),
            ("fa-envelope",     "Email:",   v.get("email")),
        ])
    with bc:
        b = data.get("bill_to") or {}
        _info_card("fa-user", "Bill To", [
            ("fa-signature",    "",         b.get("name")),
            ("fa-location-dot", "Address:", b.get("address")),
        ])
    if data.get("notes"):
        st.caption(data["notes"])

    if show_totals:
        _section("fa-chart-simple", "Financial Summary")
        f1, f2, f3, f4 = st.columns(4)
        with f1:
            st.metric("Subtotal",  _fmt(data.get("subtotal"),  cur))
        with f2:
            st.metric("Tax",       _fmt(data.get("tax"),       cur))
        with f3:
            st.metric("Shipping",  _fmt(data.get("shipping"),  cur))
        with f4:
            st.metric("Discount",  _fmt(data.get("discount"),  cur))
        gap = data.get("_reconciliation_gap")
        gap_html = (
            f"<div style='font-size:0.75rem;color:#f87171;margin-top:4px'>"
            f"<i class='fa-solid fa-triangle-exclamation'></i> "
            f"Reconciliation gap: {gap:+.2f} — reconciliation line added."
            f"</div>"
        ) if gap is not None else ""
        st.markdown(
            "<div class='section-label' style='margin-top:14px'>"
            "<i class='fa-solid fa-sack-dollar'></i> Grand Total</div>"
            f"<div class='grand-total'>{_fmt(data.get('total'), cur)}</div>"
            + gap_html,
            unsafe_allow_html=True,
        )
    else:
        st.markdown(
            "<div style='font-size:0.78rem;opacity:.45;margin-top:8px'>"
            "<i class='fa-solid fa-hourglass-half'></i> "
            "Totals will appear after the final page is processed."
            "</div>",
            unsafe_allow_html=True,
        )

    # ── Line items (below details in same column) ──────────────────────────
    st.divider()
    count = len(items)
    _section("fa-list", f"Line Items  ({count})")
    if not items:
        st.info("No line items yet.")
    else:
        df = _build_line_items_df(items, standardized=(extraction_mode == "Standardized"))
        if extraction_mode == "Standardized":
            for col in ("Unit Price", "Total"):
                if col in df.columns:
                    df[col] = df[col].apply(lambda v: _fmt(v, cur) if v is not None else "—")
            computable = [it["total"] for it in items if it.get("total") is not None]
            if computable:
                footer = pd.DataFrame([{
                    "Description": "TOTAL", "Qty": "",
                    "Unit Price": "", "Total": _fmt(sum(float(x) for x in computable), cur),
                }])
                df = pd.concat([df, footer], ignore_index=True)
            st.dataframe(
                df, use_container_width=True, hide_index=True,
                column_config={
                    "Description": st.column_config.TextColumn("Description", width="large"),
                    "Qty":         st.column_config.TextColumn("Qty",         width="small"),
                    "Unit Price":  st.column_config.TextColumn("Unit Price",  width="medium"),
                    "Total":       st.column_config.TextColumn("Total",       width="medium"),
                },
            )
        else:
            skip_fmt = {"description", "code", "unit"}
            for col in df.columns:
                if col not in skip_fmt:
                    df[col] = df[col].apply(
                        lambda v: _fmt(v, cur) if isinstance(v, (int, float))
                        else (v if v is not None else "—")
                    )
            st.dataframe(df, use_container_width=True, hide_index=True)


# ── Job tracker ───────────────────────────────────────────────────────────────

_MAX_JOBS = 8


def _jobs() -> list:
    if "job_log" not in st.session_state:
        st.session_state["job_log"] = []
    return st.session_state["job_log"]


def _job_start(filename: str, mode: str, n_pages: int) -> None:
    job = {
        "filename": filename,
        "mode": mode,
        "n_pages": n_pages,
        "status": "running",
        "step": "Starting…",
        "started_at": time.time(),
        "ended_at": None,
        "items": None,
        "error": None,
    }
    log = _jobs()
    log.insert(0, job)
    if len(log) > _MAX_JOBS:
        log.pop()


def _job_update(step: str) -> None:
    log = _jobs()
    if log and log[0]["status"] == "running":
        log[0]["step"] = step


def _job_done(items: int) -> None:
    log = _jobs()
    if log and log[0]["status"] == "running":
        log[0].update(status="done", ended_at=time.time(), items=items, step="")


def _job_fail(error: str) -> None:
    log = _jobs()
    if log and log[0]["status"] == "running":
        log[0].update(status="failed", ended_at=time.time(), error=error, step="")


def _render_jobs() -> None:
    log = _jobs()
    if not log:
        st.caption("No jobs yet.")
        return
    rows_html = ""
    for job in log:
        elapsed = (
            (job["ended_at"] or time.time()) - job["started_at"]
        )
        name = job["filename"]
        name_display = name if len(name) <= 22 else name[:19] + "…"

        if job["status"] == "running":
            icon = "<i class='fa-solid fa-circle-notch fa-spin' style='color:#60a5fa'></i>"
            meta = f"{job['n_pages']}p · {elapsed:.0f}s"
            step_html = f"<div class='job-step'>{job['step']}</div>"
        elif job["status"] == "done":
            icon = "<i class='fa-solid fa-circle-check' style='color:#4ade80'></i>"
            items_str = f"{job['items']} item{'s' if job['items'] != 1 else ''}"
            meta = f"{items_str} · {elapsed:.1f}s · {job['mode']}"
            step_html = ""
        else:
            icon = "<i class='fa-solid fa-circle-xmark' style='color:#f87171'></i>"
            meta = job.get("error") or "Failed"
            step_html = ""

        rows_html += (
            f"<div class='job-row'>"
            f"  <div class='job-icon'>{icon}</div>"
            f"  <div>"
            f"    <div class='job-name'>{name_display}</div>"
            f"    <div class='job-meta'>{meta}</div>"
            f"    {step_html}"
            f"  </div>"
            f"</div>"
        )
    st.markdown(rows_html, unsafe_allow_html=True)


# ── Sidebar ───────────────────────────────────────────────────────────────────

with st.sidebar:
    st.markdown("<h3 style='margin:0'><i class='fa-solid fa-gear'></i> Settings</h3>",
                unsafe_allow_html=True)
    st.markdown("")

    st.markdown("")
    _cuda_ok = detect_cuda()
    _device_label = st.segmented_control(
        "Docling device",
        options=["Auto", "CPU", "GPU"],
        default="Auto",
        key="_device_switch",
        help=(
            "**Auto** — Docling picks the best available device.\n\n"
            "**CPU** — force CPU (slower but always works).\n\n"
            "**GPU** — force CUDA; only select if torch+CUDA is installed."
            + ("  RTX 4050 detected." if _cuda_ok else "  No CUDA detected — GPU may fail.")
        ),
    ) or "Auto"
    _device_map = {"Auto": DEVICE_AUTO, "CPU": DEVICE_CPU, "GPU": DEVICE_GPU}
    device = _device_map[_device_label]

    if _cuda_ok:
        st.caption("CUDA available — GPU ready")
    else:
        st.caption("CPU only (torch has no CUDA)")

    st.markdown("")
    extraction_mode = st.radio(
        "Line-items display",
        ["Dynamic", "Standardized"],
        help=(
            "**Standardized** — Description | Qty | Unit Price | Total\n\n"
            "**Dynamic** — all raw fields returned by the model"
        ),
    )

    st.divider()
    st.markdown(
        "<div style='font-size:0.82rem;font-weight:700;opacity:.7;margin-bottom:6px'>"
        "<i class='fa-solid fa-list-check'></i> Jobs</div>",
        unsafe_allow_html=True,
    )
    _jobs_ph = st.empty()
    with _jobs_ph.container():
        _render_jobs()

    st.divider()
    st.markdown(
        "<small><i class='fa-solid fa-diagram-project'></i> "
        "OpenCV &rarr; Docling &rarr; Ollama</small>",
        unsafe_allow_html=True,
    )
    st.markdown("")
    if st.button("Clear everything", use_container_width=True):
        for k in ("data", "text", "file_hash", "pdf_bytes", "img_bytes",
                  "is_pdf", "qa_history", "processing_mode_used"):
            st.session_state.pop(k, None)
        st.rerun()


# ── Header ────────────────────────────────────────────────────────────────────

st.markdown(
    "<h1 style='margin-bottom:0'>"
    "<i class='fa-solid fa-file-invoice-dollar'></i> DocProcessor"
    "</h1>"
    "<p style='opacity:.5;margin-top:2px'>Invoice &middot; Bill &middot; Receipt &mdash; AI-powered extraction</p>",
    unsafe_allow_html=True,
)
st.divider()

# ── Mode switcher ─────────────────────────────────────────────────────────────

_sw_l, _sw_c, _sw_r = st.columns([1, 1, 1])
with _sw_c:
    processing_mode: str = st.segmented_control(
        "Processing Mode",
        options=MODES,
        default=MODES[0],
        key="_mode_switch",
    ) or MODES[0]

st.divider()

# ── Two-column top section ────────────────────────────────────────────────────

col_preview, col_fields = st.columns([1, 1], gap="large")
_right_ph = col_fields.empty()   # defined here so button handler (left col) can write to it


# ─────────────────── LEFT: upload + preview ───────────────────────────────────

with col_preview:
    _section("fa-cloud-arrow-up", "Upload Document")
    uploaded = st.file_uploader(
        "PDF or image",
        type=["pdf", "png", "jpg", "jpeg", "tiff", "bmp", "webp"],
        label_visibility="collapsed",
    )

    if uploaded:
        raw = uploaded.read()
        ext = Path(uploaded.name).suffix.lower()
        fhash = _file_hash(raw)

        # New file → reset results and Q&A
        if st.session_state.get("file_hash") != fhash:
            for k in ("data", "text", "qa_history", "processing_mode_used"):
                st.session_state.pop(k, None)
            st.session_state["file_hash"] = fhash
            st.session_state["is_pdf"] = (ext == ".pdf")
            if ext == ".pdf":
                st.session_state["pdf_bytes"] = raw
                st.session_state.pop("img_bytes", None)
            else:
                st.session_state["img_bytes"] = raw
                st.session_state.pop("pdf_bytes", None)

        _section("fa-eye", "Preview")

        if st.session_state.get("is_pdf"):
            _pdf_iframe(st.session_state["pdf_bytes"])
        else:
            st.image(st.session_state.get("img_bytes", raw), use_container_width=True)

        st.divider()

        high = (processing_mode == "High Accuracy")
        btn_label = "Extract Data  [3 passes + voting]" if high else "Extract Data"

        if st.button(btn_label, type="primary", use_container_width=True):
            raw_bytes = (
                st.session_state.get("pdf_bytes")
                or st.session_state.get("img_bytes")
                or raw
            )
            tmp = _save_tmp(raw_bytes, ext)
            set_num_gpu(0 if device == DEVICE_CPU else 99)
            _active = [PRIMARY_MODEL]

            def _call(fn, *args, default=None, **kwargs):
                """Try fn with the active model; switch to FALLBACK_MODEL once on any failure."""
                try:
                    return fn(*args, model=_active[0], **kwargs)
                except Exception as exc:
                    if _active[0] == FALLBACK_MODEL:
                        logger.error("Fallback %s also failed: %s", FALLBACK_MODEL, exc)
                        return default
                    logger.warning("%s failed (%s) — switching to %s",
                                   _active[0], exc, FALLBACK_MODEL)
                    _active[0] = FALLBACK_MODEL
                    try:
                        return fn(*args, model=_active[0], **kwargs)
                    except Exception as exc2:
                        logger.error("Fallback %s also failed: %s", FALLBACK_MODEL, exc2)
                        return default

            try:
                t_start = time.time()
                n_pages = count_pdf_pages(tmp)
                full_pl  = high  # high accuracy uses TableFormer

                _job_start(uploaded.name, processing_mode, n_pages)
                with _jobs_ph.container():
                    _render_jobs()

                # ── Progress bar (page scan) ───────────────────────────────
                prog_bar = st.progress(0, text=f"Page 1 / {n_pages} — scanning")

                if n_pages == 1:
                    # ── Single page: normal extraction ─────────────────────
                    _job_update("OCR — scanning")
                    with _jobs_ph.container():
                        _render_jobs()
                    with st.spinner(f"Extracting [{_device_label}]"):
                        _job_update("LLM — extracting")
                        with _jobs_ph.container():
                            _render_jobs()
                        result = _call(
                            process_file_high if high else process_file,
                            tmp, device=device, default=None,
                        )
                    if result is None:
                        text, data = "", {"line_items": []}
                        st.warning(
                            f"Both {PRIMARY_MODEL} and {FALLBACK_MODEL} failed "
                            f"— no data could be extracted."
                        )
                    else:
                        text, data = result
                    prog_bar.progress(1.0, text="Done")
                    all_items = data.get("line_items") or []
                    if not all_items:
                        _job_update("LLM — pattern fallback")
                        with _jobs_ph.container():
                            _render_jobs()
                        all_items = _call(extract_line_items_fallback, text, default=[]) or []
                        data["line_items"] = all_items

                else:
                    # ── Multi-page: one request per page ───────────────────
                    all_text_parts: list[str] = []
                    all_items: list[dict] = []
                    data: dict = {}

                    for pg in range(1, n_pages + 1):
                        prog_bar.progress(
                            (pg - 0.5) / n_pages,
                            text=f"Page {pg} / {n_pages} — scanning with Docling",
                        )
                        _job_update(f"OCR — page {pg} / {n_pages}")
                        with _jobs_ph.container():
                            _render_jobs()
                        pg_text = get_page_text(tmp, pg, full_pipeline=full_pl, device=device)
                        all_text_parts.append(pg_text)

                        prog_bar.progress(
                            (pg - 0.2) / n_pages,
                            text=f"Page {pg} / {n_pages} — sending to {_active[0]}",
                        )
                        _job_update(f"LLM — page {pg} / {n_pages}")
                        with _jobs_ph.container():
                            _render_jobs()

                        is_last = (pg == n_pages)

                        if pg == 1:
                            # Header-only call (fast — no line items, no totals)
                            data = _call(extract_header, pg_text, default={}) or {}
                            # Line items from page 1
                            pg1_items = _call(extract_line_items_page, pg_text, 1, default=[]) or []
                            if not pg1_items:
                                _job_update(f"LLM — pattern fallback p{pg}")
                                with _jobs_ph.container():
                                    _render_jobs()
                                pg1_items = _call(extract_line_items_fallback, pg_text, default=[]) or []
                            all_items = pg1_items
                            # Show header immediately; totals hidden (they're on the last page)
                            with _right_ph.container():
                                _render_right(data, all_items, extraction_mode,
                                              elapsed=None, mode_used=processing_mode,
                                              model_used=_active[0], show_totals=False)
                        elif is_last:
                            # Line items from last page
                            new_items = _call(extract_line_items_page, pg_text, pg, default=[]) or []
                            if not new_items:
                                new_items = _call(extract_line_items_fallback, pg_text, default=[]) or []
                            all_items.extend(new_items)
                            # Totals-only call on last page
                            totals = _call(extract_totals, pg_text, default={}) or {}
                            data.update({k: v for k, v in totals.items() if v is not None})
                        else:
                            new_items = _call(extract_line_items_page, pg_text, pg, default=[]) or []
                            if not new_items:
                                new_items = _call(extract_line_items_fallback, pg_text, default=[]) or []
                            all_items.extend(new_items)

                        # Update right column; show totals only on last page
                        with _right_ph.container():
                            _render_right(data, all_items, extraction_mode,
                                          elapsed=None, mode_used=processing_mode,
                                          model_used=_active[0], show_totals=is_last)

                        prog_bar.progress(pg / n_pages, text=f"Page {pg} / {n_pages} done")

                    # Merge all pages' text for raw text tab
                    text = "\n\n".join(all_text_parts)
                    data["line_items"] = all_items

                # Math validation for multi-page (high mode does it internally)
                if n_pages > 1:
                    from llmplus import _validate_math
                    data = _validate_math(data)

                elapsed = time.time() - t_start
                prog_bar.progress(1.0, text=f"Done in {elapsed:.1f}s")

                item_count = len(data.get("line_items") or [])
                _job_done(item_count)
                with _jobs_ph.container():
                    _render_jobs()

                st.session_state.update(
                    text=text, data=data,
                    qa_history=[],
                    processing_mode_used=processing_mode,
                    model_used=_active[0],
                    elapsed=elapsed,
                )
                st.success(
                    f"Done in {elapsed:.1f}s · {_active[0]} · "
                    f"{item_count} line item{'s' if item_count != 1 else ''} extracted."
                )
                st.rerun()

            except json.JSONDecodeError as e:
                logger.error(
                    "Streamlit: malformed JSON [file=%s  model=%s  mode=%s]: %s",
                    uploaded.name, _active[0], processing_mode, e,
                )
                _job_fail("Malformed JSON")
                with _jobs_ph.container():
                    _render_jobs()
                st.error(f"Model returned malformed JSON — try again.\n\n{e}")
            except Exception as e:
                logger.error(
                    "Streamlit: extraction failed [file=%s  model=%s  mode=%s]: %s",
                    uploaded.name, _active[0], processing_mode, e,
                    exc_info=True,
                )
                _job_fail(str(e)[:60])
                with _jobs_ph.container():
                    _render_jobs()
                st.error(f"Extraction failed: {e}")
            finally:
                try:
                    os.unlink(tmp)
                except OSError:
                    pass


# ─────────────────── RIGHT: extracted fields ──────────────────────────────────

if "data" not in st.session_state:
    _right_ph.info("Upload a document and click Extract Data to see results here.")
else:
    with _right_ph.container():
        _render_right(
            st.session_state["data"],
            items=st.session_state["data"].get("line_items") or [],
            extraction_mode=extraction_mode,
            elapsed=st.session_state.get("elapsed"),
            mode_used=st.session_state.get("processing_mode_used", ""),
            model_used=st.session_state.get("model_used", ""),
        )


# ── Bottom tabs ───────────────────────────────────────────────────────────────

if "data" in st.session_state:
    st.divider()
    data = st.session_state["data"]
    cur = data.get("currency") or ""
    items: list = data.get("line_items") or []

    tab_qa, tab_json, tab_raw = st.tabs(["Q & A", "JSON", "Raw Text"])

    # ── Q & A ─────────────────────────────────────────────────
    with tab_qa:
        if "qa_history" not in st.session_state:
            st.session_state["qa_history"] = []
        history: list = st.session_state["qa_history"]

        if not history:
            _section("fa-lightbulb", "Suggested Questions")
            cols = st.columns(len(QA_SUGGESTIONS))
            for col, suggestion in zip(cols, QA_SUGGESTIONS):
                with col:
                    if st.button(suggestion, use_container_width=True, key=f"sq_{suggestion}"):
                        st.session_state["qa_pending"] = suggestion
                        st.rerun()
            st.markdown("")

        for msg in history:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])

        pending = st.session_state.pop("qa_pending", None)
        user_input = st.chat_input("Ask anything about this document…") or pending

        if user_input:
            history.append({"role": "user", "content": user_input})
            with st.chat_message("user"):
                st.markdown(user_input)

            with st.chat_message("assistant"):
                response = st.write_stream(
                    stream_qa(
                        question=user_input,
                        doc_text=st.session_state.get("text", ""),
                        doc_data=data,
                        history=history[:-1],
                        model=st.session_state.get("model_used", PRIMARY_MODEL),
                    )
                )

            history.append({"role": "assistant", "content": response})
            st.session_state["qa_history"] = history

        if history:
            if st.button("Clear conversation", key="clear_qa"):
                st.session_state["qa_history"] = []
                st.rerun()

    # ── JSON ──────────────────────────────────────────────────
    with tab_json:
        # Strip internal meta key before display
        display_data = {k: v for k, v in data.items() if not k.startswith("_")}
        st.json(display_data, expanded=True)
        st.download_button(
            "Download JSON",
            data=json.dumps(display_data, indent=2, ensure_ascii=False),
            file_name="extracted_invoice.json",
            mime="application/json",
            use_container_width=True,
        )

    # ── Raw Text ──────────────────────────────────────────────
    with tab_raw:
        if "text" in st.session_state:
            _section("fa-file-lines", "Docling Markdown Output")
            st.text_area(
                "raw", value=st.session_state["text"],
                height=500, label_visibility="collapsed",
            )
            st.download_button(
                "Download raw text",
                data=st.session_state["text"],
                file_name="extracted_text.md",
                mime="text/markdown",
                use_container_width=True,
            )
        else:
            st.info("No raw text available.")