gexp/routers/upload.py

import asyncio
import json
import uuid
from datetime import date
from pathlib import Path

from fastapi import APIRouter, File, Form, Request, UploadFile
from fastapi.responses import RedirectResponse
from fastapi.templating import Jinja2Templates

from config import UPLOAD_DIR
from database.connection import get_db
from services.extractor import extract_document, get_available_providers, get_default_provider
from services.pdf_handler import pdf_to_images

router = APIRouter()
templates = Jinja2Templates(directory="templates")

ALLOWED_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"}
ALLOWED_PDF_EXTS = {".pdf"}
ALLOWED_EXTENSIONS = ALLOWED_IMAGE_EXTS | ALLOWED_PDF_EXTS


def _save_image(file_bytes: bytes, original_name: str) -> str:
    """Save image to uploads/{date}/{uuid}_{name} and return relative path."""
    today = date.today().isoformat()
    dest_dir = Path(UPLOAD_DIR) / today
    dest_dir.mkdir(parents=True, exist_ok=True)
    suffix = Path(original_name).suffix.lower() or ".jpg"
    filename = f"{uuid.uuid4().hex}{suffix}"
    dest = dest_dir / filename
    dest.write_bytes(file_bytes)
    return str(Path(today) / filename)


async def _extract_and_save(doc_id: int, image_path: str, provider: str = ""):
    """Background task: call extractor, parse result, update DB."""
    try:
        data = await extract_document(image_path, provider=provider)
        raw_json = json.dumps(data, ensure_ascii=False)

        with get_db() as conn:
            conn.execute(
                """UPDATE documents SET status='extracted',
                   raw_extraction_json=?,
                   request_number=?,
                   request_date=?,
                   applicant_name_raw=?,
                   request_purpose=?,
                   data_valid_until=?,
                   registry_office=?,
                   owns_properties=?,
                   declared_property_count=?,
                   page_info=?,
                   search_scope=?,
                   updated_at=CURRENT_TIMESTAMP
                   WHERE id=?""",
                (
                    raw_json,
                    data.get("request_number"),
                    data.get("request_date"),
                    data.get("applicant_name_raw"),
                    data.get("request_purpose"),
                    data.get("data_valid_until"),
                    data.get("registry_office"),
                    data.get("owns_properties"),
                    data.get("declared_property_count"),
                    data.get("page_info"),
                    data.get("search_scope"),
                    doc_id,
                ),
            )

            for i, prop in enumerate(data.get("properties", [])):
                conn.execute(
                    """INSERT INTO properties
                       (document_id, row_order, party_name, property_number,
                        section, block, real_estate_district, qaza, num_shares, ownership_type)
                       VALUES (?,?,?,?,?,?,?,?,?,?)""",
                    (
                        doc_id, i,
                        prop.get("party_name"),
                        prop.get("property_number"),
                        prop.get("section"),
                        prop.get("block"),
                        prop.get("real_estate_district"),
                        prop.get("qaza"),
                        prop.get("num_shares"),
                        prop.get("ownership_type"),
                    ),
                )

    except Exception as e:
        with get_db() as conn:
            conn.execute(
                """UPDATE documents SET status='error', extraction_error=?,
                   updated_at=CURRENT_TIMESTAMP WHERE id=?""",
                (str(e), doc_id),
            )


@router.get("/")
async def upload_page(request: Request):
    with get_db() as conn:
        stats = conn.execute(
            """SELECT
               COUNT(*) AS total,
               SUM(CASE WHEN status='confirmed' THEN 1 ELSE 0 END) AS confirmed,
               SUM(CASE WHEN status='extracted' THEN 1 ELSE 0 END) AS pending_review,
               SUM(CASE WHEN status='error' THEN 1 ELSE 0 END) AS errors
            FROM documents WHERE status != 'staged'"""
        ).fetchone()
    return templates.TemplateResponse(
        request, "index.html", {
            "stats": dict(stats) if stats else {},
            "providers": get_available_providers(),
            "default_provider": get_default_provider(),
        }
    )


@router.post("/upload")
async def upload_files(
    request: Request,
    files: list[UploadFile] = File(...),
    provider: str = Form(""),
):
    if not provider:
        provider = get_default_provider()

    doc_ids = []

    for upload in files:
        suffix = Path(upload.filename).suffix.lower()
        if suffix not in ALLOWED_EXTENSIONS:
            continue

        file_bytes = await upload.read()

        if suffix in ALLOWED_PDF_EXTS:
            # PDF: split into per-page images
            pages = pdf_to_images(file_bytes, upload.filename)
            for page_info in pages:
                with get_db() as conn:
                    cursor = conn.execute(
                        """INSERT INTO documents
                           (image_path, status, provider, pdf_group_id, page_number)
                           VALUES (?, 'pending', ?, ?, ?)""",
                        (
                            page_info["image_path"],
                            provider,
                            page_info["pdf_group_id"],
                            page_info["page_number"],
                        ),
                    )
                    doc_ids.append((cursor.lastrowid, page_info["image_path"]))
        else:
            # Image file
            rel_path = _save_image(file_bytes, upload.filename)
            with get_db() as conn:
                cursor = conn.execute(
                    "INSERT INTO documents (image_path, status, provider) VALUES (?, 'pending', ?)",
                    (rel_path, provider),
                )
                doc_ids.append((cursor.lastrowid, rel_path))

    # Fire background extractions
    for doc_id, rel_path in doc_ids:
        asyncio.create_task(_extract_and_save(doc_id, rel_path, provider))

    if len(doc_ids) == 1:
        return RedirectResponse(f"/review/{doc_ids[0][0]}?wait=1", status_code=303)
    return RedirectResponse("/documents?uploaded=1", status_code=303)


# ─── Two-step workflow: stage images, then process ─────────────

@router.post("/upload/stage")
async def stage_file(
    request: Request,
    files: list[UploadFile] = File(...),
):
    """Save uploaded images without triggering AI extraction."""
    staged = []
    for upload in files:
        suffix = Path(upload.filename).suffix.lower()
        if suffix not in ALLOWED_EXTENSIONS:
            continue

        file_bytes = await upload.read()

        if suffix in ALLOWED_PDF_EXTS:
            pages = pdf_to_images(file_bytes, upload.filename)
            for page_info in pages:
                with get_db() as conn:
                    cursor = conn.execute(
                        """INSERT INTO documents
                           (image_path, status, pdf_group_id, page_number)
                           VALUES (?, 'staged', ?, ?)""",
                        (
                            page_info["image_path"],
                            page_info["pdf_group_id"],
                            page_info["page_number"],
                        ),
                    )
                    staged.append({
                        "id": cursor.lastrowid,
                        "image_path": page_info["image_path"],
                        "name": f"{upload.filename} (p{page_info['page_number']})",
                    })
        else:
            rel_path = _save_image(file_bytes, upload.filename)
            with get_db() as conn:
                cursor = conn.execute(
                    "INSERT INTO documents (image_path, status) VALUES (?, 'staged')",
                    (rel_path,),
                )
                staged.append({
                    "id": cursor.lastrowid,
                    "image_path": rel_path,
                    "name": upload.filename,
                })

    from fastapi.responses import JSONResponse
    return JSONResponse({"staged": staged})


@router.post("/upload/process-staged")
async def process_staged(
    request: Request,
    doc_ids: str = Form(...),
    provider: str = Form(""),
):
    """Trigger AI extraction for previously staged documents."""
    if not provider:
        provider = get_default_provider()

    ids = [int(x) for x in doc_ids.split(",") if x.strip().isdigit()]

    with get_db() as conn:
        rows = conn.execute(
            f"SELECT id, image_path FROM documents WHERE id IN ({','.join('?' * len(ids))}) AND status='staged'",
            ids,
        ).fetchall()
        for row in rows:
            conn.execute(
                "UPDATE documents SET status='pending', provider=?, updated_at=CURRENT_TIMESTAMP WHERE id=?",
                (provider, row["id"]),
            )

    for row in rows:
        asyncio.create_task(_extract_and_save(row["id"], row["image_path"], provider))

    if len(rows) == 1:
        return RedirectResponse(f"/review/{rows[0]['id']}?wait=1", status_code=303)
    return RedirectResponse("/documents?uploaded=1", status_code=303)


@router.delete("/upload/staged/{doc_id}")
async def remove_staged(doc_id: int):
    """Remove a single staged document before processing."""
    from fastapi.responses import JSONResponse
    with get_db() as conn:
        row = conn.execute("SELECT image_path FROM documents WHERE id=? AND status='staged'", (doc_id,)).fetchone()
        if row:
            # Delete the file
            file_path = Path(UPLOAD_DIR) / row["image_path"]
            if file_path.exists():
                file_path.unlink()
            conn.execute("DELETE FROM documents WHERE id=?", (doc_id,))
            return JSONResponse({"ok": True})
    return JSONResponse({"ok": False}, status_code=404)