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 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, 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),
            },
        )


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
    while not stop.is_set():
        # Wait first: the initial state was already covered by the backfill.
        with contextlib.suppress(TimeoutError):
            await asyncio.wait_for(stop.wait(), timeout=period)
        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()
    while not stop.is_set():
        with contextlib.suppress(TimeoutError):
            await asyncio.wait_for(stop.wait(), timeout=interval)
        if stop.is_set():
            break
        await poll_folders_once(account_id, include_inbox=True)