webmail-client/backend/app/api/compose.py

"""Compose endpoints: drafts + send + undo (SPEC §9.1)."""

from __future__ import annotations

import asyncio
import contextlib
import uuid
from datetime import UTC, datetime, timedelta
from typing import Annotated

from fastapi import APIRouter, Depends, File, Query, UploadFile
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.errors import raise_api_error
from app.api.schemas import AttachmentOut
from app.config import get_settings
from app.db import get_session
from app.identity.deps import CurrentUser, require_csrf
from app.indexer import attachments as att_store
from app.indexer import blobstore, rawstore
from app.models import Attachment, BodyPart, ComposeAttachment, Draft, Mailbox, Message, Outbox
from app.outbox.jobs import RETRY_WINDOW
from app.outbox.schemas import (
    DraftIn,
    DraftOut,
    OutboxItemOut,
    RemoteDraftOut,
    SendRequest,
    SendResponse,
)
from app.outbox.service import (
    DraftError,
    cancel_outbox,
    create_outbox,
    delete_mirrored_draft_row,
    not_being_sent,
    upsert_draft,
)
from app.queue import enqueue_job

router = APIRouter(tags=["compose"], dependencies=[Depends(require_csrf)])

DbSession = Annotated[AsyncSession, Depends(get_session)]

MAX_ATTACHMENT_BYTES = 25 * 1024 * 1024  # 25 MB per file (compose Attach limit)


def _att_out(att: ComposeAttachment) -> AttachmentOut:
    return AttachmentOut(
        id=att.id, filename=att.filename, mime_type=att.mime_type, size_bytes=att.size_bytes
    )


@router.post("/attachments", response_model=AttachmentOut)
async def upload_attachment(
    user: CurrentUser, db: DbSession, file: UploadFile = File(...)
) -> AttachmentOut:
    """Store a file for an outgoing message (≤ 25 MB) and return its metadata."""
    data = await file.read()
    if not data:
        raise_api_error(400, "empty_file", "the file is empty")
    if len(data) > MAX_ATTACHMENT_BYTES:
        raise_api_error(413, "too_large", "attachment exceeds the 25 MB limit")
    sha = blobstore.store(user.id, data)
    att = ComposeAttachment(
        user_id=user.id,
        filename=file.filename or "attachment",
        mime_type=file.content_type or "application/octet-stream",
        size_bytes=len(data),
        blob_sha256=sha,
    )
    db.add(att)
    await db.commit()
    return _att_out(att)


@router.get("/attachments", response_model=list[AttachmentOut])
async def list_compose_attachments(
    user: CurrentUser, db: DbSession, ids: str = Query("")
) -> list[AttachmentOut]:
    """Resolve metadata for the given compose-attachment ids (reopening a draft)."""
    id_list: list[uuid.UUID] = []
    for value in ids.split(","):
        value = value.strip()
        if not value:
            continue
        try:
            id_list.append(uuid.UUID(value))
        except ValueError:
            continue
    if not id_list:
        return []
    rows = await db.scalars(
        select(ComposeAttachment).where(
            ComposeAttachment.user_id == user.id, ComposeAttachment.id.in_(id_list)
        )
    )
    by_id = {r.id: r for r in rows}
    return [_att_out(by_id[i]) for i in id_list if i in by_id]


def _draft_out(draft: Draft) -> DraftOut:
    return DraftOut(
        id=draft.id,
        account_id=draft.account_id,
        thread_id=draft.thread_id,
        reply_to_message_id=draft.reply_to_message_id,
        kind=draft.kind,
        payload=draft.payload,
        updated_at=draft.updated_at.isoformat(),
    )


@router.get("/drafts", response_model=list[DraftOut])
async def list_drafts(user: CurrentUser, db: DbSession) -> list[DraftOut]:
    drafts = await db.scalars(
        select(Draft)
        .where(Draft.user_id == user.id, not_being_sent())
        .order_by(Draft.updated_at.desc())
    )
    return [_draft_out(d) for d in drafts]


def _fmt_addr(addr: dict[str, str]) -> str:
    email = addr.get("email", "")
    name = addr.get("name", "")
    return f"{name} <{email}>" if name else email


async def _own_draft_uids(db: AsyncSession, user_id: uuid.UUID) -> set[tuple[uuid.UUID, int]]:
    """(account_id, imap_uid) of drafts we mirrored — excluded from the remote list."""
    rows = await db.execute(
        select(Draft.account_id, Draft.imap_uid).where(
            Draft.user_id == user_id, Draft.imap_uid.isnot(None)
        )
    )
    return {(r.account_id, r.imap_uid) for r in rows}


@router.get("/drafts/remote", response_model=list[RemoteDraftOut])
async def list_remote_drafts(user: CurrentUser, db: DbSession) -> list[RemoteDraftOut]:
    """Drafts sitting in the IMAP Drafts folder that we haven't already adopted."""
    own = await _own_draft_uids(db, user.id)
    messages = await db.scalars(
        select(Message)
        .join(Mailbox, Mailbox.id == Message.mailbox_id)
        .where(Message.user_id == user.id, Mailbox.special_use == "drafts")
        .order_by(Message.internal_date.desc())
    )
    out: list[RemoteDraftOut] = []
    for m in messages:
        if (m.account_id, m.uid) in own:
            continue
        out.append(
            RemoteDraftOut(
                message_id=m.id,
                account_id=m.account_id,
                imap_uid=m.uid,
                subject=m.subject,
                to=[_fmt_addr(a) for a in m.to_addr],
                cc=[_fmt_addr(a) for a in m.cc_addr],
                bcc=[_fmt_addr(a) for a in m.bcc_addr],
                snippet=m.snippet,
                updated_at=(m.date_hdr or m.internal_date).isoformat(),
            )
        )
    return out


@router.post("/drafts/adopt/{message_id}", response_model=DraftOut)
async def adopt_remote_draft(message_id: uuid.UUID, user: CurrentUser, db: DbSession) -> DraftOut:
    """Turn a Drafts-folder message into an editable local draft (linked by UID)."""
    msg = await db.scalar(
        select(Message)
        .join(Mailbox, Mailbox.id == Message.mailbox_id)
        .where(
            Message.id == message_id,
            Message.user_id == user.id,
            Mailbox.special_use == "drafts",
        )
    )
    if msg is None:
        raise_api_error(404, "not_found", "draft message not found")

    # Idempotent: if already adopted (a local draft links this UID), return it.
    existing = await db.scalar(
        select(Draft).where(
            Draft.user_id == user.id,
            Draft.account_id == msg.account_id,
            Draft.imap_uid == msg.uid,
        )
    )
    if existing is not None:
        return _draft_out(existing)

    parts = list(await db.scalars(select(BodyPart).where(BodyPart.message_id == msg.id)))
    html = next((p.content for p in parts if p.kind == "html_sanitized"), "")
    text = next((p.content for p in parts if p.kind == "text"), "")

    # Carry the draft's real (non-inline) attachments over as compose attachments,
    # so they show in the editor and are re-attached on save/send. The underlying
    # blobs are shared by sha256 — we only add new ComposeAttachment references.
    att_rows = list(
        await db.scalars(
            select(Attachment)
            .where(Attachment.message_id == msg.id, Attachment.content_id.is_(None))
            .order_by(Attachment.filename)
        )
    )
    attachment_ids: list[str] = []
    if att_rows:
        with contextlib.suppress(FileNotFoundError):  # fetched if a hosted mailbox's
            await rawstore.ensure(msg)
    for a in att_rows:
        # Kept only inside the message: give the compose attachment a blob of its own.
        if not await asyncio.to_thread(
            att_store.materialize, user.id, msg.blob_sha256, a.blob_sha256, a.part_index
        ):
            continue
        ca = ComposeAttachment(
            user_id=user.id,
            filename=a.filename,
            mime_type=a.mime_type,
            size_bytes=a.size_bytes,
            blob_sha256=a.blob_sha256,
        )
        db.add(ca)
        await db.flush()
        attachment_ids.append(str(ca.id))

    payload = {
        "to": [_fmt_addr(a) for a in msg.to_addr],
        "cc": [_fmt_addr(a) for a in msg.cc_addr],
        "bcc": [_fmt_addr(a) for a in msg.bcc_addr],
        "subject": msg.subject,
        "html": html,
        "text": text,
        "attachment_ids": attachment_ids,
    }
    draft = Draft(
        user_id=user.id,
        account_id=msg.account_id,
        kind="new",
        payload=payload,
        imap_uid=msg.uid,
    )
    db.add(draft)
    await db.commit()
    return _draft_out(draft)


@router.post("/drafts", response_model=DraftOut)
async def create_draft(body: DraftIn, user: CurrentUser, db: DbSession) -> DraftOut:
    try:
        draft = await upsert_draft(db, user_id=user.id, body=body, draft_id=None)
    except DraftError as exc:
        raise_api_error(400, "draft_error", str(exc))
    await db.commit()
    return _draft_out(draft)


@router.put("/drafts/{draft_id}", response_model=DraftOut)
async def update_draft(
    draft_id: uuid.UUID, body: DraftIn, user: CurrentUser, db: DbSession
) -> DraftOut:
    try:
        draft = await upsert_draft(db, user_id=user.id, body=body, draft_id=draft_id)
    except DraftError as exc:
        code = 404 if "not found" in str(exc) else 400
        raise_api_error(code, "draft_error", str(exc))
    await db.commit()
    return _draft_out(draft)


@router.post("/drafts/{draft_id}/publish")
async def publish_draft(draft_id: uuid.UUID, user: CurrentUser, db: DbSession) -> dict[str, str]:
    """Mirror the draft to the account's IMAP Drafts folder (continue on any client)."""
    draft = await db.scalar(select(Draft).where(Draft.id == draft_id, Draft.user_id == user.id))
    if draft is None:
        raise_api_error(404, "not_found", "draft not found")
    await enqueue_job("push_draft_to_imap_job", str(draft_id))
    return {"status": "queued"}


@router.delete("/drafts/{draft_id}")
async def delete_draft(draft_id: uuid.UUID, user: CurrentUser, db: DbSession) -> dict[str, str]:
    draft = await db.scalar(select(Draft).where(Draft.id == draft_id, Draft.user_id == user.id))
    if draft is None:
        raise_api_error(404, "not_found", "draft not found")
    account_id, imap_uid = draft.account_id, draft.imap_uid
    await db.delete(draft)
    # Drop the mirrored Drafts-folder message row now so the draft doesn't
    # resurface as an "other client" remote draft after the local row is gone.
    if imap_uid:
        await delete_mirrored_draft_row(db, account_id=account_id, imap_uid=imap_uid)
    await db.commit()
    # Remove the mirrored copy from the server Drafts folder too.
    if imap_uid:
        await enqueue_job("remove_imap_draft_job", str(account_id), imap_uid)
    return {"status": "deleted"}


@router.post("/send", response_model=SendResponse)
async def send(body: SendRequest, user: CurrentUser, db: DbSession) -> SendResponse:
    draft = await db.scalar(
        select(Draft).where(Draft.id == body.draft_id, Draft.user_id == user.id)
    )
    if draft is None:
        raise_api_error(404, "not_found", "draft not found")
    outbox = await create_outbox(db, user_id=user.id, draft=draft)
    await db.commit()

    await enqueue_job(
        "send_job",
        str(outbox.id),
        _defer_by=timedelta(seconds=get_settings().undo_seconds),
    )
    left = max(0.0, (outbox.send_after - datetime.now(UTC)).total_seconds())
    return SendResponse(
        outbox_id=outbox.id, undo_until=outbox.send_after.isoformat(), undo_seconds=left
    )


@router.post("/send/{outbox_id}/undo")
async def undo_send(outbox_id: uuid.UUID, user: CurrentUser, db: DbSession) -> dict[str, str]:
    ok = await cancel_outbox(db, user_id=user.id, outbox_id=outbox_id)
    await db.commit()
    if not ok:
        raise_api_error(409, "undo_too_late", "the message has already left the undo window")
    return {"status": "cancelled"}


# --- the Scheduled folder's server side: sends being retried, or failed ------------


@router.get("/outbox", response_model=list[OutboxItemOut])
async def list_outbox(user: CurrentUser, db: DbSession) -> list[OutboxItemOut]:
    """Sends that did not go out: retrying (the server could not take them yet)
    or failed for good. They wait in Scheduled, not in Drafts."""
    rows = await db.execute(
        select(Outbox, Draft)
        .join(Draft, Draft.id == Outbox.draft_id)
        .where(Outbox.user_id == user.id, Outbox.state.in_(("retrying", "failed")))
        .order_by(Outbox.send_after.desc())
    )
    out: list[OutboxItemOut] = []
    for ob, draft in rows.tuples():
        payload = draft.payload or {}
        to = payload.get("to")
        out.append(
            OutboxItemOut(
                outbox_id=ob.id,
                draft_id=draft.id,
                account_id=draft.account_id,
                state=ob.state,
                attempts=ob.attempts,
                error=ob.error,
                next_attempt_at=ob.next_attempt_at.isoformat() if ob.next_attempt_at else None,
                retry_until=ob.retry_until.isoformat() if ob.retry_until else None,
                sent_at=ob.send_after.isoformat(),
                to=[str(t) for t in to] if isinstance(to, list) else [],
                subject=str(payload.get("subject") or ""),
                preview=str(payload.get("text") or "")[:200],
            )
        )
    return out


async def _waiting_send(db: AsyncSession, user_id: uuid.UUID, outbox_id: uuid.UUID) -> Outbox:
    ob = await db.scalar(select(Outbox).where(Outbox.id == outbox_id, Outbox.user_id == user_id))
    if ob is None or ob.state not in ("retrying", "failed"):
        raise_api_error(409, "not_waiting", "this message is not waiting to be sent")
    return ob


@router.post("/outbox/{outbox_id}/retry")
async def retry_outbox(outbox_id: uuid.UUID, user: CurrentUser, db: DbSession) -> dict[str, str]:
    """Retry now. A send that had failed for good gets a fresh 24-hour window."""
    ob = await _waiting_send(db, user.id, outbox_id)
    now = datetime.now(UTC)
    if ob.state == "failed" or ob.retry_until is None or ob.retry_until <= now:
        ob.retry_until = now + RETRY_WINDOW
    ob.state = "retrying"
    ob.next_attempt_at = now
    await db.commit()
    await enqueue_job("send_job", str(ob.id), _job_id=f"send:{ob.id}:now:{ob.attempts}")
    return {"status": "retrying"}


@router.post("/outbox/{outbox_id}/cancel")
async def cancel_waiting_send(
    outbox_id: uuid.UUID, user: CurrentUser, db: DbSession
) -> dict[str, str]:
    """Stop trying: the message goes back to Drafts (Edit, or before Delete)."""
    ob = await _waiting_send(db, user.id, outbox_id)
    ob.state = "cancelled"
    ob.next_attempt_at = None
    await db.commit()
    return {"status": "cancelled", "draft_id": str(ob.draft_id)}