socaichat/backend/app/api/attachments.py

"""Attachment API endpoints -- upload, metadata, download, delete.

Files are stored in Azure Blob Storage under the path:
    attachments/{conversation_id}/{uuid}_{filename}

After upload, if the content type is parseable (PDF, images, etc.),
a task is dispatched to Service Bus for async processing.
"""

from __future__ import annotations

import logging
import uuid

from litestar import Response, delete, get, post
from litestar.datastructures import UploadFile
from litestar.enums import RequestEncodingType
from litestar.exceptions import NotFoundException
from litestar.params import Body

from app.schemas import AttachmentOut
from app.storage.blob import delete_blob, generate_sas_url, upload_blob
from app.store.postgres import Attachment, async_session_factory
from app.tasks.bus import send_task

logger = logging.getLogger(__name__)

# Content types eligible for async parsing via Service Bus
PARSEABLE_TYPES = {
    "application/pdf",
    "image/png",
    "image/jpeg",
    "image/webp",
    "image/gif",
    "text/csv",
    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
}

# 50 MB upload limit
MAX_UPLOAD_BYTES = 50 * 1024 * 1024


def _attachment_to_out(att: Attachment) -> AttachmentOut:
    """Convert an ORM Attachment to the response schema."""
    return AttachmentOut(
        id=att.id,
        filename=att.filename,
        content_type=att.content_type,
        blob_url=att.blob_url,
        size_bytes=att.size_bytes,
        conversation_id=att.conversation_id,
        message_id=att.message_id,
        created_at=att.created_at.isoformat() if att.created_at else "",
    )


@post("/api/attachments/upload")
async def upload_attachment(
    data: UploadFile = Body(media_type=RequestEncodingType.MULTI_PART),
    conversation_id: str | None = None,
    message_id: str | None = None,
) -> AttachmentOut:
    """POST /api/attachments/upload -- Upload a file to Blob Storage.

    Accepts multipart/form-data with a single file field named ``data``.
    Optional query parameters ``conversation_id`` and ``message_id``
    associate the attachment with a conversation/message.

    Returns the attachment metadata including the blob URL.
    """
    file_bytes = await data.read()
    file_size = len(file_bytes)

    if file_size > MAX_UPLOAD_BYTES:
        return Response(
            content={"detail": f"File too large. Max {MAX_UPLOAD_BYTES // (1024*1024)} MB."},
            status_code=413,
        )

    filename = data.filename or "untitled"
    content_type = data.content_type or "application/octet-stream"
    attachment_id = str(uuid.uuid4())

    # Build a unique blob path
    conv_segment = conversation_id or "_unlinked"
    blob_name = f"attachments/{conv_segment}/{attachment_id}_{filename}"

    # Upload to Azure Blob Storage
    blob_url = await upload_blob(
        data=file_bytes,
        filename=blob_name,
        content_type=content_type,
    )

    # Persist metadata in PostgreSQL
    attachment = Attachment(
        id=attachment_id,
        conversation_id=conversation_id,
        message_id=message_id,
        filename=filename,
        content_type=content_type,
        blob_url=blob_url,
        size_bytes=file_size,
    )
    async with async_session_factory() as session:
        session.add(attachment)
        await session.commit()
        await session.refresh(attachment)

    # Dispatch async parsing task if the file type is parseable
    if content_type in PARSEABLE_TYPES:
        try:
            await send_task(
                task_type="parse_attachment",
                payload={
                    "attachment_id": attachment_id,
                    "content_type": content_type,
                    "blob_name": blob_name,
                    "filename": filename,
                },
                conversation_id=conversation_id,
            )
            logger.info(
                "Dispatched parse_attachment task for %s (type=%s)",
                attachment_id,
                content_type,
            )
        except Exception:
            # Parsing is best-effort; upload itself is already persisted
            logger.warning(
                "Failed to dispatch parse_attachment task for %s",
                attachment_id,
                exc_info=True,
            )

    return _attachment_to_out(attachment)


@get("/api/attachments/{attachment_id:str}")
async def get_attachment(attachment_id: str) -> AttachmentOut:
    """GET /api/attachments/:id -- Get attachment metadata."""
    async with async_session_factory() as session:
        attachment = await session.get(Attachment, attachment_id)
    if attachment is None:
        raise NotFoundException(detail=f"Attachment {attachment_id} not found")
    return _attachment_to_out(attachment)


@get("/api/attachments/{attachment_id:str}/download")
async def download_attachment(attachment_id: str) -> Response:
    """GET /api/attachments/:id/download -- Redirect to a time-limited SAS URL.

    Generates a SAS token valid for 1 hour and returns a 302 redirect.
    """
    async with async_session_factory() as session:
        attachment = await session.get(Attachment, attachment_id)
    if attachment is None:
        raise NotFoundException(detail=f"Attachment {attachment_id} not found")

    # Extract blob name from the full URL by removing the base
    # blob_url looks like https://<account>.blob.core.windows.net/<container>/<blob_name>
    # We stored the blob_name as attachments/{conv}/{uuid}_{filename}
    # Re-derive it from the stored URL
    blob_url = attachment.blob_url
    # Find the blob name after the container segment
    container_marker = "/soc-files/"
    if container_marker in blob_url:
        blob_name = blob_url.split(container_marker, 1)[1]
    else:
        # Fallback: reconstruct from known pattern
        conv_segment = attachment.conversation_id or "_unlinked"
        blob_name = f"attachments/{conv_segment}/{attachment.id}_{attachment.filename}"

    sas_url = await generate_sas_url(blob_name, expiry_hours=1)

    return Response(
        content=None,
        status_code=302,
        headers={"Location": sas_url},
    )


@delete("/api/attachments/{attachment_id:str}", status_code=204)
async def delete_attachment(attachment_id: str) -> None:
    """DELETE /api/attachments/:id -- Delete attachment (blob + DB record)."""
    async with async_session_factory() as session:
        attachment = await session.get(Attachment, attachment_id)
        if attachment is None:
            raise NotFoundException(detail=f"Attachment {attachment_id} not found")

        # Extract blob name
        blob_url = attachment.blob_url
        container_marker = "/soc-files/"
        if container_marker in blob_url:
            blob_name = blob_url.split(container_marker, 1)[1]
        else:
            conv_segment = attachment.conversation_id or "_unlinked"
            blob_name = f"attachments/{conv_segment}/{attachment.id}_{attachment.filename}"

        # Delete from Blob Storage (best-effort; DB record deleted regardless)
        try:
            await delete_blob(blob_name)
        except Exception:
            logger.warning(
                "Failed to delete blob %s, proceeding with DB deletion",
                blob_name,
                exc_info=True,
            )

        # Delete from database
        await session.delete(attachment)
        await session.commit()