app/api/upload.py

"""Upload endpoints: single file, bulk multi-part, and ZIP archives."""

from __future__ import annotations

import logging
import tempfile
import uuid
from pathlib import Path

from typing import Annotated

from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
from sqlalchemy import delete, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.rate_limit import check_read
from app.config import settings
from app.db.database import get_db
from app.db.models import Document, IngestStatus, Report
from app.ingest.schedule import schedule_ingest
from app.ingest.zip_extract import extract_reference_documents
from app.services.photo_policy_corpus import invalidate_tenant_photo_policy_cache
from app.models.schemas import (
    BatchUploadItem,
    BulkUploadResponse,
    DocumentDeleteResponse,
    DocumentSurveyLevelResponse,
    DocumentSurveyLevelUpdate,
    UploadResponse,
)

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

_ALLOWED_SUFFIXES: frozenset[str] = frozenset({".docx", ".pdf"})
_ZIP_SUFFIX: str = ".zip"


async def _read_upload_limited(upload: UploadFile, max_bytes: int) -> bytes:
    data = await upload.read(max_bytes + 1)
    if len(data) > max_bytes:
        raise HTTPException(
            status_code=413,
            detail=f"File too large. Maximum allowed size is {max_bytes // (1024 * 1024)} MB.",
        )
    return data


def _save_document_record(
    db: AsyncSession,
    tenant_id: str,
    filename: str,
    contents: bytes,
    suffix: str,
    *,
    survey_level: int | None = None,
) -> tuple[str, Path]:
    doc_id = str(uuid.uuid4())
    dest_dir = settings.upload_dir / tenant_id
    dest_dir.mkdir(parents=True, exist_ok=True)
    dest_path = dest_dir / f"{doc_id}{suffix}"
    dest_path.write_bytes(contents)
    doc = Document(
        id=doc_id,
        tenant_id=tenant_id,
        filename=filename,
        file_path=str(dest_path),
        status=IngestStatus.pending,
        survey_level=survey_level,
    )
    db.add(doc)
    return doc_id, dest_path


def _parse_optional_survey_level(raw: str | None) -> int | None:
    """Parse multipart ``survey_level`` field: empty → ``None``, else 1–3."""
    if raw is None:
        return None
    s = str(raw).strip()
    if not s:
        return None
    try:
        v = int(s)
    except ValueError:
        raise HTTPException(
            status_code=422,
            detail="survey_level must be an integer 1, 2, or 3 when provided.",
        ) from None
    if v not in (1, 2, 3):
        raise HTTPException(
            status_code=422,
            detail="survey_level must be 1 (Condition Report), 2 (HomeBuyer), or 3 (Building Survey).",
        )
    return v


@router.post("/upload", response_model=UploadResponse, status_code=202)
async def upload_file(
    request: Request,
    file: UploadFile = File(...),
    tenant_id: str = Form(...),
    survey_level: Annotated[str | None, Form()] = None,
    db: AsyncSession = Depends(get_db),
) -> UploadResponse:
    """Accept a single document upload and schedule async ingestion."""
    suffix = Path(file.filename or "").suffix.lower()
    if suffix not in _ALLOWED_SUFFIXES:
        raise HTTPException(
            status_code=422,
            detail=f"Unsupported file type '{suffix}'. Allowed: {sorted(_ALLOWED_SUFFIXES)}",
        )

    contents = await _read_upload_limited(file, settings.max_single_upload_bytes)
    sl = _parse_optional_survey_level(survey_level)
    doc_id, dest_path = _save_document_record(
        db, tenant_id, file.filename or "unknown", contents, suffix, survey_level=sl
    )
    await db.flush()
    await db.commit()

    schedule_ingest(doc_id=doc_id, file_path=dest_path)
    # Invalidate style cache so the next generation re-learns from the new doc
    from app.cache import style_cache as _sc
    _sc.invalidate(tenant_id)
    logger.info("Queued ingestion doc=%s tenant=%s (style cache invalidated)", doc_id, tenant_id)

    return UploadResponse(
        document_id=doc_id,
        tenant_id=tenant_id,
        filename=file.filename or "unknown",
        status="pending",
        message="File accepted; ingestion queued.",
    )


@router.post("/upload/batch", response_model=BulkUploadResponse, status_code=202)
async def upload_batch(
    request: Request,
    files: list[UploadFile] = File(...),
    tenant_id: str = Form(...),
    survey_level: Annotated[str | None, Form()] = None,
    db: AsyncSession = Depends(get_db),
) -> BulkUploadResponse:
    """Accept many reference documents in one request (and/or ZIP archives).

    Each accepted ``.docx`` / ``.pdf`` is written to disk and inserted as its own
    ``documents`` row **before** the response returns, then ingestion is queued in
    the background (subject to ``max_concurrent_ingests``).

    Optional ``survey_level`` (``1``, ``2``, or ``3``) applies to **every** file in
    this request (use PATCH ``/documents/{id}/survey-level`` to adjust individual
    rows after upload, or split mixed-tier libraries across multiple batch calls).

    ZIP files are expanded; only ``.docx`` and ``.pdf`` members are ingested.
    For very large libraries (millions of files), split work across multiple
    batch requests and/or several ZIPs — each stays within ``max_upload_batch_files``.
    """
    if not files:
        raise HTTPException(status_code=422, detail="No files uploaded")

    batch_sl = _parse_optional_survey_level(survey_level)

    items: list[BatchUploadItem] = []
    accepted = 0
    rejected = 0
    logical_count = 0
    cap = settings.max_upload_batch_files
    to_ingest: list[tuple[str, Path]] = []

    async def _accept_bytes(original_name: str, body: bytes, suffix: str) -> None:
        nonlocal accepted, rejected, logical_count
        if logical_count >= cap:
            rejected += 1
            items.append(
                BatchUploadItem(
                    filename=original_name,
                    status="rejected",
                    message=f"Batch limit reached ({cap} documents per request). Send another batch.",
                )
            )
            return
        doc_id, dest_path = _save_document_record(
            db, tenant_id, original_name, body, suffix, survey_level=batch_sl
        )
        await db.flush()
        to_ingest.append((doc_id, dest_path))
        accepted += 1
        logical_count += 1
        items.append(
            BatchUploadItem(
                document_id=doc_id,
                filename=original_name,
                status="accepted",
                message="Stored; ingestion queued.",
            )
        )

    for upload in files:
        name = upload.filename or "unknown"
        suffix = Path(name).suffix.lower()

        if suffix == _ZIP_SUFFIX:
            zmax = settings.max_archive_upload_bytes
            zbytes = await _read_upload_limited(upload, zmax)
            with tempfile.TemporaryDirectory() as tmp:
                zpath = Path(tmp) / "bundle.zip"
                zpath.write_bytes(zbytes)
                try:
                    extracted = extract_reference_documents(zpath, Path(tmp) / "out")
                except ValueError as exc:
                    rejected += 1
                    items.append(
                        BatchUploadItem(
                            filename=name,
                            status="rejected",
                            message=str(exc),
                        )
                    )
                    continue
                if not extracted:
                    rejected += 1
                    items.append(
                        BatchUploadItem(
                            filename=name,
                            status="rejected",
                            message="ZIP contained no .docx or .pdf files.",
                        )
                    )
                    continue
                for inner_name, inner_path in extracted:
                    body = inner_path.read_bytes()
                    suf = inner_path.suffix.lower()
                    await _accept_bytes(inner_name, body, suf)
            continue

        if suffix not in _ALLOWED_SUFFIXES:
            rejected += 1
            items.append(
                BatchUploadItem(
                    filename=name,
                    status="rejected",
                    message=f"Unsupported type {suffix!r}; allowed .docx, .pdf, .zip",
                )
            )
            continue

        body = await _read_upload_limited(upload, settings.max_single_upload_bytes)
        await _accept_bytes(name, body, suffix)

    await db.commit()
    for doc_id, dest_path in to_ingest:
        schedule_ingest(doc_id=doc_id, file_path=dest_path)

    # Invalidate cached style profile so the next generation re-learns from new docs
    if accepted > 0:
        from app.cache import style_cache as _sc
        _sc.invalidate(tenant_id)

    msg = (
        f"Accepted {accepted} document(s), rejected {rejected}. "
        "Each accepted file has its own database row; ingestion runs in the background."
    )
    logger.info(
        "Batch upload tenant=%s accepted=%s rejected=%s", tenant_id, accepted, rejected
    )
    return BulkUploadResponse(
        tenant_id=tenant_id,
        accepted=accepted,
        rejected=rejected,
        items=items,
        message=msg,
    )


@router.patch(
    "/documents/{document_id}/survey-level",
    response_model=DocumentSurveyLevelResponse,
    summary="Set RICS survey product tier for an uploaded document",
)
async def patch_document_survey_level(
    document_id: str,
    body: DocumentSurveyLevelUpdate,
    request: Request,
    db: AsyncSession = Depends(get_db),
    _: None = Depends(check_read),
) -> DocumentSurveyLevelResponse:
    """Label whether this upload is Level 1, 2, or 3 so library RAG can filter by tier.

    Does **not** re-ingest automatically; chunk metadata still carries the prior
    stamp. Retrieval filtering uses the database ``survey_level`` column — safe
    without re-indexing.
    """
    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")
    doc.survey_level = int(body.survey_level)
    await db.commit()
    invalidate_tenant_photo_policy_cache(tenant_id)
    return DocumentSurveyLevelResponse(document_id=document_id, survey_level=doc.survey_level)


@router.delete(
    "/documents/{document_id}",
    response_model=DocumentDeleteResponse,
    summary="Remove a reference document from the tenant library",
)
async def delete_uploaded_document(
    document_id: str,
    request: Request,
    db: AsyncSession = Depends(get_db),
    _: None = Depends(check_read),
) -> DocumentDeleteResponse:
    """Delete an uploaded file, its DB row, and all vector-index chunks for this document.

    Blocked with HTTP 409 when any report still references the document (FK integrity).
    """
    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")

    cnt = await db.execute(
        select(func.count()).select_from(Report).where(Report.document_id == document_id)
    )
    if int(cnt.scalar_one() or 0) > 0:
        raise HTTPException(
            status_code=409,
            detail=(
                "This file is still linked to one or more reports. "
                "Finish or abandon those jobs first, or upload replacements under a new document."
            ),
        )

    try:
        from app.vectorstore.factory import get_vectorstore

        get_vectorstore().delete_document(document_id)
    except Exception as exc:  # noqa: BLE001
        logger.warning("Vector store delete failed for doc=%s: %s", document_id, exc)

    fp = Path(doc.file_path)
    try:
        if fp.is_file():
            fp.unlink()
    except OSError as exc:
        logger.warning("Could not remove file %s: %s", fp, exc)

    await db.execute(delete(Document).where(Document.id == document_id))
    await db.commit()
    from app.retrieval.semantic_cache import invalidate_semantic_cache_for_tenant

    await invalidate_semantic_cache_for_tenant(tenant_id)
    invalidate_tenant_photo_policy_cache(tenant_id)
    return DocumentDeleteResponse(
        document_id=document_id,
        detail="Document removed from disk, database, and search index.",
    )