app/api/status.py

"""Status polling endpoint for documents and reports."""

from collections import Counter
from datetime import UTC, datetime
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, Query, Request
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.db.database import get_db
from app.db.models import Document, IngestStatus, Report
from app.models.schemas import (
    DocumentBatchStatusRequest,
    DocumentBatchStatusResponse,
    DocumentStatusItem,
    ReportStatusResponse,
)

router = APIRouter()


def _utc_age_seconds(now_utc: datetime, then: datetime | None) -> float:
    """Compute age in seconds, tolerant of naive DB datetimes.

    SQLite commonly returns naive datetimes. We treat naive values as UTC.
    """
    if then is None:
        return 0.0
    if then.tzinfo is None:
        then = then.replace(tzinfo=UTC)
    return (now_utc - then).total_seconds()


@router.get("/documents/{document_id}/info")
async def document_info(
    document_id: str,
    request: Request,
    db: AsyncSession = Depends(get_db),
) -> dict[str, Any]:
    """Return document metadata including ingested chunk count.

    Args:
        document_id: UUID of the document.
        request: Provides ``state.tenant_id``.
        db: Injected database session.

    Returns:
        Dict with document metadata and chunk count.
    """
    tenant_id: str = request.state.tenant_id
    doc = await db.get(Document, document_id)
    if doc is None or doc.tenant_id != tenant_id:
        raise HTTPException(status_code=404, detail="Document not found")

    from app.vectorstore.factory import get_vectorstore
    try:
        vs = get_vectorstore()
        chunk_count = vs.count_for_doc(document_id)
    except Exception:
        chunk_count = 0

    return {
        "document_id": document_id,
        "filename": doc.filename,
        "status": doc.status.value,
        "chunk_count": chunk_count,
        "created_at": doc.created_at.isoformat(),
        "survey_level": doc.survey_level,
    }


@router.get("/documents/{document_id}/status")
async def document_status(
    document_id: str,
    request: Request,
    db: AsyncSession = Depends(get_db),
) -> dict[str, str]:
    """Return ingestion status of an uploaded document.

    Args:
        document_id: UUID of the document.
        request: Provides ``state.tenant_id``.
        db: Injected database session.

    Returns:
        Dict with ``document_id``, ``status``, and optional ``error``.

    Raises:
        HTTPException: 404 if not found or wrong tenant.
    """
    tenant_id: str = request.state.tenant_id
    doc = await db.get(Document, document_id)
    if doc is None or doc.tenant_id != tenant_id:
        raise HTTPException(status_code=404, detail="Document not found")

    payload: dict[str, str] = {
        "document_id": document_id,
        "status": doc.status.value,
    }
    if doc.error_message:
        payload["error"] = doc.error_message
    return payload


@router.get("/documents")
async def list_documents(
    request: Request,
    db: AsyncSession = Depends(get_db),
    limit: int = Query(default=100, ge=1),
    offset: int = Query(default=0, ge=0),
) -> dict[str, Any]:
    """List uploaded documents for the tenant (newest first).

    Use pagination for large libraries; increase ``max_upload_batch_files`` /
    ``documents_list_max_limit`` via settings when you need higher throughput.
    """
    tenant_id: str = request.state.tenant_id
    cap = settings.documents_list_max_limit
    limit = min(limit, cap)

    result = await db.execute(
        select(Document)
        .where(Document.tenant_id == tenant_id)
        .order_by(Document.created_at.desc())
        .limit(limit)
        .offset(offset)
    )
    rows = result.scalars().all()
    return {
        "tenant_id": tenant_id,
        "limit": limit,
        "offset": offset,
        "documents": [
            {
                "document_id": d.id,
                "filename": d.filename,
                "status": d.status.value,
                "created_at": d.created_at.isoformat(),
                "survey_level": d.survey_level,
                "error": d.error_message,
            }
            for d in rows
        ],
    }


@router.post("/documents/batch-status", response_model=DocumentBatchStatusResponse)
async def documents_batch_status(
    body: DocumentBatchStatusRequest,
    request: Request,
    db: AsyncSession = Depends(get_db),
) -> DocumentBatchStatusResponse:
    """Return ingestion status for many document IDs in one request."""
    tenant_id: str = request.state.tenant_id
    if not body.document_ids:
        return DocumentBatchStatusResponse(
            items=[],
            pending=0,
            processing=0,
            complete=0,
            failed=0,
        )

    result = await db.execute(
        select(Document).where(
            Document.tenant_id == tenant_id,
            Document.id.in_(body.document_ids),
        )
    )
    found = {d.id: d for d in result.scalars().all()}

    items: list[DocumentStatusItem] = []
    counts: Counter[str] = Counter()
    now = datetime.now(UTC)
    stale_cutoff_s = int(settings.ingest_timeout_seconds)
    mutated = False
    for did in body.document_ids:
        doc = found.get(did)
        if doc is None:
            items.append(
                DocumentStatusItem(
                    document_id=did,
                    status="not_found",
                    filename="",
                    error="Document not found or not owned by this tenant",
                )
            )
            continue
        st = doc.status.value
        if st == "processing":
            # If a worker gets stuck (network hang, loader deadlock), do not block
            # the UI forever: mark as failed after a safe timeout.
            age_s = _utc_age_seconds(now, doc.updated_at)
            if age_s > stale_cutoff_s:
                doc.status = IngestStatus.failed
                doc.error_message = (
                    f"Ingestion timed out after {stale_cutoff_s}s. "
                    "The file may be too large/corrupted, or embedding/indexing may be unavailable."
                )
                mutated = True
                st = "failed"
        # Treat unknown / not_found as failed for aggregates so the UI never
        # shows "0 failed" while every row is a terminal error (e.g. tenant mismatch).
        if st == "not_found":
            counts["failed"] += 1
        elif st in ("pending", "processing", "complete", "failed"):
            counts[st] += 1
        items.append(
            DocumentStatusItem(
                document_id=did,
                status=st,
                filename=doc.filename,
                error=doc.error_message,
            )
        )

    if mutated:
        await db.commit()

    return DocumentBatchStatusResponse(
        items=items,
        pending=counts["pending"],
        processing=counts["processing"],
        complete=counts["complete"],
        failed=counts["failed"],
    )


@router.get("/documents/tenant-chunk-summary")
async def tenant_chunk_summary(request: Request) -> dict[str, str | int]:
    """Return total vector-store chunks for this tenant (all reference documents)."""
    tenant_id: str = request.state.tenant_id
    from app.vectorstore.factory import get_vectorstore

    try:
        vs = get_vectorstore()
        n = vs.count(tenant_id)
    except Exception:
        n = 0
    return {"tenant_id": tenant_id, "indexed_chunk_count": n}


@router.get("/reports/{report_id}/status", response_model=ReportStatusResponse)
async def report_status(
    report_id: str,
    request: Request,
    db: AsyncSession = Depends(get_db),
) -> ReportStatusResponse:
    """Return generation status of a report.

    Args:
        report_id: UUID of the report.
        request: Provides ``state.tenant_id``.
        db: Injected database session.

    Returns:
        :class:`~app.models.schemas.ReportStatusResponse`.

    Raises:
        HTTPException: 404 if not found or wrong tenant.
    """
    tenant_id: str = request.state.tenant_id
    report = await db.get(Report, report_id)
    if report is None or report.tenant_id != tenant_id:
        raise HTTPException(status_code=404, detail="Report not found")

    return ReportStatusResponse(
        report_id=report_id,
        status=report.status.value,
        created_at=report.created_at,
        updated_at=report.updated_at,
        error_message=report.error_message,
        survey_level=report.survey_level,
    )