app/api/status.py

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

from collections import Counter
from datetime import UTC, datetime
from pathlib import Path
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.services.generation_progress import build_status_progress_fields
from app.models.schemas import (
    DocumentBatchStatusRequest,
    DocumentBatchStatusResponse,
    DocumentStatusItem,
    ReportStatusResponse,
)

router = APIRouter()


def _iso_utc(dt: datetime | None) -> str | None:
    """Serialize a DB timestamp as UTC ISO-8601 with a ``Z`` suffix.

    SQLite returns naive datetimes; we treat those as UTC because ORM
    defaults use ``datetime.now(UTC)``. Without ``Z``, browsers parse the
    value as local time and display the wrong clock time.
    """
    if dt is None:
        return None
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=UTC)
    else:
        dt = dt.astimezone(UTC)
    return dt.isoformat().replace("+00:00", "Z")


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": _iso_utc(doc.created_at),
        "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()

    def _file_size(path_str: str) -> int | None:
        try:
            p = Path(path_str)
            return p.stat().st_size if p.is_file() else None
        except OSError:
            return None

    return {
        "tenant_id": tenant_id,
        "limit": limit,
        "offset": offset,
        "documents": [
            {
                "document_id": d.id,
                "filename": d.filename,
                "status": d.status.value,
                "created_at": _iso_utc(d.created_at),
                "updated_at": _iso_utc(d.updated_at),
                "survey_level": d.survey_level,
                "document_purpose": (
                    d.document_purpose.value
                    if hasattr(d.document_purpose, "value")
                    else d.document_purpose
                ),
                "file_size": _file_size(d.file_path),
                "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()}

    from app.ingest.schedule import schedule_ingest
    from app.vectorstore.factory import get_vectorstore

    vs = get_vectorstore()
    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"
        if st == "complete" and vs.count_for_doc(doc.id) == 0:
            upload_path = Path(doc.file_path)
            if upload_path.is_file():
                doc.status = IngestStatus.pending
                doc.error_message = None
                schedule_ingest(doc_id=doc.id, file_path=upload_path)
                mutated = True
                st = "pending"
            elif not doc.error_message:
                doc.status = IngestStatus.failed
                doc.error_message = f"Upload file not found: {upload_path}"
                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")

    progress = await build_status_progress_fields(db, report)
    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,
        **progress,
    )