app.py

import json
import logging
import threading
from pathlib import Path

import gradio as gr
import uvicorn

import config
from core.extractor import extract_document
from core.structurer import structure_extraction
from core.database import save_document, get_all_documents, get_stats, init_db
from utils.helpers import validate_file, build_summary_display, format_file_size
from api.routes import api_app
from benchmark.compare import run_benchmark

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
)
logger = logging.getLogger(__name__)

init_db()


# ── Processing pipeline ───────────────────────────────────────────────────────

def process_document(file_obj, language_hint: str = "auto"):
    if file_obj is None:
        return (
            "❌ No file uploaded. Please upload a PDF or image.",
            "{}",
            "",
        )

    try:
        file_path = file_obj
        filename = Path(file_path).name
        ext = Path(file_path).suffix.lower()

        valid, msg = validate_file(file_path)
        if not valid:
            return (f"❌ Validation failed: {msg}", "{}", "")

        with open(file_path, "rb") as f:
            file_bytes = f.read()

        file_size = format_file_size(len(file_bytes))
        logger.info(f"Processing: {filename} ({file_size})")

        raw_result = extract_document(file_bytes, ext, filename)
        structured = structure_extraction(raw_result)
        saved = save_document(structured)

        save_status = "✅ Saved to database" if saved else "⚠️ Save failed"
        summary = build_summary_display(structured)
        summary += f"\n\n💾 **DB Status:** {save_status}"

        json_output = json.dumps(structured, ensure_ascii=False, indent=2)

        full_text = structured.get("content", {}).get("full_text", "")
        if not full_text.strip():
            text_preview = (
                "⚠️ No text extracted. "
                "Try uploading a clearer image or a text-based PDF."
            )
        else:
            preview = full_text[:2000]
            lang = structured["document_analysis"]["language"]
            text_preview = f"**Extracted Text** ({lang}):\n\n{preview}"
            if len(full_text) > 2000:
                text_preview += f"\n\n... [+{len(full_text)-2000} more characters]"

        return summary, json_output, text_preview

    except Exception as e:
        logger.error(f"Processing error: {e}", exc_info=True)
        return (
            f"❌ Processing failed: {str(e)}",
            "{}",
            "",
        )


def load_documents_table():
    docs = get_all_documents(limit=50)
    if not docs:
        return [["No documents yet", "", "", "", "", ""]]
    return [
        [
            d.get("document_id", ""),
            d.get("filename", ""),
            d.get("document_type", ""),
            d.get("language", ""),
            d.get("total_words", 0),
            f"{d.get('confidence', 0):.0%}",
        ]
        for d in docs
    ]


def load_stats():
    stats = get_stats()
    if not stats:
        return "No documents processed yet."
    lines = [
        "## 📊 Database Statistics",
        "",
        f"**Total Documents:** {stats.get('total_documents', 0)}",
        f"**Arabic Documents:** {stats.get('arabic_documents', 0)}",
        f"**English Documents:** {stats.get('english_documents', 0)}",
        f"**Avg Confidence:** {stats.get('average_confidence', 0):.2%}",
        f"**Total Words Processed:** {stats.get('total_words_processed', 0):,}",
        "",
        "**By Document Type:**",
    ]
    for doc_type, count in stats.get("by_document_type", {}).items():
        lines.append(f"  - {doc_type}: {count}")
    return "\n".join(lines)


def run_benchmark_ui():
    try:
        results = run_benchmark()
        s = results["summary"]
        return f"""
## 🏁 Benchmark Results

| Metric | Arabic | English |
|--------|--------|---------|
| Avg Entity Score | {s['arabic_avg_score']:.0%} | {s['english_avg_score']:.0%} |
| Avg Processing Time | {s['arabic_avg_time_ms']:.1f}ms | {s['english_avg_time_ms']:.1f}ms |
| Documents Tested | {s['arabic_doc_count']} | {s['english_doc_count']} |

*Mock benchmark using sample data. Upload real documents for live results.*
""".strip()
    except Exception as e:
        return f"❌ Benchmark error: {e}"


# ── Gradio UI ─────────────────────────────────────────────────────────────────

def create_ui():
    with gr.Blocks(
        title=config.APP_TITLE,
        theme=gr.themes.Soft(),
    ) as demo:

        gr.HTML("""
        <div style="background:linear-gradient(135deg,#1a1a2e,#16213e);
                    padding:24px;border-radius:12px;color:white;margin-bottom:16px;">
          <h1 style="margin:0;font-size:1.7em;">
            🔍 Arabic Document Intelligence Pipeline
          </h1>
          <p style="margin:8px 0 0;opacity:.8;">
            Upload Arabic / English PDF or image → Extract → Structure → Query
          </p>
          <p style="margin:4px 0 0;font-size:.82em;opacity:.55;">
            Tesseract OCR engine (Arabic + English) · SQLite storage · FastAPI
          </p>
        </div>
        """)

        with gr.Tabs():

            # ── Tab 1: Upload & Process ──────────────────────────────────────
            with gr.Tab("📤 Upload & Process"):
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Upload Document")
                        file_input = gr.File(
                            label="PDF or Image",
                            file_types=[".pdf", ".png", ".jpg", ".jpeg",
                                        ".tiff", ".bmp"],
                            type="filepath",
                        )
                        language_radio = gr.Radio(
                            choices=["auto", "ar", "en"],
                            value="auto",
                            label="Language Hint",
                            info="'auto' detects automatically",
                        )
                        process_btn = gr.Button(
                            "🚀 Process Document",
                            variant="primary",
                            size="lg",
                        )
                        gr.Markdown(
                            "**Formats:** PDF · PNG · JPG · TIFF · BMP  \n"
                            "**Max size:** 20 MB  \n"
                            "**Engine:** Tesseract OCR (ara+eng)"
                        )

                    with gr.Column(scale=2):
                        gr.Markdown("### Results")
                        summary_out = gr.Markdown(
                            value="Upload a document and click **Process Document**."
                        )
                        text_preview = gr.Markdown(value="")

                json_out = gr.Textbox(
                    label="Structured JSON Output",
                    lines=22,
                    max_lines=40,
                    show_copy_button=True,
                )

                process_btn.click(
                    fn=process_document,
                    inputs=[file_input, language_radio],
                    outputs=[summary_out, json_out, text_preview],
                )

            # ── Tab 2: History ───────────────────────────────────────────────
            with gr.Tab("📚 Document History"):
                gr.Markdown("### Processed Documents")
                refresh_btn = gr.Button("🔄 Refresh", variant="secondary")
                docs_table = gr.Dataframe(
                    headers=[
                        "Document ID", "Filename", "Type",
                        "Language", "Words", "Confidence"
                    ],
                    value=load_documents_table(),
                    interactive=False,
                    wrap=True,
                )
                refresh_btn.click(fn=load_documents_table, outputs=docs_table)

            # ── Tab 3: Statistics ────────────────────────────────────────────
            with gr.Tab("📊 Statistics"):
                gr.Markdown("### Database Statistics")
                stats_btn = gr.Button("🔄 Refresh Stats", variant="secondary")
                stats_out = gr.Markdown(value=load_stats())
                stats_btn.click(fn=load_stats, outputs=stats_out)

            # ── Tab 4: Benchmark ─────────────────────────────────────────────
            with gr.Tab("🏁 Benchmark"):
                gr.Markdown(
                    "### Arabic vs English Extraction Benchmark\n"
                    "Compares entity extraction on sample documents."
                )
                bench_btn = gr.Button("▶️ Run Benchmark", variant="primary")
                bench_out = gr.Markdown(
                    value="Click **Run Benchmark** to start."
                )
                bench_btn.click(fn=run_benchmark_ui, outputs=bench_out)

            # ── Tab 5: API ───────────────────────────────────────────────────
            with gr.Tab("🔌 API"):
                gr.Markdown("""
### FastAPI Endpoints

The REST API runs alongside this UI on port **7861**.

| Method | Endpoint | Description |
|--------|----------|-------------|
| GET | `/` | API info |
| GET | `/stats` | DB statistics |
| GET | `/documents` | List all documents |
| GET | `/documents/{id}` | Get by ID |
| GET | `/documents/{id}/json` | Full JSON |
| POST | `/search` | Search by text |

**Search example:**
```json
POST /search
{ "query": "فاتورة", "language": "arabic" }
```
                """)

            # ── Tab 6: About ─────────────────────────────────────────────────
            with gr.Tab("ℹ️ About"):
                gr.Markdown("""
## Arabic Document Intelligence Pipeline

### Architecture