src/api/routers/batch.py

"""POST /extract/batch, GET /extract/batch/{job_id} — async batch extraction.

Use case
--------
An enterprise integration that wants to submit 100 invoices at once. Instead
of 100 sequential POST /extract calls (~10 minutes wall time), the caller
submits one batch request, gets a job_id, and polls for results while the
server processes them with bounded concurrency.

Semantics
---------
- POST /extract/batch is multipart with N files + doc_type + optional model.
  Returns 202 Accepted, {job_id, status: "pending"} — nothing has run yet.
  Actual extraction happens in a background task.
- GET /extract/batch/{job_id} returns the current snapshot: overall status,
  per-item status, and per-item results/errors as they land. Poll every 1-3s.
- Concurrency is capped globally at BATCH_CONCURRENCY (5) via an asyncio
  semaphore — sized to stay well under OpenAI rate limits even if multiple
  jobs run simultaneously.

Not implemented (deliberate)
----------------------------
- No persistent storage (Redis/DB). Jobs live in-process; restart wipes them.
- No auth, no per-user quotas.
- No webhook completion callbacks (the caller polls). Streaming completion
  updates would be the natural next step — but SSE for job progress is
  covered by the /extract/stream endpoint for the single-doc case.
"""
from __future__ import annotations

import asyncio
from pathlib import Path

from fastapi import APIRouter, BackgroundTasks, Depends, File, Form, UploadFile
from fastapi.responses import JSONResponse

from src.api.batch_store import JobStore, get_job_store
from src.api.deps import (
    ALLOWED_EXTENSIONS,
    ALLOWED_MIME_TYPES,
    MAX_UPLOAD_BYTES,
    get_extractor,
)
from src.api.errors import (
    FileTooLarge,
    UnsupportedDocType,
    UnsupportedFileType,
)
from src.extractors.extractor import DocumentExtractor
from src.schemas.registry import get_schema
from src.utils.logging import logger

router = APIRouter(prefix="/extract", tags=["extract-batch"])


# --- POST /extract/batch --------------------------------------------------

@router.post("/batch", status_code=202)
async def create_batch(
    background: BackgroundTasks,
    files:    list[UploadFile] = File(..., description="List of documents to extract."),
    doc_type: str    = Form(..., description="Registered doc type (invoice|receipt|filing)."),
    model:    str | None = Form(None, description="Optional model override."),
    extractor: DocumentExtractor = Depends(get_extractor),
    store:     JobStore = Depends(get_job_store),
) -> JSONResponse:
    # Reject unknown doc types up front — one 400 for the whole batch.
    try:
        get_schema(doc_type)
    except KeyError as e:
        raise UnsupportedDocType(str(e), details={"doc_type": doc_type}) from e

    if not files:
        # Return an empty done job rather than 400 — the semantics are clear.
        job = await store.create_job([], doc_type=doc_type, model=model)
        return JSONResponse(
            status_code=202,
            content={"job_id": job.id, "status": "done", "progress": job.progress},
        )

    # Read + validate each file *before* enqueueing so a bad file fails fast.
    items: list[tuple[str, bytes]] = []
    for uf in files:
        fn  = uf.filename or "upload"
        ext = Path(fn).suffix.lower()
        if ext not in ALLOWED_EXTENSIONS:
            raise UnsupportedFileType(
                f"Extension {ext!r} is not supported for {fn!r}.",
                details={"filename": fn, "extension": ext},
            )
        if uf.content_type and uf.content_type not in ALLOWED_MIME_TYPES:
            raise UnsupportedFileType(
                f"MIME type {uf.content_type!r} not accepted for {fn!r}.",
                details={"content_type": uf.content_type, "filename": fn},
            )
        payload = await uf.read()
        if len(payload) > MAX_UPLOAD_BYTES:
            raise FileTooLarge(
                f"{fn!r} is {len(payload)} bytes; max is {MAX_UPLOAD_BYTES}.",
                details={"filename": fn, "size_bytes": len(payload)},
            )
        items.append((fn, payload))

    job = await store.create_job(items, doc_type=doc_type, model=model)
    logger.info(f"[batch] created job {job.id} with {len(items)} items ({doc_type}, model={model})")

    # Fan out into background tasks. FastAPI runs BackgroundTasks after the
    # response is returned — so the caller gets an immediate 202.
    for i, (fn, payload) in enumerate(items):
        background.add_task(
            _run_item, store, extractor, job.id, i, fn, payload, doc_type, model
        )

    return JSONResponse(
        status_code=202,
        content={
            "job_id":   job.id,
            "status":   job.status,
            "progress": job.progress,
        },
    )


# --- GET /extract/batch/{job_id} -----------------------------------------

@router.get("/batch/{job_id}")
async def get_batch(
    job_id: str,
    store:  JobStore = Depends(get_job_store),
) -> dict:
    job = await store.get(job_id)
    if job is None:
        # 404 as a plain JSON envelope — this happens after a restart wipes memory.
        return JSONResponse(
            status_code=404,
            content={"error": {"code": "job_not_found", "message": f"No job with id {job_id!r}."}},
        )
    return job.snapshot()


# --- Worker --------------------------------------------------------------

async def _run_item(
    store: JobStore,
    extractor: DocumentExtractor,
    job_id: str,
    index: int,
    filename: str,
    payload: bytes,
    doc_type: str,
    model: str | None,
) -> None:
    """Run a single item under the global concurrency semaphore.

    The semaphore caps ALL concurrent extractions across every in-flight job.
    OpenAI rate limits are per-org, so job-scoped limits wouldn\'t protect us.
    """
    async with store.semaphore:
        await store.mark_running(job_id, index)
        try:
            # DocumentExtractor.extract() is sync + CPU/IO-bound (openai SDK is
            # synchronous under .parse). Push it to a thread so the event loop
            # keeps servicing status polls + other batches.
            result, metrics = await asyncio.to_thread(
                extractor.extract,
                payload,
                filename=filename,
                doc_type=doc_type,
                model_override=model,
            )
            await store.set_result(
                job_id, index,
                result=result.model_dump(mode="json"),
                metrics=metrics.to_dict(),
            )
        except Exception as e:  # noqa: BLE001
            logger.warning(f"[batch] job={job_id} item={index} failed: {e}")
            await store.set_error(job_id, index, f"{type(e).__name__}: {e}")