backend/app/api/documents.py

"""
Documents API — Document Management
=====================================
GET  /api/documents        — list all indexed documents
GET  /api/documents/{id}   — document metadata
POST /api/documents/upload — upload and ingest a new PDF
"""

import logging
import json
import shutil
from pathlib import Path
from typing import Any, Optional
from fastapi import APIRouter, HTTPException, UploadFile, File, BackgroundTasks
from pydantic import BaseModel

logger = logging.getLogger(__name__)
router = APIRouter()

BASE_DIR      = Path(__file__).parent.parent.parent
DATA_DIR      = BASE_DIR / "data"
RAW_DIR       = DATA_DIR / "raw"
DOCS_META_FILE = DATA_DIR / "documents.json"

RAW_DIR.mkdir(parents=True, exist_ok=True)


class DocumentMeta(BaseModel):
    doc_id: str
    name: str
    doc_type: str
    period: str
    institution: Optional[str] = None
    total_pages: int
    status: str        # "indexing" | "indexed" | "error"
    filename: str
    chunks_indexed: Optional[int] = None
    tables_indexed: Optional[int] = None
    colpali_pages: Optional[int] = None
    pagemap_file: Optional[str] = None
    page_section_map: Optional[dict[str, str]] = None
    page_metadata_map: Optional[dict[str, Any]] = None
    retrieval_config: Optional[str] = None


def _load_docs() -> list[dict]:
    if DOCS_META_FILE.exists():
        with open(DOCS_META_FILE) as f:
            return json.load(f)
    # Seed with the existing PDF if present
    seed_pdf = BASE_DIR / "emiratesnbd_investor_presentation_2026_q1.pdf"
    if seed_pdf.exists():
        return [{
            "doc_id": "emiratesnbd_q1_2026",
            "name": "Emirates NBD Investor Presentation Q1 2026",
            "doc_type": "Investor Presentation",
            "period": "Q1 2026",
            "total_pages": 48,
            "status": "indexed",
            "filename": seed_pdf.name,
        }]
    return []


def _save_docs(docs: list[dict]):
    DOCS_META_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(DOCS_META_FILE, "w") as f:
        json.dump(docs, f, indent=2)


@router.get("/", response_model=list[DocumentMeta])
async def list_documents():
    return _load_docs()


@router.get("/{doc_id}", response_model=DocumentMeta)
async def get_document(doc_id: str):
    docs = _load_docs()
    for d in docs:
        if d["doc_id"] == doc_id:
            return d
    raise HTTPException(status_code=404, detail=f"Document '{doc_id}' not found")


@router.post("/upload")
async def upload_document(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
):
    """Upload a PDF and trigger background ingestion."""
    if not file.filename.endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are accepted")

    # Save raw file
    safe_name = file.filename.replace(" ", "_").lower()
    dest = RAW_DIR / safe_name
    with open(dest, "wb") as f:
        shutil.copyfileobj(file.file, f)

    doc_id = safe_name.replace(".pdf", "").replace("-", "_")
    docs = _load_docs()
    docs.append({
        "doc_id": doc_id,
        "name": file.filename.replace(".pdf", "").replace("_", " ").title(),
        "doc_type": "Financial Document",
        "period": "Unknown",
        "total_pages": 0,
        "status": "indexing",
        "filename": safe_name,
    })
    _save_docs(docs)

    # Trigger background ingestion
    background_tasks.add_task(_ingest_background, doc_id, str(dest))

    return {"doc_id": doc_id, "status": "indexing", "message": "Ingestion started"}


async def _ingest_background(doc_id: str, pdf_path: str):
    """Run full ingestion pipeline in background."""
    try:
        import sys
        sys.path.insert(0, str(BASE_DIR))
        from ingest import ingest_document
        await ingest_document(doc_id, pdf_path)

        docs = _load_docs()
        for d in docs:
            if d["doc_id"] == doc_id:
                d["status"] = "indexed"
        _save_docs(docs)

    except Exception as e:
        logger.error(f"Background ingestion failed for {doc_id}: {e}")
        docs = _load_docs()
        for d in docs:
            if d["doc_id"] == doc_id:
                d["status"] = "error"
        _save_docs(docs)