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 dataclasses import dataclass, field
from datetime import date, timedelta
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
from app.threading_.threader import denormalize_thread

logger = logging.getLogger(__name__)

# Backfill priority: inbox, then sent/drafts, then everything else (SPEC §6.3).
# archive (Gmail All Mail) sorts last so its dedupe pass runs after the system folders.
_PRIORITY = {"inbox": 0, "sent": 1, "drafts": 2}
# Gmail sync (SPEC: Gmail sync): the only folders backfilled — All Mail (archive) is
# the complete store, the rest are the system folders needed for their views. Every
# custom-label folder is skipped; labels come from X-GM-LABELS instead.
_GMAIL_SYNC_FOLDERS = {"inbox", "sent", "drafts", "trash", "junk", "archive"}
# Progress denominator for Gmail: All Mail already contains INBOX/Sent/Drafts/labels,
# so only these three are counted (Trash/Spam sit outside All Mail).
_GMAIL_DENOM_FOLDERS = {"archive", "trash", "junk"}

# Sync-window import strategy (SPEC §12 M2). "all" imports the most recent
# _RECENT_DAYS first (so the account is usable fast), then the older remainder.
_RECENT_DAYS = 30
_WINDOW_DAYS = {"3months": 90}


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 _index_and_notify(results: list[PersistResult], account: Account) -> None:
    """Index each persisted message and publish its ``message.new`` event."""
    backend = get_search_backend()
    bus = get_event_bus()
    for r in results:
        await backend.index(r.search_doc)
        payload: dict[str, object] = {
            "user_id": str(account.user_id),
            "account_id": str(account.id),
            "message_id": str(r.message.id),
            "thread_id": str(r.message.thread_id),
        }
        # Gmail custom labels ride along so the labels plugin can map them to chips
        # (SPEC: Gmail sync). Only present on Gmail accounts with labelled messages.
        if r.gmail_labels:
            payload["gmail_labels"] = list(r.gmail_labels)
        await bus.publish("message.new", payload)


async def _emit_batch(
    results: list[PersistResult], account: Account, mailbox_path: str, done: int
) -> None:
    await _index_and_notify(results, account)
    await get_event_bus().publish(
        "sync.progress",
        {
            "user_id": str(account.user_id),
            "account_id": str(account.id),
            "mailbox": mailbox_path,
            "done": done,
        },
    )


@dataclass
class _BackfillProgress:
    """Account-level backfill counter: persists ``backfill_done`` and emits
    ``sync.progress`` with the running total so the UI can show a progress bar."""

    account: Account
    total: int
    done: int = 0
    _mailboxes: set[str] = field(default_factory=set)

    async def record(self, results: list[PersistResult], mailbox_path: str) -> None:
        from app.db import session_scope

        await _index_and_notify(results, self.account)
        self.done += len(results)
        self._mailboxes.add(mailbox_path)
        async with session_scope() as session:
            await session.execute(
                update(Account)
                .where(Account.id == self.account.id)
                .values(backfill_done=self.done)
            )
        await get_event_bus().publish(
            "sync.progress",
            {
                "user_id": str(self.account.user_id),
                "account_id": str(self.account.id),
                "mailbox": mailbox_path,
                "done": self.done,
                "total": self.total,
            },
        )


async def backfill_mailbox(
    provider: ImapSmtpProvider,
    account: Account,
    mailbox: Mailbox,
    progress: _BackfillProgress | None = None,
    *,
    since: date | None = None,
    before: date | None = None,
    dedupe_account_scope: bool = False,
) -> int:
    """Backfill one mailbox newest-first, indexing inline per batch (FR-30).

    ``since``/``before`` bound the import by date (sync window, SPEC §12 M2).
    When a ``progress`` tracker is given, batches count toward the account-level
    backfill total; otherwise each batch emits a standalone ``sync.progress``.
    ``dedupe_account_scope`` drops messages already synced under another mailbox of
    the account (Gmail All Mail pass — SPEC: Gmail sync / no-duplication).
    """
    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, since=since, before=before):
        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,
                    dedupe_account_scope=dedupe_account_scope,
                )
                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)
        if progress is not None:
            await progress.record(results, path)
        else:
            await _emit_batch(results, account, path, done)
    return done


async def _set_backfill_counts(account_id: uuid.UUID, *, total: int, done: int) -> None:
    from app.db import session_scope

    async with session_scope() as session:
        await session.execute(
            update(Account)
            .where(Account.id == account_id)
            .values(backfill_total=total, backfill_done=done)
        )


async def _sync_new_only(provider: ImapSmtpProvider, mailboxes: list[Mailbox]) -> None:
    """``new`` window: import no history — just move each mailbox watermark to the
    current tip so only mail arriving after setup is pulled (SPEC §12 M2)."""
    from app.db import session_scope

    for mailbox in mailboxes:
        tip = await provider.newest_uid(mailbox.path)
        async with session_scope() as session:
            await session.execute(
                update(Mailbox)
                .where(Mailbox.id == mailbox.id, Mailbox.last_seen_uid < tip)
                .values(last_seen_uid=tip)
            )


async def bootstrap_and_backfill(account_id: uuid.UUID) -> int:
    """Full initial sync for an account: mailboxes + windowed backfill; 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
        window = account.sync_window
        mailbox_list = list(mailboxes)
        # Gmail: sync only the system folders + All Mail and read each message's
        # labels (X-GM-LABELS) instead of visiting every label folder. This avoids
        # the INBOX/All-Mail/label duplication and lets labels map to chips lazily,
        # tied to the sync window (SPEC: Gmail sync).
        gmail = provider.gmail_labels_supported
        if gmail:
            mailbox_list = [m for m in mailbox_list if m.special_use in _GMAIL_SYNC_FOLDERS]
        # Spam is opt-in: skip the Junk folder unless the account enabled it (§6.3).
        if not account.sync_spam:
            mailbox_list = [m for m in mailbox_list if m.special_use != "junk"]

    try:
        if window == "new":
            # No backfill: set watermarks to the tip, go live immediately.
            await _sync_new_only(provider, mailbox_list)
            await _set_backfill_counts(account_id, total=0, done=0)
            await _set_state(account_id, "live")
            return 0

        since = (
            date.today() - timedelta(days=_WINDOW_DAYS[window]) if window in _WINDOW_DAYS else None
        )
        # Progress denominator: count what we're about to import up front so the
        # bar has a stable total (SPEC §12 M2 backfill-progress UX). For Gmail only
        # All Mail + Trash + Spam are counted — INBOX/Sent/Drafts are subsets of All
        # Mail, so counting them too would inflate the total past the deduped `done`.
        total = 0
        for mailbox in mailbox_list:
            if gmail and mailbox.special_use not in _GMAIL_DENOM_FOLDERS:
                continue
            total += await provider.message_count(mailbox.path, since=since)
        await _set_backfill_counts(account_id, total=total, done=0)
        progress = _BackfillProgress(account_obj, total)

        recent_cutoff = date.today() - timedelta(days=_RECENT_DAYS)
        count = 0
        for mailbox in mailbox_list:
            # Dedupe the Gmail All Mail pass against messages already stored under a
            # system folder (INBOX/Sent/…), which the priority order backfills first.
            dedupe = gmail and mailbox.special_use == "archive"
            if since is not None:  # 3months: single windowed pass
                count += await backfill_mailbox(
                    provider, account_obj, mailbox, progress, since=since,
                    dedupe_account_scope=dedupe,
                )
            else:  # all: recent window first (usable fast), then the remainder
                count += await backfill_mailbox(
                    provider, account_obj, mailbox, progress, since=recent_cutoff,
                    dedupe_account_scope=dedupe,
                )
                count += await backfill_mailbox(
                    provider, account_obj, mailbox, progress, before=recent_cutoff,
                    dedupe_account_scope=dedupe,
                )
        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 count


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
    gone_ids: list[uuid.UUID] = []
    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)
        # Threads whose denormalized counters this batch invalidates. A flag set
        # on the server (read on a phone, filed by a server-side rule) and an
        # expunge both change what `threads.unread_count` should say, and neither
        # goes through apply_flags_to_messages — so without this the row keeps its
        # old count and the list shows a message as unread that the reading pane
        # shows as read.
        touched_threads: set[uuid.UUID] = set()

        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.flag_changes:
            touched_threads.update(
                await session.scalars(
                    select(Message.thread_id).where(
                        Message.mailbox_id == mailbox.id,
                        Message.uid.in_([c.uid for c in delta.flag_changes]),
                        Message.thread_id.is_not(None),
                    )
                )
            )
        if delta.expunged_uids:
            # Server-side expunges (QRESYNC/VANISHED, SPEC §6.3). Capture the row ids
            # first so we can drop them from the search index after the delete.
            gone = (
                await session.execute(
                    select(Message.id, Message.thread_id).where(
                        Message.mailbox_id == mailbox.id, Message.uid.in_(delta.expunged_uids)
                    )
                )
            ).all()
            gone_ids = [row[0] for row in gone]
            touched_threads.update(row[1] for row in gone if row[1] is not None)
            await session.execute(
                delete(Message).where(
                    Message.mailbox_id == mailbox.id, Message.uid.in_(delta.expunged_uids)
                )
            )
            # A local draft whose mirrored copy vanished reverts to local-only.
            await session.execute(
                update(Draft)
                .where(Draft.account_id == account.id, Draft.imap_uid.in_(delta.expunged_uids))
                .values(imap_uid=None)
            )
        # After the deletes, so a thread left with no messages is removed rather
        # than recounted (denormalize_thread handles that case).
        for tid in touched_threads:
            await denormalize_thread(session, tid)

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

    if gone_ids:
        # Drop expunged messages from the search index (PG keeps its tsvector on the
        # row, so this is a no-op there; Meili needs the explicit removal) and tell
        # the UI to refresh, mirroring the drafts-reconcile notification.
        backend = get_search_backend()
        for mid in gone_ids:
            await backend.delete(mid)
        await bus.publish(
            "message.deleted",
            {
                "user_id": str(account.user_id),
                "account_id": str(account.id),
                "mailbox": mailbox.path,
            },
        )

    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). ``archive`` (Gmail All Mail) is never live-polled — it is
    huge and cold, and backfill already covers it.
    """
    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),
            Mailbox.special_use != "archive",
        )
        if not include_inbox:
            stmt = stmt.where(Mailbox.special_use != "inbox")
        if not account.sync_spam:  # Spam is opt-in (§6.3)
            stmt = stmt.where(Mailbox.special_use != "junk")
        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)