smoke_signal/scripts/06_review_workbench.py

#!/usr/bin/env python3
"""
Smoke Signal — Stage 8: Human Review Workbench
================================================
Gradio app for reviewing low-confidence OCR pages.

Reviewer actions per page:
  - ACCEPT   : text is correct, pass to Codex export
  - EDIT     : correct the text, then accept
  - REJECT   : unusable, exclude from export
  - QUARANTINE: flag for specialist review
  - ILLUSTRATION ONLY: no text on this page

Captures: reviewer ID, timestamp, edits, reason codes, final status.

To run locally:
    pip install gradio
    python scripts/06_review_workbench.py

To deploy on HF Spaces:
    This file should be copied to the root as smoke_signal_review.py
    or integrated into the main app.py as a new tab.
"""

import csv
import json
import os
from datetime import datetime
from pathlib import Path
from typing import Optional

import gradio as gr

# ── Paths ──────────────────────────────────────────────────────────────────────
ROOT        = Path(__file__).resolve().parents[1]
REVIEW_DIR  = ROOT / "review"
REGIONS_DIR = ROOT / "regions"
CLEANED_DIR = ROOT / "regions" / "cleaned"
RENDERS_DIR = ROOT / "renders"
EXPORTS_DIR = ROOT / "exports"

REVIEW_DIR.mkdir(parents=True, exist_ok=True)
EXPORTS_DIR.mkdir(parents=True, exist_ok=True)

QUEUE_CSV    = REVIEW_DIR / "review_queue.csv"
DECISIONS_CSV = REVIEW_DIR / "review_decisions.csv"

# ── Reason codes ───────────────────────────────────────────────────────────────
REASON_CODES = [
    "OCR_MISS",
    "OCR_WRONG_WORD",
    "REGION_MISSING",
    "REGION_FALSE_POSITIVE",
    "READING_ORDER_ERROR",
    "DECORATIVE_FONT",
    "SPEECH_BUBBLE_ERROR",
    "LOW_CONTRAST",
    "SCAN_SKEW_BLUR",
    "NON_STORY_TEXT",
    "RIGHTS_UNCLEAR",
    "DUPLICATE_SOURCE",
    "LLM_OVER_CORRECTION",
    "MANUAL_TRANSCRIPTION_REQUIRED",
    "OTHER",
]

REVIEW_STATUSES = ["pending", "accepted", "edited", "rejected", "quarantined", "illustration-only"]

# ── CSS ────────────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Playfair+Display:wght@700;900&family=Source+Code+Pro:wght@400;600&family=Lato:wght@300;400;700&display=swap');

:root {
    --ink:      #1a1a2e;
    --paper:    #f5f0e8;
    --smoke:    #2d3561;
    --signal:   #e94560;
    --ash:      #8892b0;
    --accepted: #00b894;
    --rejected: #e17055;
    --quarantine: #fdcb6e;
    --pending:  #74b9ff;
}

.gradio-container {
    background: var(--paper) !important;
    font-family: 'Lato', sans-serif !important;
    max-width: none !important;
}

footer { display: none !important; }

#ss-header {
    background: var(--ink);
    padding: 20px 32px;
    border-bottom: 3px solid var(--signal);
    display: flex;
    align-items: center;
    gap: 20px;
}

#ss-title {
    font-family: 'Playfair Display', serif;
    font-size: 28px;
    font-weight: 900;
    color: var(--paper);
    letter-spacing: -0.5px;
    margin: 0;
}

#ss-subtitle {
    font-family: 'Source Code Pro', monospace;
    font-size: 11px;
    color: var(--ash);
    letter-spacing: 3px;
    text-transform: uppercase;
    margin: 0;
}

#ss-signal {
    color: var(--signal);
    font-size: 36px;
    font-weight: 900;
}

.queue-panel {
    background: white;
    border: 1px solid #e0d9cc;
    border-radius: 8px;
    padding: 16px;
    height: 600px;
    overflow-y: auto;
}

.queue-item {
    padding: 12px 14px;
    border-radius: 6px;
    margin-bottom: 8px;
    cursor: pointer;
    border: 2px solid transparent;
    transition: all 0.15s;
    font-size: 13px;
}

.queue-item:hover { border-color: var(--smoke); }
.queue-item.active { border-color: var(--signal); background: #fff5f7; }
.queue-item.pending   { border-left: 4px solid var(--pending); }
.queue-item.accepted  { border-left: 4px solid var(--accepted); opacity: 0.6; }
.queue-item.rejected  { border-left: 4px solid var(--rejected); opacity: 0.6; }
.queue-item.quarantined { border-left: 4px solid var(--quarantine); }

.page-image-panel {
    background: #2a2a2a;
    border-radius: 8px;
    min-height: 400px;
    display: flex;
    align-items: center;
    justify-content: center;
}

.confidence-badge {
    display: inline-block;
    padding: 3px 10px;
    border-radius: 999px;
    font-size: 12px;
    font-weight: 700;
    font-family: 'Source Code Pro', monospace;
}

.conf-high    { background: #d4f5e9; color: #00695c; }
.conf-medium  { background: #fff3cd; color: #856404; }
.conf-low     { background: #fde8e8; color: #c62828; }
.conf-quarantine { background: #2a2a2a; color: #fdcb6e; }

.action-btn {
    font-weight: 700 !important;
    font-size: 14px !important;
    border-radius: 6px !important;
    min-height: 44px !important;
    transition: transform 0.1s !important;
}

.action-btn:active { transform: scale(0.97) !important; }

.accept-btn  { background: var(--accepted) !important; color: white !important; }
.reject-btn  { background: var(--rejected) !important; color: white !important; }
.quar-btn    { background: var(--quarantine) !important; color: var(--ink) !important; }
.illus-btn   { background: var(--smoke) !important; color: white !important; }

.stats-bar {
    background: var(--ink);
    color: var(--paper);
    padding: 10px 20px;
    border-radius: 6px;
    font-family: 'Source Code Pro', monospace;
    font-size: 12px;
    display: flex;
    gap: 24px;
    margin-bottom: 12px;
}

.stat-item { display: flex; flex-direction: column; gap: 2px; }
.stat-value { font-size: 20px; font-weight: 600; }
.stat-label { color: var(--ash); font-size: 10px; letter-spacing: 1px; }

.ocr-text-box textarea {
    font-family: 'Source Code Pro', monospace !important;
    font-size: 14px !important;
    background: #fafaf8 !important;
    border: 2px solid #e0d9cc !important;
    border-radius: 6px !important;
}

.ocr-text-box textarea:focus {
    border-color: var(--signal) !important;
}
"""

# ── Data loading ───────────────────────────────────────────────────────────────
def load_queue() -> list:
    """Load the review queue CSV."""
    if not QUEUE_CSV.exists():
        return []
    with open(QUEUE_CSV, newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def load_decisions() -> dict:
    """Load existing decisions keyed by region_id."""
    decisions = {}
    if not DECISIONS_CSV.exists():
        return decisions
    with open(DECISIONS_CSV, newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            decisions[row.get("region_id", "")] = row
    return decisions


def save_decision(
    region_id: str,
    book_id: str,
    page: int,
    status: str,
    final_text: str,
    reason_code: str,
    reviewer: str,
    notes: str,
) -> None:
    """Append or update a decision record."""
    fields = [
        "region_id", "book_id", "page", "status",
        "final_text", "reason_code", "reviewer", "notes", "decided_at"
    ]
    existing = load_decisions()
    existing[region_id] = {
        "region_id":  region_id,
        "book_id":    book_id,
        "page":       page,
        "status":     status,
        "final_text": final_text,
        "reason_code": reason_code,
        "reviewer":   reviewer,
        "notes":      notes,
        "decided_at": datetime.utcnow().isoformat() + "Z",
    }
    write_header = not DECISIONS_CSV.exists()
    with open(DECISIONS_CSV, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        writer.writerows(existing.values())


def get_queue_stats(queue: list, decisions: dict) -> dict:
    total     = len(queue)
    decided   = len(decisions)
    pending   = total - decided
    accepted  = sum(1 for d in decisions.values() if d["status"] == "accepted")
    edited    = sum(1 for d in decisions.values() if d["status"] == "edited")
    rejected  = sum(1 for d in decisions.values() if d["status"] == "rejected")
    quarantined = sum(1 for d in decisions.values() if d["status"] == "quarantined")
    return {
        "total": total, "pending": pending, "accepted": accepted,
        "edited": edited, "rejected": rejected, "quarantined": quarantined,
    }


# ── Image loader ───────────────────────────────────────────────────────────────
def get_page_image(book_id: str, page: str) -> Optional[str]:
    """Find the rendered page image."""
    try:
        page_num = int(page)
    except (ValueError, TypeError):
        return None

    book_dir = RENDERS_DIR / str(book_id)
    if book_dir.exists():
        candidates = sorted(book_dir.glob(f"{book_id}_page_{page_num:04d}_*.png"))
        if candidates:
            return str(candidates[0])
    return None


# ── Queue HTML builder ─────────────────────────────────────────────────────────
def build_queue_html(queue: list, decisions: dict, active_idx: int = 0) -> str:
    if not queue:
        return "<div style='padding:20px;color:#888;font-family:monospace'>No items in review queue.<br>Run the pipeline first.</div>"

    items = []
    for i, item in enumerate(queue):
        region_id = item.get("region_id", "")
        decision  = decisions.get(region_id, {})
        status    = decision.get("status", item.get("status", "pending"))
        conf      = float(item.get("confidence", 0))
        page      = item.get("page", "?")
        book_id   = item.get("book_id", "?")
        conf_str  = f"{conf:.0%}"
        active    = "active" if i == active_idx else ""
        items.append(f"""
        <div class="queue-item {status} {active}" onclick="selectItem({i})" id="qi-{i}">
            <div style="display:flex;justify-content:space-between;align-items:center">
                <span style="font-weight:700;color:#1a1a2e">{book_id} · p{page}</span>
                <span style="font-size:11px;color:#888">{status.upper()}</span>
            </div>
            <div style="margin-top:4px;font-size:12px;color:#555">
                conf: <b style="color:{'#c62828' if conf < 0.6 else '#856404' if conf < 0.85 else '#00695c'}">{conf_str}</b>
                &nbsp;·&nbsp; {item.get('region_class','?')}
            </div>
        </div>""")

    return f"<div class='queue-panel'>{''.join(items)}</div>"


def build_stats_html(stats: dict) -> str:
    progress = (stats['accepted'] + stats['edited']) / max(stats['total'], 1) * 100
    return f"""
    <div class="stats-bar">
        <div class="stat-item"><span class="stat-value">{stats['total']}</span><span class="stat-label">TOTAL</span></div>
        <div class="stat-item"><span class="stat-value" style="color:var(--pending)">{stats['pending']}</span><span class="stat-label">PENDING</span></div>
        <div class="stat-item"><span class="stat-value" style="color:var(--accepted)">{stats['accepted'] + stats['edited']}</span><span class="stat-label">APPROVED</span></div>
        <div class="stat-item"><span class="stat-value" style="color:var(--rejected)">{stats['rejected']}</span><span class="stat-label">REJECTED</span></div>
        <div class="stat-item"><span class="stat-value" style="color:var(--quarantine)">{stats['quarantined']}</span><span class="stat-label">QUARANTINED</span></div>
        <div class="stat-item" style="flex:1">
            <span class="stat-label">PROGRESS</span>
            <div style="background:#333;border-radius:4px;height:8px;margin-top:6px">
                <div style="background:var(--accepted);width:{progress:.0f}%;height:8px;border-radius:4px;transition:width 0.3s"></div>
            </div>
        </div>
    </div>"""


# ── Gradio app ────────────────────────────────────────────────────────────────
def build_app():

    queue     = load_queue()
    decisions = load_decisions()
    state     = {"idx": 0, "queue": queue, "decisions": decisions}

    def get_current_item():
        q = state["queue"]
        if not q:
            return None
        idx = min(state["idx"], len(q) - 1)
        return q[idx]

    def refresh_view():
        queue     = state["queue"]
        decisions = state["decisions"]
        item      = get_current_item()
        stats     = get_queue_stats(queue, decisions)
        stats_html = build_stats_html(stats)
        queue_html = build_queue_html(queue, decisions, state["idx"])

        if not item:
            return (
                stats_html, queue_html,
                None, "", "", "", "pending", "", "",
                "No items in queue"
            )

        region_id = item.get("region_id", "")
        decision  = decisions.get(region_id, {})
        book_id   = item.get("book_id", "")
        page      = item.get("page", "")
        conf      = float(item.get("confidence", 0))
        conf_class = "conf-low" if conf < 0.6 else "conf-medium" if conf < 0.85 else "conf-high"

        raw_text   = item.get("raw_ocr", "")
        final_text = decision.get("final_text", raw_text)
        status     = decision.get("status", "pending")
        reason     = decision.get("reason_code", "")
        reviewer   = decision.get("reviewer", "")
        notes      = decision.get("notes", "")

        img_path = get_page_image(book_id, page)

        info = f"<span class='confidence-badge {conf_class}'>conf: {conf:.0%}</span> &nbsp; {book_id} · page {page} · {item.get('region_class','?')}"

        return (
            stats_html, queue_html,
            img_path, raw_text, final_text,
            info, status, reason, reviewer, notes
        )

    def navigate(direction: int):
        q = state["queue"]
        if not q:
            return refresh_view()
        state["idx"] = max(0, min(state["idx"] + direction, len(q) - 1))
        return refresh_view()

    def submit_decision(final_text, status, reason_code, reviewer, notes):
        item = get_current_item()
        if not item:
            return refresh_view()

        region_id = item.get("region_id", "")
        book_id   = item.get("book_id", "")
        page      = item.get("page", "")

        # Determine actual status
        raw_text   = item.get("raw_ocr", "")
        act_status = status
        if status == "accepted" and final_text.strip() != raw_text.strip():
            act_status = "edited"

        save_decision(
            region_id=region_id,
            book_id=book_id,
            page=page,
            status=act_status,
            final_text=final_text,
            reason_code=reason_code,
            reviewer=reviewer or "reviewer",
            notes=notes,
        )

        state["decisions"] = load_decisions()

        # Auto-advance to next pending item
        q = state["queue"]
        for i in range(state["idx"] + 1, len(q)):
            rid = q[i].get("region_id", "")
            if rid not in state["decisions"]:
                state["idx"] = i
                break

        return refresh_view()

    def quick_action(action: str, reviewer_name: str):
        item = get_current_item()
        if not item:
            return refresh_view()
        region_id = item.get("region_id", "")
        book_id   = item.get("book_id", "")
        page      = item.get("page", "")
        raw_text  = item.get("raw_ocr", "")

        status_map = {
            "accept":    "accepted",
            "reject":    "rejected",
            "quarantine": "quarantined",
            "illus":     "illustration-only",
        }
        save_decision(
            region_id=region_id, book_id=book_id, page=page,
            status=status_map.get(action, "accepted"),
            final_text=raw_text, reason_code="", reviewer=reviewer_name or "reviewer", notes="",
        )
        state["decisions"] = load_decisions()

        # Auto-advance
        q = state["queue"]
        for i in range(state["idx"] + 1, len(q)):
            rid = q[i].get("region_id", "")
            if rid not in state["decisions"]:
                state["idx"] = i
                break

        return refresh_view()

    def export_approved():
        """Export all accepted/edited decisions to JSONL for Codex."""
        decisions = load_decisions()
        approved  = [d for d in decisions.values() if d["status"] in ("accepted", "edited")]

        if not approved:
            return "No approved items to export yet."

        ts = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        out_path = EXPORTS_DIR / f"review_export_{ts}.jsonl"
        with open(out_path, "w", encoding="utf-8") as f:
            for d in approved:
                f.write(json.dumps(d) + "\n")

        return f"Exported {len(approved)} approved records → {out_path.relative_to(ROOT)}"

    # ── Layout ────────────────────────────────────────────────────────────────
    with gr.Blocks(title="Smoke Signal — Review Workbench", css=CSS) as app:

        gr.HTML("""
        <div id="ss-header">
            <span id="ss-signal">◈</span>
            <div>
                <p id="ss-title">Smoke Signal</p>
                <p id="ss-subtitle">OCR Review Workbench · Human-in-the-Loop</p>
            </div>
        </div>
        """)

        # Stats bar
        stats_html = gr.HTML()

        with gr.Row():
            # ── Left: queue ───────────────────────────────────────────────────
            with gr.Column(scale=1):
                gr.Markdown("### Review Queue")
                queue_html = gr.HTML()

                with gr.Row():
                    prev_btn = gr.Button("← Prev", size="sm")
                    next_btn = gr.Button("Next →", size="sm")

                reviewer_name = gr.Textbox(
                    label="Your name / ID",
                    placeholder="e.g. jamal",
                    scale=1,
                )

            # ── Centre: page image ────────────────────────────────────────────
            with gr.Column(scale=2):
                gr.Markdown("### Page Image")
                page_image = gr.Image(
                    label="",
                    type="filepath",
                    height=480,
                    show_download_button=False,
                )
                item_info = gr.HTML()

            # ── Right: text + actions ─────────────────────────────────────────
            with gr.Column(scale=2):
                gr.Markdown("### OCR Text")
                raw_text_box = gr.Textbox(
                    label="Raw OCR (read-only)",
                    lines=5,
                    interactive=False,
                    elem_classes=["ocr-text-box"],
                )
                final_text_box = gr.Textbox(
                    label="Final Text (edit if needed)",
                    lines=8,
                    interactive=True,
                    elem_classes=["ocr-text-box"],
                )

                with gr.Row():
                    accept_btn = gr.Button("✓ Accept", elem_classes=["action-btn", "accept-btn"])
                    reject_btn = gr.Button("✗ Reject", elem_classes=["action-btn", "reject-btn"])

                with gr.Row():
                    quar_btn  = gr.Button("⚑ Quarantine", elem_classes=["action-btn", "quar-btn"])
                    illus_btn = gr.Button("◉ Illus Only", elem_classes=["action-btn", "illus-btn"])

                gr.Markdown("### Decision Details")
                status_dd = gr.Dropdown(
                    label="Status",
                    choices=REVIEW_STATUSES,
                    value="pending",
                )
                reason_dd = gr.Dropdown(
                    label="Reason Code",
                    choices=[""] + REASON_CODES,
                    value="",
                )
                notes_box = gr.Textbox(label="Notes", lines=2)

                submit_btn = gr.Button(
                    "Save Decision",
                    variant="primary",
                    size="lg",
                )

                export_btn = gr.Button("⬇ Export Approved to Codex", variant="secondary")
                export_status = gr.Textbox(label="Export status", interactive=False)

        # ── Wire events ───────────────────────────────────────────────────────
        outputs = [
            stats_html, queue_html,
            page_image, raw_text_box, final_text_box,
            item_info, status_dd, reason_dd, reviewer_name, notes_box,
        ]

        app.load(refresh_view, outputs=outputs)

        prev_btn.click(lambda: navigate(-1), outputs=outputs)
        next_btn.click(lambda: navigate(1), outputs=outputs)

        accept_btn.click(
            lambda rev: quick_action("accept", rev),
            inputs=[reviewer_name], outputs=outputs
        )
        reject_btn.click(
            lambda rev: quick_action("reject", rev),
            inputs=[reviewer_name], outputs=outputs
        )
        quar_btn.click(
            lambda rev: quick_action("quarantine", rev),
            inputs=[reviewer_name], outputs=outputs
        )
        illus_btn.click(
            lambda rev: quick_action("illus", rev),
            inputs=[reviewer_name], outputs=outputs
        )

        submit_btn.click(
            submit_decision,
            inputs=[final_text_box, status_dd, reason_dd, reviewer_name, notes_box],
            outputs=outputs,
        )

        export_btn.click(export_approved, outputs=[export_status])

    return app


if __name__ == "__main__":
    app = build_app()
    app.launch(share=False)