Ten31-Portal/backend/ten31portal/routers/document_router.py

"""Document upload, listing, download, and deletion with per-account access control."""

from datetime import datetime, timedelta

from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from sqlmodel import Session, select

from ten31portal.audit import record_audit
from ten31portal.auth import (
    accessible_entity_ids, can_access_entity, get_current_user, household_user_ids,
)
from ten31portal.database import get_session
from ten31portal.models import (
    Document, DocumentCategory, Entity, User, UserRole,
)
from ten31portal.schemas import DocumentResponse
from ten31portal import storage

router = APIRouter(prefix="/api/documents", tags=["documents"])


def _can_upload(user: User, entity_id: int, session: Session) -> bool:
    """Internal admins upload anywhere; external fund admins upload for their entities."""
    if user.role in (UserRole.approver, UserRole.cfo, UserRole.operations):
        return True
    if user.role == UserRole.fund_administrator:
        return can_access_entity(user, entity_id, session)
    return False


def _can_view(user: User, doc: Document, session: Session) -> bool:
    allowed = accessible_entity_ids(user, session)
    if allowed is None:  # internal staff
        return True
    if doc.entity_id not in allowed:
        return False
    # Investors only see shared docs or those addressed to any of their linked names;
    # fund admins see all docs for the entity.
    if user.role == UserRole.investor:
        return (
            doc.investor_user_id is None
            or doc.investor_user_id in household_user_ids(user, session)
        )
    return True


@router.get("")
def list_documents(
    entity_id: int | None = None,
    investor_user_id: int | None = None,
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> list[DocumentResponse]:
    query = select(Document)
    if entity_id is not None:
        query = query.where(Document.entity_id == entity_id)
    if investor_user_id is not None:
        query = query.where(Document.investor_user_id == investor_user_id)
    rows = session.exec(query.order_by(Document.created_at.desc())).all()  # type: ignore[union-attr]
    visible = [d for d in rows if _can_view(user, d, session)]

    # Badge docs that arrived since the investor's previous visit. Badges show on the first
    # page-load after new documents arrive; the watermark advances at most once per 30 minutes
    # so rapid refetches don't rewrite the row. First ever visit (no watermark) badges nothing —
    # everything would be "new".
    seen_before = user.docs_seen_at if user.role == UserRole.investor else None
    if user.role == UserRole.investor:
        now = datetime.utcnow()
        if user.docs_seen_at is None or (now - user.docs_seen_at) > timedelta(minutes=30):
            user.docs_seen_at = now
            session.add(user)
            session.commit()

    return [
        DocumentResponse.model_validate(d, from_attributes=True).model_copy(
            update={"is_new": seen_before is not None and d.created_at > seen_before}
        )
        for d in visible
    ]


@router.post("", status_code=201)
def upload_document(
    entity_id: int = Form(...),
    category: DocumentCategory = Form(DocumentCategory.other),
    title: str | None = Form(None),
    investor_user_id: int | None = Form(None),
    file: UploadFile = File(...),
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> DocumentResponse:
    entity = session.get(Entity, entity_id)
    if entity is None:
        raise HTTPException(status_code=404, detail="Entity not found")
    if not _can_upload(user, entity_id, session):
        raise HTTPException(status_code=403, detail="Insufficient permissions")

    # If targeting a specific investor, they must have access to this entity.
    if investor_user_id is not None:
        target = session.get(User, investor_user_id)
        if target is None:
            raise HTTPException(status_code=404, detail="Investor not found")
        if not can_access_entity(target, entity_id, session):
            raise HTTPException(
                status_code=400,
                detail="That investor does not have access to this entity.",
            )

    try:
        storage_name, size = storage.save_upload(file)
    except storage.UploadTooLarge as exc:
        raise HTTPException(status_code=413, detail=str(exc))
    doc = Document(
        entity_id=entity_id,
        investor_user_id=investor_user_id,
        category=category,
        title=title or (file.filename or "Untitled"),
        original_filename=file.filename or storage_name,
        content_type=file.content_type or "application/octet-stream",
        size_bytes=size,
        storage_path=storage_name,
        uploaded_by=user.id,
    )
    session.add(doc)
    session.flush()
    record_audit(session, user.id, "upload", "document", doc.id, {
        "entity_id": entity_id,
        "investor_user_id": investor_user_id,
        "category": category.value,
        "filename": doc.original_filename,
    })
    session.commit()
    session.refresh(doc)
    return DocumentResponse.model_validate(doc, from_attributes=True)


@router.get("/{document_id}/download")
def download_document(
    document_id: int,
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> FileResponse:
    doc = session.get(Document, document_id)
    if doc is None or not _can_view(user, doc, session):
        raise HTTPException(status_code=404, detail="Document not found")
    try:
        path = storage.full_path(doc.storage_path)
    except ValueError:
        raise HTTPException(status_code=404, detail="Document not found")
    if not path.is_file():
        raise HTTPException(status_code=404, detail="File missing on disk")
    return FileResponse(
        path,
        media_type=doc.content_type,
        filename=doc.original_filename,
        # K-1s and capital statements must not linger in a shared-machine browser cache.
        headers={"Cache-Control": "no-store"},
    )


@router.delete("/{document_id}")
def delete_document(
    document_id: int,
    admin: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> dict[str, str]:
    doc = session.get(Document, document_id)
    if doc is None:
        raise HTTPException(status_code=404, detail="Document not found")
    # Same reach as upload: internal admins anywhere, an Administrator on their entities.
    if not _can_upload(admin, doc.entity_id, session):
        raise HTTPException(status_code=403, detail="Insufficient permissions")
    storage.delete_file(doc.storage_path)
    record_audit(session, admin.id, "delete", "document", document_id, {
        "filename": doc.original_filename,
    })
    session.delete(doc)
    session.commit()
    return {"status": "deleted"}