ForcePilot/backend/package/yuxi/channel/extensions/bluebubbles/catchup.py

import hashlib
import json
import logging
import re
import time
from dataclasses import dataclass, field
from pathlib import Path

logger = logging.getLogger(__name__)


@dataclass
class CatchupCursor:
    last_seen_ms: int | None = None
    last_seen_rowid: int | None = None
    updated_at: float = 0.0
    failure_retries: dict[str, int] = field(default_factory=dict)

    MAX_FAILURE_MAP = 5000


@dataclass
class CatchupSummary:
    query_succeeded: bool = False
    replayed: int = 0
    skipped_from_me: int = 0
    skipped_pre_cursor: int = 0
    skipped_given_up: int = 0
    failed: int = 0
    given_up: int = 0
    cursor_before: int | None = None
    cursor_after: int = 0
    window_start_ms: int = 0
    window_end_ms: int = 0
    fetched_count: int = 0


def _cursor_path(state_dir: Path, account_id: str) -> Path:
    safe = re.sub(r"[^a-zA-Z0-9_-]", "_", account_id)
    h = hashlib.sha256(account_id.encode()).hexdigest()[:12]
    return state_dir / "bluebubbles" / "catchup" / f"{safe}__{h}.json"


def load_cursor(state_dir: Path, account_id: str) -> CatchupCursor:
    path = _cursor_path(state_dir, account_id)
    if not path.exists():
        return CatchupCursor()
    try:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
    except (json.JSONDecodeError, OSError):
        return CatchupCursor()
    return CatchupCursor(
        last_seen_ms=data.get("lastSeenMs"),
        last_seen_rowid=data.get("lastSeenRowid"),
        updated_at=data.get("updatedAt", 0),
        failure_retries=data.get("failureRetries", {}),
    )


def save_cursor(state_dir: Path, account_id: str, cursor: CatchupCursor):
    path = _cursor_path(state_dir, account_id)
    path.parent.mkdir(parents=True, exist_ok=True)
    if len(cursor.failure_retries) > CatchupCursor.MAX_FAILURE_MAP:
        sorted_items = sorted(cursor.failure_retries.items(), key=lambda x: x[1], reverse=True)
        cursor.failure_retries = dict(sorted_items[: CatchupCursor.MAX_FAILURE_MAP])
    with open(path, "w", encoding="utf-8") as f:
        json.dump(
            {
                "lastSeenMs": cursor.last_seen_ms,
                "lastSeenRowid": cursor.last_seen_rowid,
                "updatedAt": time.time(),
                "failureRetries": cursor.failure_retries,
            },
            f,
        )


async def run_catchup(
    client,
    state_dir: Path,
    account_id: str,
    catchup_config,
    process_message: callable,
    dedupe_store,
) -> CatchupSummary:
    if not catchup_config or not catchup_config.enabled:
        return CatchupSummary()

    cursor = load_cursor(state_dir, account_id)
    summary = CatchupSummary(
        cursor_before=cursor.last_seen_ms,
        window_end_ms=int(time.time() * 1000),
    )

    now_ms = int(time.time() * 1000)
    lookback_limit_ms = catchup_config.max_age_minutes * 60 * 1000

    if cursor.last_seen_ms:
        summary.window_start_ms = cursor.last_seen_ms
    else:
        summary.window_start_ms = now_ms - catchup_config.first_run_lookback_minutes * 60 * 1000

    if summary.window_start_ms < now_ms - lookback_limit_ms:
        summary.window_start_ms = now_ms - lookback_limit_ms

    if summary.window_start_ms >= summary.window_end_ms:
        return summary

    try:
        params = {
            "limit": catchup_config.per_run_limit,
            "sort": "ASC",
        }
        resp = await client.get("/api/v1/messages", params=params)
        resp.raise_for_status()
        summary.query_succeeded = True

        data = resp.json()
        messages = data if isinstance(data, list) else data.get("data", data.get("messages", []))

        summary.fetched_count = len(messages)

        max_failures = catchup_config.max_failure_retries

        for msg_data in messages:
            guid = msg_data.get("guid", "")
            delivered = msg_data.get("dateDelivered", msg_data.get("date_delivered", 0))

            if cursor.last_seen_ms and delivered <= cursor.last_seen_ms:
                summary.skipped_pre_cursor += 1
                continue

            if msg_data.get("isFromMe", msg_data.get("is_from_me", False)):
                summary.skipped_from_me += 1
                continue

            failures = cursor.failure_retries.get(guid, 0)
            if failures >= max_failures:
                summary.skipped_given_up += 1
                continue

            if dedupe_store and dedupe_store.is_duplicate(guid, account_id):
                continue

            try:
                await process_message(msg_data)
                summary.replayed += 1
            except Exception:
                summary.failed += 1
                cursor.failure_retries[guid] = failures + 1
                if failures + 1 >= max_failures:
                    summary.given_up += 1

            if delivered > summary.cursor_after:
                summary.cursor_after = delivered

        cursor.last_seen_ms = summary.cursor_after or summary.window_end_ms
        save_cursor(state_dir, account_id, cursor)

    except Exception as e:
        logger.warning("Catchup query failed for account %s: %s", account_id, e)

    return summary