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

logger = logging.getLogger(__name__)

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

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


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,
) -> 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``.
    """
    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,
                )
                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)

    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).
        total = 0
        for mailbox in mailbox_list:
            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:
            if since is not None:  # 3months: single windowed pass
                count += await backfill_mailbox(
                    provider, account_obj, mailbox, progress, since=since
                )
            else:  # all: recent window first (usable fast), then the remainder
                count += await backfill_mailbox(
                    provider, account_obj, mailbox, progress, since=recent_cutoff
                )
                count += await backfill_mailbox(
                    provider, account_obj, mailbox, progress, before=recent_cutoff
                )
        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
    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)