webmail-client/backend/app/sync/engine.py

"""Sync engine: bootstrap mailboxes, backfill, and the live IDLE/poll loop (SPEC §6.3)."""

from __future__ import annotations

import asyncio
import contextlib
import logging
import random
import uuid
from typing import Any

from sqlalchemy import delete, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import get_settings
from app.events import get_event_bus
from app.indexer.indexer import PersistResult, persist_message
from app.models import Account, Draft, Mailbox, Message
from app.providers.base import MailboxDelta, ProviderNotSupported
from app.providers.imap_smtp import ImapSmtpProvider
from app.search.factory import get_search_backend
from app.sync.folder_names import infer_special_use
from app.sync.provider_factory import build_provider

logger = logging.getLogger(__name__)

# Backfill priority: inbox, then sent/drafts, then everything else (SPEC §6.3).
_PRIORITY = {"inbox": 0, "sent": 1, "drafts": 2}


async def _load_account(session: AsyncSession, account_id: uuid.UUID) -> Account:
    account = await session.get(Account, account_id)
    if account is None:
        raise ValueError(f"account {account_id} not found")
    return account


async def _set_state(account_id: uuid.UUID, state: str, detail: str = "") -> None:
    from app.db import session_scope

    async with session_scope() as session:
        await session.execute(
            update(Account).where(Account.id == account_id).values(state=state, state_detail=detail)
        )
    # Keep the sync:live registry in step with the account's state so the sync
    # coordinator claims it once live and drops it otherwise (refactor §2).
    from app.redis_client import get_redis
    from app.sync import leases

    redis = get_redis()
    if state == "live":
        await leases.mark_live(redis, account_id)
    else:
        await leases.mark_not_live(redis, account_id)

    bus = get_event_bus()
    account = None
    async with session_scope() as session:
        account = await session.get(Account, account_id)
    if account is not None:
        await bus.publish(
            "account.state_changed",
            {"user_id": str(account.user_id), "account_id": str(account_id), "state": state},
        )


async def bootstrap_mailboxes(
    session: AsyncSession, provider: ImapSmtpProvider, account: Account
) -> list[Mailbox]:
    """List remote mailboxes and upsert local ``mailboxes`` rows (SPEC §6.3)."""
    remote = await provider.list_mailboxes()
    result: list[Mailbox] = []
    for rm in remote:
        special = rm.special_use or infer_special_use(rm.path)
        existing = await session.scalar(
            select(Mailbox).where(Mailbox.account_id == account.id, Mailbox.path == rm.path)
        )
        if existing is None:
            mailbox = Mailbox(
                account_id=account.id,
                user_id=account.user_id,
                path=rm.path,
                special_use=special,
                uidvalidity=rm.uidvalidity,
            )
            session.add(mailbox)
        else:
            existing.special_use = special
            existing.uidvalidity = rm.uidvalidity
            mailbox = existing
        result.append(mailbox)
    await session.flush()
    return result


def _mailbox_order(mailboxes: list[Mailbox]) -> list[Mailbox]:
    return sorted(mailboxes, key=lambda m: _PRIORITY.get(m.special_use or "", 99))


async def _emit_batch(
    results: list[PersistResult], account: Account, mailbox_path: str, done: int
) -> None:
    backend = get_search_backend()
    bus = get_event_bus()
    for r in results:
        await backend.index(r.search_doc)
        await bus.publish(
            "message.new",
            {
                "user_id": str(account.user_id),
                "account_id": str(account.id),
                "message_id": str(r.message.id),
                "thread_id": str(r.message.thread_id),
            },
        )
    await bus.publish(
        "sync.progress",
        {
            "user_id": str(account.user_id),
            "account_id": str(account.id),
            "mailbox": mailbox_path,
            "done": done,
        },
    )


async def backfill_mailbox(provider: ImapSmtpProvider, account: Account, mailbox: Mailbox) -> int:
    """Backfill one mailbox newest-first, indexing inline per batch (FR-30)."""
    from app.db import session_scope

    account_id, user_id, mailbox_id = account.id, account.user_id, mailbox.id
    special_use, path = mailbox.special_use, mailbox.path
    done = 0
    async for batch in provider.backfill(path, batch_size=200):
        results: list[PersistResult] = []
        max_uid = 0
        async with session_scope() as session:
            for remote in batch:
                r = await persist_message(
                    session,
                    user_id=user_id,
                    account_id=account_id,
                    mailbox_id=mailbox_id,
                    mailbox_special_use=special_use,
                    remote=remote,
                )
                if r is not None:
                    results.append(r)
                max_uid = max(max_uid, remote.uid)
            await session.execute(
                update(Mailbox)
                .where(Mailbox.id == mailbox_id, Mailbox.last_seen_uid < max_uid)
                .values(last_seen_uid=max_uid)
            )
        done += len(results)
        await _emit_batch(results, account, path, done)
    return done


async def bootstrap_and_backfill(account_id: uuid.UUID) -> int:
    """Full initial sync for an account: mailboxes + backfill; sets state=live."""
    from app.db import session_scope

    async with session_scope() as session:
        account = await _load_account(session, account_id)
        provider = build_provider(account)
        mailboxes = _mailbox_order(await bootstrap_mailboxes(session, provider, account))
        # Detach: reuse account/mailboxes values after the session closes.
        account_obj = account
        mailbox_list = list(mailboxes)

    total = 0
    try:
        for mailbox in mailbox_list:
            total += await backfill_mailbox(provider, account_obj, mailbox)
        await _set_state(account_id, "live")
    except Exception as exc:  # noqa: BLE001 - surface provider errors as account state
        logger.exception("backfill failed for account %s", account_id)
        await _set_state(account_id, "error", str(exc)[:500])
        raise
    finally:
        await provider.close()
    return total


async def apply_delta(
    provider: ImapSmtpProvider, account: Account, mailbox: Mailbox, delta: MailboxDelta
) -> None:
    """Apply a poll delta: new messages, flag changes, uidvalidity resets (SPEC §6.3)."""
    from app.db import session_scope

    bus = get_event_bus()

    if delta.new_uidvalidity is not None:
        # UIDVALIDITY changed: wipe this mailbox's rows and re-backfill (SPEC §6.3).
        async with session_scope() as session:
            await session.execute(delete(Message).where(Message.mailbox_id == mailbox.id))
            await session.execute(
                update(Mailbox)
                .where(Mailbox.id == mailbox.id)
                .values(uidvalidity=delta.new_uidvalidity, last_seen_uid=0)
            )
        await backfill_mailbox(provider, account, mailbox)
        return

    results: list[PersistResult] = []
    max_uid = mailbox.last_seen_uid
    async with session_scope() as session:
        for remote in delta.new:
            r = await persist_message(
                session,
                user_id=account.user_id,
                account_id=account.id,
                mailbox_id=mailbox.id,
                mailbox_special_use=mailbox.special_use,
                remote=remote,
            )
            if r is not None:
                results.append(r)
            max_uid = max(max_uid, remote.uid)
        for change in delta.flag_changes:
            await session.execute(
                update(Message)
                .where(Message.mailbox_id == mailbox.id, Message.uid == change.uid)
                .values(flags=sorted(change.flags))
            )
        if delta.expunged_uids:
            await session.execute(
                delete(Message).where(
                    Message.mailbox_id == mailbox.id, Message.uid.in_(delta.expunged_uids)
                )
            )
        updates = {"last_seen_uid": max_uid}
        if delta.new_highestmodseq is not None:
            updates["highestmodseq"] = delta.new_highestmodseq
        await session.execute(update(Mailbox).where(Mailbox.id == mailbox.id).values(**updates))

    await _emit_batch(results, account, mailbox.path, len(results))
    for change in delta.flag_changes:
        await bus.publish(
            "flags.changed",
            {
                "user_id": str(account.user_id),
                "account_id": str(account.id),
                "mailbox": mailbox.path,
                "uid": change.uid,
                "flags": sorted(change.flags),
            },
        )

    # Expunge reconciliation for Drafts. The mirror mechanism (and other clients)
    # expunge server copies, but a CONDSTORE poll doesn't report vanished UIDs, so
    # expunged drafts would linger in the DB and resurface as phantom "remote"
    # drafts. Drafts folders are small, so a full UID reconcile per poll is cheap.
    if mailbox.special_use == "drafts":
        await _reconcile_drafts(provider, account, mailbox)


async def _reconcile_drafts(
    provider: ImapSmtpProvider, account: Account, mailbox: Mailbox
) -> None:
    """Drop DB rows for drafts that no longer exist on the server (were expunged)."""
    from app.db import session_scope

    try:
        server_uids = await provider.list_uids(mailbox.path)
    except Exception:
        logger.exception("drafts reconcile: list_uids failed for %s", mailbox.path)
        return

    async with session_scope() as session:
        db_uids = set(
            await session.scalars(select(Message.uid).where(Message.mailbox_id == mailbox.id))
        )
        gone = db_uids - server_uids
        if not gone:
            return
        await session.execute(
            delete(Message).where(Message.mailbox_id == mailbox.id, Message.uid.in_(gone))
        )
        # A local draft whose mirrored copy vanished becomes local-only again.
        await session.execute(
            update(Draft)
            .where(Draft.account_id == account.id, Draft.imap_uid.in_(gone))
            .values(imap_uid=None)
        )

    logger.info("drafts reconcile: removed %d expunged row(s) in %s", len(gone), mailbox.path)
    await get_event_bus().publish(
        "message.deleted",
        {"user_id": str(account.user_id), "account_id": str(account.id), "mailbox": mailbox.path},
    )


async def live_loop(
    account_id: uuid.UUID, mailbox_path: str = "INBOX", stop: asyncio.Event | None = None
) -> None:
    """Maintain a live IDLE/poll loop for one mailbox (SPEC §6.3)."""
    from app.db import session_scope

    stop = stop or asyncio.Event()
    idle_fallback = get_settings().poll_warm

    async with session_scope() as session:
        account = await _load_account(session, account_id)
        account_obj = account

    provider = build_provider(account_obj)
    try:
        while not stop.is_set():
            try:
                await provider.idle(mailbox_path)
            except ProviderNotSupported:
                await asyncio.sleep(idle_fallback)
            if stop.is_set():
                break
            async with session_scope() as session:
                mailbox = await session.scalar(
                    select(Mailbox).where(
                        Mailbox.account_id == account_id, Mailbox.path == mailbox_path
                    )
                )
                mailbox_obj = mailbox
            if mailbox_obj is None:
                break
            delta = await provider.poll(
                mailbox_path,
                last_seen_uid=mailbox_obj.last_seen_uid,
                highestmodseq=mailbox_obj.highestmodseq,
            )
            await apply_delta(provider, account_obj, mailbox_obj, delta)
    finally:
        await provider.close()


async def poll_folders_once(account_id: uuid.UUID, *, include_inbox: bool) -> None:
    """One CONDSTORE poll pass over an account's special-use folders.

    ``include_inbox`` also polls INBOX (used by the manual Refresh); the periodic
    secondary loop leaves INBOX to its dedicated IDLE loop. Custom (non special-use)
    folders are not polled. Drafts are polled so drafts made in other clients appear
    (persisted un-threaded).
    """
    from app.db import session_scope

    async with session_scope() as session:
        account = await _load_account(session, account_id)
        account_obj = account
        stmt = select(Mailbox.id).where(
            Mailbox.account_id == account_id, Mailbox.special_use.isnot(None)
        )
        if not include_inbox:
            stmt = stmt.where(Mailbox.special_use != "inbox")
        mailbox_ids = list(await session.scalars(stmt))
    if not mailbox_ids:
        return

    provider = build_provider(account_obj)
    try:
        for mailbox_id in mailbox_ids:
            async with session_scope() as session:
                mailbox = await session.get(Mailbox, mailbox_id)
                mailbox_obj = mailbox
            if mailbox_obj is None:
                continue
            try:
                delta = await provider.poll(
                    mailbox_obj.path,
                    last_seen_uid=mailbox_obj.last_seen_uid,
                    highestmodseq=mailbox_obj.highestmodseq,
                )
                await apply_delta(provider, account_obj, mailbox_obj, delta)
            except Exception:
                logger.exception("poll failed: account=%s mailbox=%s", account_id, mailbox_obj.path)
    finally:
        await provider.close()


async def _poll_secondary_once(account_id: uuid.UUID) -> None:
    await poll_folders_once(account_id, include_inbox=False)


async def refresh_account_job(ctx: dict[str, Any], account_id: str) -> None:
    """Poll every special-use folder (incl. INBOX) once — the manual Refresh action."""
    await poll_folders_once(uuid.UUID(account_id), include_inbox=True)


async def secondary_poll_loop(
    account_id: uuid.UUID, stop: asyncio.Event | None = None, interval: float | None = None
) -> None:
    """Periodically CONDSTORE-poll non-INBOX folders (Sent/Drafts/Archive/Trash/Junk).

    INBOX has its own IDLE loop (real-time). This keeps the other folders in step
    with server-side changes (e.g. Gmail-filed sent mail) and changes made by other
    clients. CONDSTORE/QRESYNC makes each pass a cheap MODSEQ delta. Used by hot-tier
    runners alongside the INBOX IDLE loop.
    """
    stop = stop or asyncio.Event()
    period = interval if interval is not None else get_settings().poll_warm
    first = True
    while not stop.is_set():
        # Wait first (initial state was covered by backfill). Spread the fleet: the
        # first wait is a random point in the period, later waits jitter ±10%, so
        # accounts claimed together don't poll in lockstep (thundering herd).
        wait = random.uniform(0, period) if first else period * random.uniform(0.9, 1.1)
        first = False
        with contextlib.suppress(TimeoutError):
            await asyncio.wait_for(stop.wait(), timeout=wait)
        if stop.is_set():
            break
        await _poll_secondary_once(account_id)


async def poll_all_loop(
    account_id: uuid.UUID, interval: float, stop: asyncio.Event | None = None
) -> None:
    """Transient CONDSTORE poll of INBOX + special folders every ``interval`` seconds.

    The warm/cold tiers use this instead of a persistent IDLE connection: no socket
    is held between passes. Correctness matches the hot tier — only latency differs.
    """
    stop = stop or asyncio.Event()
    first = True
    while not stop.is_set():
        # Spread the fleet: first poll at a random point in the interval, then
        # jitter ±10% so a batch claimed together doesn't poll in lockstep.
        wait = random.uniform(0, interval) if first else interval * random.uniform(0.9, 1.1)
        first = False
        with contextlib.suppress(TimeoutError):
            await asyncio.wait_for(stop.wait(), timeout=wait)
        if stop.is_set():
            break
        await poll_folders_once(account_id, include_inbox=True)