app.py

"""Gradio app: OCR documents into structured notification/data records.

Upload scanned PDFs or images (construction reports, compliance
certificates, invoices, etc.), scan the extracted text for domain
"notification" keywords, pull out labeled data fields (standard IDs/dates/
amounts plus any custom labels the user names), and export everything to
CSV as tidy (one-fact-per-row) tables, alongside an audit record of the
settings that produced the run.
"""
import os
import tempfile
import uuid
from datetime import datetime

import gradio as gr
import pandas as pd

import chat_config
import extractors
import ocr_engine

CATEGORY_CHOICES = chat_config.CATEGORY_CHOICES

DEFAULT_SETTINGS = {
    "categories": CATEGORY_CHOICES,
    "custom_keywords": [],
    "custom_fields": [],
    "only_matches": True,
}

FIELDS_COLUMNS = [
    "run_id", "file_name", "page", "field_type", "field_name", "field_value",
    "text_source", "ocr_confidence",
]
NOTIFICATIONS_COLUMNS = [
    "run_id", "file_name", "page", "category", "keyword", "context",
    "text_source", "ocr_confidence",
]
METADATA_COLUMNS = [
    "run_id", "timestamp", "files", "categories", "custom_keywords",
    "custom_fields", "only_matches_filter", "chat_model", "chat_transcript",
]

from sample_docs.generate_samples import ensure_samples

# The demo PDFs are generated on first launch rather than shipped in git
# (Hugging Face rejects binary files outside LFS storage).
SAMPLE_DIR = ensure_samples()

# One-click demo documents (all fictitious, clearly watermarked). Each button
# loads the file AND presets the matching extraction settings so "Run" works
# immediately.
SAMPLE_DOCS = {
    "Construction change order": {
        "file": os.path.join(SAMPLE_DIR, "sample_construction_change_order.pdf"),
        "categories": ["construction", "compliance", "general"],
        "keywords": "",
        "fields": "contract ref",
    },
    "Accounting / tax invoice": {
        "file": os.path.join(SAMPLE_DIR, "sample_accounting_tax_invoice.pdf"),
        "categories": ["finance"],
        "keywords": "",
        "fields": "tax year",
    },
    "Bill of lading (scanned)": {
        "file": os.path.join(SAMPLE_DIR, "sample_bill_of_lading.pdf"),
        "categories": ["construction", "general"],
        "keywords": "",
        "fields": "bol number, reference number, gross weight, net weight, delivery date, carrier",
    },
}

QUICK_PROMPTS = {
    "Construction & SAP data": (
        "Flag change orders, delay notices, and safety violations. "
        "Pull PO numbers, SAP document numbers, and reference numbers."
    ),
    "Regulatory compliance": (
        "Flag non-compliance, violations, and expired certifications. "
        "Pull certification codes and expiration dates."
    ),
    "Finance & accounts receivable": (
        "Flag overdue and past-due invoices, credit holds, and disputes. "
        "Pull invoice numbers, amounts, and due dates."
    ),
}

CUSTOM_CSS = """
#hero { text-align: center; padding: 28px 16px 8px 16px; }
#hero h1 {
    font-size: 2.1rem; font-weight: 700; letter-spacing: -0.02em;
    margin-bottom: 8px; color: var(--body-text-color);
}
#hero p {
    color: var(--body-text-color-subdued); max-width: 640px; margin: 0 auto;
    font-size: 1.02rem; line-height: 1.5;
}
.pipeline-strip {
    display: flex; align-items: center; justify-content: center; flex-wrap: wrap;
    gap: 6px; margin: 4px auto 28px auto; max-width: 900px; padding: 0 16px;
}
.pipeline-step {
    display: flex; align-items: center; gap: 7px;
    font-size: 0.78rem; font-weight: 600; letter-spacing: 0.01em;
    color: var(--body-text-color-subdued);
    background: var(--background-fill-secondary);
    border: 1px solid var(--border-color-primary);
    border-radius: 999px; padding: 5px 13px; white-space: nowrap;
}
.pipeline-step .num {
    display: inline-flex; align-items: center; justify-content: center;
    width: 16px; height: 16px; border-radius: 50%;
    background: var(--primary-500); color: white; font-size: 0.65rem;
}
.pipeline-arrow { color: var(--body-text-color-subdued); font-size: 0.85rem; opacity: 0.6; }

.step-card {
    border-radius: 14px !important; padding: 24px !important; margin-bottom: 18px;
    box-shadow: 0 1px 2px rgba(0,0,0,0.04), 0 1px 8px rgba(0,0,0,0.03) !important;
}
.step-eyebrow {
    font-size: 0.72rem; font-weight: 700; text-transform: uppercase;
    letter-spacing: 0.08em; color: var(--primary-500); margin-bottom: 4px;
}
.step-kicker { font-weight: 600; font-size: 1.15rem; margin-bottom: 2px; }
.step-subtitle { color: var(--body-text-color-subdued); margin-bottom: 14px; font-size: 0.92rem; }
#run-button button { font-size: 1.02rem !important; padding: 14px !important; font-weight: 600 !important; }
#quick-prompts button, #sample-docs button { font-size: 0.85rem !important; }
.sample-note {
    font-size: 0.8rem; color: var(--body-text-color-subdued);
    margin: 10px 0 6px 0;
}
.sample-note strong { color: var(--body-text-color); }
"""

PIPELINE_HTML = """
<div class="pipeline-strip">
  <div class="pipeline-step"><span class="num">1</span>Upload</div>
  <div class="pipeline-arrow">&rarr;</div>
  <div class="pipeline-step"><span class="num">2</span>OCR Recognition</div>
  <div class="pipeline-arrow">&rarr;</div>
  <div class="pipeline-step"><span class="num">3</span>AI-Assisted Extraction</div>
  <div class="pipeline-arrow">&rarr;</div>
  <div class="pipeline-step"><span class="num">4</span>Structured Data</div>
  <div class="pipeline-arrow">&rarr;</div>
  <div class="pipeline-step"><span class="num">5</span>CSV Export</div>
</div>
"""


def _file_path(f):
    return f if isinstance(f, str) else f.name


def _parse_comma_list(s):
    return [k.strip() for k in s.split(",") if k.strip()] if s and s.strip() else []


def _format_transcript(history):
    if not history:
        return ""
    lines = []
    for turn in history:
        role = turn.get("role", "?")
        content = turn.get("content", "")
        lines.append(f"{role}: {content}")
    return "\n".join(lines)


def process_files(files, categories, custom_keywords_str, custom_fields_str, only_matches, chat_history):
    empty = (pd.DataFrame(columns=FIELDS_COLUMNS), pd.DataFrame(columns=NOTIFICATIONS_COLUMNS))
    if not files:
        return "Upload at least one PDF or image file.", empty[0], empty[1], None, "", {}

    custom_keywords = _parse_comma_list(custom_keywords_str)
    custom_field_labels = _parse_comma_list(custom_fields_str)

    run_id = uuid.uuid4().hex[:8]
    field_rows = []
    notification_rows = []
    full_text_log = []
    warnings = []
    total_pages = 0
    file_names = []

    for f in files:
        path = _file_path(f)
        file_name = os.path.basename(path)
        file_names.append(file_name)
        try:
            pages = ocr_engine.extract_pages(path)
        except Exception as exc:
            warnings.append(f"{file_name}: {exc}")
            continue

        for page_data in pages:
            total_pages += 1
            text = page_data["text"]
            source = page_data["source"]
            confidence = page_data["ocr_confidence"]

            notifications = extractors.find_notifications(text, categories, custom_keywords)
            std_fields = extractors.extract_fields(text, categories)
            custom_field_values = extractors.extract_custom_fields(text, custom_field_labels)

            full_text_log.append(f"--- {file_name} | page {page_data['page']} | {source} ---\n{text}\n")

            has_content = bool(notifications) or any(std_fields.values()) or any(custom_field_values.values())
            if only_matches and not has_content:
                continue

            for field_name, values in std_fields.items():
                for value in values:
                    field_rows.append({
                        "run_id": run_id, "file_name": file_name, "page": page_data["page"],
                        "field_type": "standard", "field_name": field_name, "field_value": value,
                        "text_source": source, "ocr_confidence": confidence,
                    })
            for field_name, values in custom_field_values.items():
                for value in values:
                    field_rows.append({
                        "run_id": run_id, "file_name": file_name, "page": page_data["page"],
                        "field_type": "custom", "field_name": field_name, "field_value": value,
                        "text_source": source, "ocr_confidence": confidence,
                    })
            for n in notifications:
                notification_rows.append({
                    "run_id": run_id, "file_name": file_name, "page": page_data["page"],
                    "category": n["category"], "keyword": n["keyword"], "context": n["context"],
                    "text_source": source, "ocr_confidence": confidence,
                })

    fields_df = pd.DataFrame(field_rows, columns=FIELDS_COLUMNS)
    notifications_df = pd.DataFrame(notification_rows, columns=NOTIFICATIONS_COLUMNS)

    metadata_row = {
        "run_id": run_id,
        "timestamp": datetime.now().isoformat(timespec="seconds"),
        "files": "; ".join(file_names),
        "categories": "; ".join(categories),
        "custom_keywords": "; ".join(custom_keywords),
        "custom_fields": "; ".join(custom_field_labels),
        "only_matches_filter": only_matches,
        "chat_model": chat_config.ACTIVE_MODEL if chat_history else "",
        "chat_transcript": _format_transcript(chat_history),
    }
    metadata_df = pd.DataFrame([metadata_row], columns=METADATA_COLUMNS)

    csv_paths = []
    if not fields_df.empty or not notifications_df.empty:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        tmp_dir = tempfile.gettempdir()
        fields_path = os.path.join(tmp_dir, f"ocr_fields_{timestamp}_{run_id}.csv")
        notifications_path = os.path.join(tmp_dir, f"ocr_notifications_{timestamp}_{run_id}.csv")
        metadata_path = os.path.join(tmp_dir, f"ocr_run_metadata_{timestamp}_{run_id}.csv")
        fields_df.to_csv(fields_path, index=False)
        notifications_df.to_csv(notifications_path, index=False)
        metadata_df.to_csv(metadata_path, index=False)
        csv_paths = [fields_path, notifications_path, metadata_path]

    status_lines = [
        f"Processed {len(files)} file(s), {total_pages} page(s) total.",
        f"{len(fields_df)} field value(s), {len(notifications_df)} notification hit(s)"
        + (" (matches-only pages)." if only_matches else " (all pages)."),
    ]
    if warnings:
        status_lines.append("Warnings: " + "; ".join(warnings))
    status = "\n".join(status_lines)

    return status, fields_df, notifications_df, csv_paths, "\n".join(full_text_log), metadata_row


def handle_chat(history, message, settings, hf_token):
    history = history or []
    if not message or not message.strip():
        return (
            history, "", settings, settings,
            settings["categories"], ", ".join(settings["custom_keywords"]),
            ", ".join(settings["custom_fields"]), settings["only_matches"],
        )

    reply, new_settings = chat_config.chat_update(history, message.strip(), settings, hf_token)
    history = history + [
        {"role": "user", "content": message.strip()},
        {"role": "assistant", "content": reply},
    ]
    return (
        history, "", new_settings, new_settings,
        new_settings["categories"], ", ".join(new_settings["custom_keywords"]),
        ", ".join(new_settings["custom_fields"]), new_settings["only_matches"],
    )


theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="slate",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)

with gr.Blocks(title="Document Intelligence — OCR & Structured Data Extraction", theme=theme, css=CUSTOM_CSS) as demo:
    gr.Markdown(
        "# Turn Documents Into Structured Business Data\n"
        "AI-assisted OCR built for construction, compliance, and finance "
        "teams. Describe what you need, upload your documents, and get "
        "clean, audit-ready data back — not just extracted text.",
        elem_id="hero",
    )
    gr.HTML(PIPELINE_HTML)

    settings_state = gr.State(dict(DEFAULT_SETTINGS))

    with gr.Group(elem_classes=["step-card"]):
        gr.Markdown('<div class="step-eyebrow">Step 1</div><div class="step-kicker">Tell us what you need</div>')
        gr.Markdown(
            '<div class="step-subtitle">Describe what to flag and what to pull as data, '
            "in your own words — or start from a template below.</div>"
        )
        with gr.Row(elem_id="quick-prompts"):
            quick_buttons = {label: gr.Button(label, size="sm") for label in QUICK_PROMPTS}

        # Evaluated at container startup - tells the operator at a glance
        # whether the HF_TOKEN secret actually reached this process.
        _token_detected = bool(os.environ.get("HF_TOKEN"))
        gr.Markdown(
            '<div class="sample-note">Space token status: '
            + ("<strong>HF_TOKEN detected</strong> - the assistant is ready to use."
               if _token_detected else
               "<strong>HF_TOKEN not detected.</strong> The secret is missing, misnamed "
               "(it must be exactly <code>HF_TOKEN</code>, saved as a Secret, not a Variable), "
               "or the Space hasn't been restarted since it was added. You can still paste "
               "a token below, or skip the assistant and use the manual settings in Step 2.")
            + "</div>"
        )

        with gr.Accordion("Use your own Hugging Face token (optional)", open=False):
            gr.Markdown("Only needed if this Space's `HF_TOKEN` secret isn't set, or you want to use your own.")
            hf_token_input = gr.Textbox(label="Hugging Face token", type="password", placeholder="hf_...", show_label=False)

        chatbot = gr.Chatbot(height=320, label="Extraction assistant")
        with gr.Row():
            chat_input = gr.Textbox(
                label="Message", scale=4, show_label=False, container=False,
                placeholder="e.g. Flag overdue invoices; pull reference number, gross weight, delivery date",
            )
            chat_send = gr.Button("Send", scale=1, variant="primary")
        settings_display = gr.JSON(label="Extraction plan", value=DEFAULT_SETTINGS)

    with gr.Group(elem_classes=["step-card"]):
        gr.Markdown('<div class="step-eyebrow">Step 2</div><div class="step-kicker">Upload your documents</div>')
        gr.Markdown('<div class="step-subtitle">Upload as many PDFs or images as you need reviewed at once.</div>')
        files_input = gr.Files(
            label="PDF or image files",
            file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"],
            file_count="multiple",
        )
        gr.Markdown(
            '<div class="sample-note"><strong>No documents handy?</strong> '
            "Load a sample below — it fills in matching extraction settings "
            "too, so you can go straight to Run. All samples are fictitious "
            "and watermarked as such.</div>"
        )
        with gr.Row(elem_id="sample-docs"):
            sample_buttons = {label: gr.Button(label, size="sm") for label in SAMPLE_DOCS}
        with gr.Accordion("Advanced: edit settings manually", open=False):
            categories_input = gr.CheckboxGroup(
                choices=CATEGORY_CHOICES,
                value=CATEGORY_CHOICES,
                label="Notification categories",
                info="Also scopes which standard fields (PO/SAP/invoice/cert) get extracted.",
            )
            custom_keywords_input = gr.Textbox(
                label="Custom notification keywords (comma-separated)",
                placeholder="e.g. re-inspection, warranty claim, hold payment",
                info="Phrases that just raise a flag - not extracted as values.",
            )
            custom_fields_input = gr.Textbox(
                label="Custom fields to extract as data (comma-separated labels)",
                placeholder="e.g. reference number, gross weight, delivery date",
                info="Labels looked up in the document; the value after each label is captured.",
            )
            only_matches_input = gr.Checkbox(
                value=True,
                label="Only export pages with a notification hit or extracted field",
            )

    with gr.Group(elem_classes=["step-card"]):
        gr.Markdown('<div class="step-eyebrow">Step 3</div><div class="step-kicker">Run &amp; review</div>')
        run_button = gr.Button("Run OCR & Extract", variant="primary", elem_id="run-button")
        status_output = gr.Textbox(label="Status", lines=2, interactive=False)
        csv_output = gr.Files(label="Download CSVs (structured fields, notifications, audit trail)")
        with gr.Tabs():
            with gr.Tab("Structured fields"):
                fields_table_output = gr.Dataframe(headers=FIELDS_COLUMNS, wrap=True)
            with gr.Tab("Notifications"):
                notifications_table_output = gr.Dataframe(headers=NOTIFICATIONS_COLUMNS, wrap=True)
            with gr.Tab("Audit trail"):
                run_metadata_output = gr.JSON(label="Exactly what settings produced this run")
            with gr.Tab("Raw OCR output (debug)"):
                text_output = gr.Textbox(label="Raw extracted text", lines=20, show_label=False)

    chat_outputs = [
        chatbot, chat_input, settings_state, settings_display,
        categories_input, custom_keywords_input, custom_fields_input, only_matches_input,
    ]
    chat_send.click(fn=handle_chat, inputs=[chatbot, chat_input, settings_state, hf_token_input], outputs=chat_outputs)
    chat_input.submit(fn=handle_chat, inputs=[chatbot, chat_input, settings_state, hf_token_input], outputs=chat_outputs)

    for label, prompt in QUICK_PROMPTS.items():
        quick_buttons[label].click(fn=lambda p=prompt: p, outputs=chat_input).then(
            fn=handle_chat, inputs=[chatbot, chat_input, settings_state, hf_token_input], outputs=chat_outputs
        )

    for label, cfg in SAMPLE_DOCS.items():
        sample_buttons[label].click(
            fn=lambda c=cfg: ([c["file"]], c["categories"], c["keywords"], c["fields"]),
            outputs=[files_input, categories_input, custom_keywords_input, custom_fields_input],
        )

    run_button.click(
        fn=process_files,
        inputs=[files_input, categories_input, custom_keywords_input, custom_fields_input, only_matches_input, chatbot],
        outputs=[status_output, fields_table_output, notifications_table_output, csv_output, text_output, run_metadata_output],
    )

if __name__ == "__main__":
    demo.launch(allowed_paths=[SAMPLE_DIR])