FabledCurator/backend/app/tasks/backup.py

"""FC-3h: backup/restore Celery tasks.

All tasks live on the maintenance queue (per celery_app.task_routes).
task_run lifecycle tracking is automatic via FC-3i signals — these
tasks just record the operator-facing artifact metadata into
BackupRun.
"""
from __future__ import annotations

import logging
import time
import uuid
from datetime import UTC, datetime, timedelta
from pathlib import Path

from celery.exceptions import SoftTimeLimitExceeded
from sqlalchemy import select, text, update
from sqlalchemy.exc import DBAPIError, OperationalError

from ..celery_app import celery
from ..config import get_config
from ..models import BackupRun, ImportSettings
from ..services import (
    backup_destination,
    backup_media,
    backup_restore,
    backup_s3,
    backup_service,
    backup_verify,
)
from ._sync_engine import get_sync_engine
from ._sync_engine import sync_session_factory as _sync_session_factory

log = logging.getLogger(__name__)
IMAGES_ROOT = Path("/images")


def _mark_failed(session, row: BackupRun, exc: BaseException) -> None:
    """Flip a BackupRun row from running/restoring to error with a
    truncated error message and finished_at. Caller already holds the
    session open."""
    row.status = "error"
    row.error = f"{type(exc).__name__}: {exc}"[:2000]
    row.finished_at = datetime.now(UTC)
    session.add(row)
    session.commit()


@celery.task(
    name="backend.app.tasks.backup.backup_db_task",
    bind=True,
    autoretry_for=(OperationalError, DBAPIError),
    retry_backoff=10, retry_backoff_max=120, max_retries=2,
    # A pg_dump can't be chunked; the 12-min limit timed out once the DB grew
    # (operator-flagged 2026-06-07). 30/35 min gives real headroom. (A long
    # backup still briefly holds the concurrency-1 maintenance lane — the
    # structural fix is a dedicated lane for the long one-shots.)
    soft_time_limit=1800, time_limit=2100,
)
def backup_db_task(self, *, tag: str | None = None,
                   triggered_by: str = "manual") -> dict:
    """Create one DB backup. Returns {'backup_run_id': N}.

    With an off-site destination set up, the dump streams to that bucket and
    nothing is written here; the row's manifest says where it went
    (`offsite`: endpoint, bucket, key). Without one, the dump is written to
    /images/_backups."""
    SessionLocal = _sync_session_factory()
    cfg = get_config()
    now = datetime.now(UTC)
    with SessionLocal() as session:
        # A restore replaces the database this would dump.
        if restore_going(session):
            return {"skipped": "a restore is going"}
        row = BackupRun(
            kind="db", status="running", tag=tag,
            triggered_by=triggered_by, started_at=now, manifest={},
        )
        session.add(row)
        session.commit()
        session.refresh(row)
        run_id = row.id

    try:
        with SessionLocal() as session:
            settings = ImportSettings.load_sync(session)
            target = backup_destination.saved_target(settings)
            limit_mbps = settings.backup_upload_limit_mbps
        if target is not None:
            result = _dump_offsite(cfg, target, limit_mbps)
        else:
            result = backup_service.backup_db(
                db_url=cfg.database_url_sync, images_root=IMAGES_ROOT,
                tag=tag, triggered_by=triggered_by,
            )
    except (SoftTimeLimitExceeded, Exception) as exc:
        with SessionLocal() as session:
            row = session.get(BackupRun, run_id)
            if row is not None:
                _mark_failed(session, row, exc)
        raise

    with SessionLocal() as session:
        row = session.get(BackupRun, run_id)
        row.status = "ok"
        row.finished_at = datetime.now(UTC)
        row.size_bytes = result["size_bytes"]
        if target is not None:
            row.manifest = {
                "ts": result["ts"],
                "offsite": {
                    "endpoint": result["endpoint"],
                    "bucket": result["bucket"],
                    "key": result["key"],
                    "parts": result["parts"],
                },
            }
        else:
            row.sql_path = result["sql_path"]
            row.manifest = {
                "manifest_path": result["manifest_path"],
                "ts": result["ts"],
            }
        session.commit()
    return {"backup_run_id": run_id}


@celery.task(
    name="backend.app.tasks.backup.backup_images_task",
    bind=True,
    autoretry_for=(OperationalError, DBAPIError),
    retry_backoff=30, retry_backoff_max=300, max_retries=1,
    soft_time_limit=21600, time_limit=23400,
)
def backup_images_task(self, *, tag: str | None = None,
                       triggered_by: str = "manual") -> dict:
    """Create one images backup. Same shape as backup_db_task; uses
    tar_path instead of sql_path."""
    SessionLocal = _sync_session_factory()
    now = datetime.now(UTC)
    with SessionLocal() as session:
        row = BackupRun(
            kind="images", status="running", tag=tag,
            triggered_by=triggered_by, started_at=now, manifest={},
        )
        session.add(row)
        session.commit()
        session.refresh(row)
        run_id = row.id

    try:
        result = backup_service.backup_images(
            images_root=IMAGES_ROOT,
            tag=tag, triggered_by=triggered_by,
        )
    except (SoftTimeLimitExceeded, Exception) as exc:
        with SessionLocal() as session:
            row = session.get(BackupRun, run_id)
            if row is not None:
                _mark_failed(session, row, exc)
        raise

    with SessionLocal() as session:
        row = session.get(BackupRun, run_id)
        row.status = "ok"
        row.finished_at = datetime.now(UTC)
        row.tar_path = result["tar_path"]
        row.size_bytes = result["size_bytes"]
        row.manifest = {
            "manifest_path": result["manifest_path"],
            "ts": result["ts"],
        }
        session.commit()
    return {"backup_run_id": run_id}


@celery.task(
    name="backend.app.tasks.backup.restore_db_task",
    bind=True,
    max_retries=0,  # NEVER auto-retry a half-applied restore.
    soft_time_limit=1200, time_limit=1800,
)
def restore_db_task(self, *, source_backup_run_id: int) -> dict:
    """Restore from a previous DB backup. Inserts a NEW BackupRun row
    (kind='db', status='restoring') linked to the source via
    restored_from_id; flips to 'restored' on success or 'error' on
    failure. Operator sees the restore as a row in the dashboard."""
    SessionLocal = _sync_session_factory()
    cfg = get_config()
    now = datetime.now(UTC)
    with SessionLocal() as session:
        src = session.get(BackupRun, source_backup_run_id)
        if src is None or src.kind != "db" or not src.sql_path:
            raise ValueError(
                f"BackupRun id={source_backup_run_id} is not a valid DB backup"
            )
        marker = BackupRun(
            kind="db", status="restoring",
            triggered_by="restore", started_at=now,
            restored_from_id=src.id,
            manifest={"source_sql_path": src.sql_path},
        )
        session.add(marker)
        session.commit()
        session.refresh(marker)
        marker_id = marker.id
        sql_path = src.sql_path

    try:
        backup_service.restore_db(
            db_url=cfg.database_url_sync, sql_path=Path(sql_path),
        )
    except (SoftTimeLimitExceeded, Exception) as exc:
        with SessionLocal() as session:
            row = session.get(BackupRun, marker_id)
            if row is not None:
                _mark_failed(session, row, exc)
        raise

    with SessionLocal() as session:
        row = session.get(BackupRun, marker_id)
        row.status = "restored"
        row.finished_at = datetime.now(UTC)
        session.commit()
    return {"backup_run_id": marker_id}


@celery.task(
    name="backend.app.tasks.backup.restore_images_task",
    bind=True, max_retries=0,
    soft_time_limit=21600, time_limit=23400,
)
def restore_images_task(self, *, source_backup_run_id: int) -> dict:
    """Mirrors restore_db_task; uses backup_service.restore_images."""
    SessionLocal = _sync_session_factory()
    now = datetime.now(UTC)
    with SessionLocal() as session:
        src = session.get(BackupRun, source_backup_run_id)
        if src is None or src.kind != "images" or not src.tar_path:
            raise ValueError(
                f"BackupRun id={source_backup_run_id} is not a valid images backup"
            )
        marker = BackupRun(
            kind="images", status="restoring",
            triggered_by="restore", started_at=now,
            restored_from_id=src.id,
            manifest={"source_tar_path": src.tar_path},
        )
        session.add(marker)
        session.commit()
        session.refresh(marker)
        marker_id = marker.id
        tar_path = src.tar_path

    try:
        backup_service.restore_images(
            images_root=IMAGES_ROOT, tar_path=Path(tar_path),
        )
    except (SoftTimeLimitExceeded, Exception) as exc:
        with SessionLocal() as session:
            row = session.get(BackupRun, marker_id)
            if row is not None:
                _mark_failed(session, row, exc)
        raise

    with SessionLocal() as session:
        row = session.get(BackupRun, marker_id)
        row.status = "restored"
        row.finished_at = datetime.now(UTC)
        session.commit()
    return {"backup_run_id": marker_id}


@celery.task(
    name="backend.app.tasks.backup.prune_backups",
    soft_time_limit=300, time_limit=600,
)
def prune_backups() -> dict:
    """Daily Beat. Per-kind retention from ImportSettings.

    Returns {"db_deleted": N, "images_deleted": M, "media_deleted": L,
    "restore_deleted": R, "reconcile_deleted": C, "test_deleted": T,
    "files_unlinked": K}. Media, restore, reconcile and test runs leave no
    files here; only their history is trimmed, to the latest MEDIA_RUNS_KEPT,
    RESTORE_RUNS_KEPT, RECONCILE_RUNS_KEPT and TEST_RUNS_KEPT.
    Tagged rows (tag IS NOT NULL) are never pruned.
    Status='running' / 'restoring' rows are never pruned — the
    recover_stalled_backup_runs sweep flips truly-stuck ones to
    'error' first. (Earlier docstring claimed the FC-3i TaskRun sweep
    handled those, but TaskRun cleanup never touched BackupRun rows.
    Audit 2026-06-02 added the dedicated sweep.)
    """
    SessionLocal = _sync_session_factory()
    counts = {"db_deleted": 0, "images_deleted": 0, "media_deleted": 0,
              "restore_deleted": 0, "reconcile_deleted": 0, "test_deleted": 0,
              "files_unlinked": 0}
    with SessionLocal() as session:
        s = ImportSettings.load_sync(session)
        for kind, keep in (
            ("db", s.backup_db_keep_last_n),
            ("images", s.backup_images_keep_last_n),
            ("media", MEDIA_RUNS_KEPT),
            ("restore", RESTORE_RUNS_KEPT),
            ("reconcile", RECONCILE_RUNS_KEPT),
            ("test", TEST_RUNS_KEPT),
        ):
            candidates = session.execute(
                select(BackupRun)
                .where(BackupRun.kind == kind)
                .where(BackupRun.tag.is_(None))
                .where(BackupRun.status.in_(["ok", "error", "restored"]))
                .order_by(BackupRun.started_at.desc())
                .offset(keep)
            ).scalars().all()
            for row in candidates:
                result = backup_service.unlink_artifact_files(
                    sql_path=row.sql_path,
                    tar_path=row.tar_path,
                    manifest_path=(row.manifest or {}).get("manifest_path"),
                )
                counts["files_unlinked"] += sum(
                    1 for v in result.values() if v
                )
                session.delete(row)
                counts[f"{kind}_deleted"] += 1
        session.commit()
    return counts


@celery.task(
    name="backend.app.tasks.backup.backup_db_nightly",
    soft_time_limit=60, time_limit=120,
)
def backup_db_nightly() -> dict:
    """Hourly tick. Dispatches a real backup ONLY if the configured
    UTC hour matches and the nightly setting is enabled. Returns
    either {'skipped': '<reason>'} or {'dispatched': '<task_id>'}."""
    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        s = ImportSettings.load_sync(session)
        nightly_enabled = s.backup_db_nightly_enabled
        configured_hour = s.backup_db_nightly_hour_utc
    if not nightly_enabled:
        return {"skipped": "nightly disabled"}
    now_hour = datetime.now(UTC).hour
    if now_hour != configured_hour:
        return {"skipped": f"hour={now_hour} != configured={configured_hour}"}
    res = backup_db_task.delay(triggered_by="nightly")
    out = {"dispatched": res.id}
    # With an off-site destination, the night's backup is the whole library:
    # the database dump, then any images and attachments not yet in the bucket.
    # Both run on the long-maintenance lane, one after the other.
    with SessionLocal() as session:
        try:
            offsite = backup_destination.saved_target(ImportSettings.load_sync(session))
        except Exception:  # noqa: BLE001 — the media run reports it itself
            offsite = True
    if offsite is not None:
        out["media_dispatched"] = backup_media_task.delay(triggered_by="nightly").id
    return out


# ── Off-site copy of the library's files ─────────────────────────────────────

# Each task invocation sends files for this long, then queues the next chunk
# of the same run, so other long-maintenance jobs get their turn in between.
MEDIA_CHUNK_SECONDS = 45 * 60
MEDIA_RUNS_KEPT = 60
# A failed file is listed on the run, up to this many; the count covers all.
MEDIA_FAILURES_LISTED = 20
# The weekly bucket check keeps about a quarter's history.
RECONCILE_RUNS_KEPT = 12


def _dump_offsite(cfg, target, limit_mbps) -> dict:
    """Dump the database to the bucket, with its row counts beside it for a
    test restore to check against. The dump stands without its counts, so a
    failure to store them is logged, not raised."""
    result = backup_service.backup_db_offsite(
        db_url=cfg.database_url_sync, target=target, limit_mbps=limit_mbps)
    if result.get("counts"):
        try:
            backup_verify.store_counts(target, result["key"], result["counts"])
        except (backup_s3.S3Error, OSError):
            log.warning("backup: the row counts for %s weren't stored", result["key"],
                        exc_info=True)
    return result


def _offsite_manifest(target) -> dict:
    return {"endpoint": target.endpoint, "bucket": target.bucket,
            "key": backup_media.PREFIX}


# ── One chunk at a time ──────────────────────────────────────────────────────
#
# A media backup and a restore's file phase run as a chain of chunks, each
# queuing the next. The chunk carrying a run writes its token to the run's
# manifest as "chunk", and hands that token to the chunk it queues as `after`.
# A chunk queued to resume a run after a restart carries `resume` instead and
# takes the run over. Any other message is stale (a killed chunk's own message
# delivered again, say) and is dropped, so a run never goes on twice and the
# chain never doubles.


def _chunk_token(task) -> str:
    return task.request.id or uuid.uuid4().hex


def _claim_run(session, run_id: int, status: str, token: str, *, after: str | None,
               resume: bool) -> str | None:
    """Take the run for this chunk. Returns None when taken, or why not."""
    row = session.get(BackupRun, run_id, with_for_update=True)
    if row is None or row.status != status:
        session.rollback()
        return "the run is no longer going"
    held = (row.manifest or {}).get("chunk")
    if not resume and held is not None and held != after:
        session.rollback()
        return "another chunk carries this run on"
    row.manifest = {**(row.manifest or {}), "chunk": token}
    session.commit()
    return None


def _still_ours(session, run_id: int, token: str) -> bool:
    """Lock the run's row and say whether this chunk still carries it. A chunk
    that was taken over stops at its next file rather than write over the
    newer chunk's place; on False the transaction is already rolled back."""
    held = session.execute(
        select(BackupRun.manifest).where(BackupRun.id == run_id).with_for_update()
    ).scalar_one_or_none()
    if (held or {}).get("chunk") == token:
        return True
    session.rollback()
    return False


def resume_offsite_runs() -> dict:
    """Queue the next chunk of every off-site run a restart cut short: a media
    backup still running, and a restore fetching its files. Called when the
    long-maintenance lane starts, where those chunks run, so none of theirs is
    alive then; the queued chunk carries on from the run's saved place. A
    restore stopped while loading its database isn't resumed: the stall sweep
    ends it, and the screen offers the restore again."""
    with _sync_session_factory()() as session:
        media = list(session.execute(
            select(BackupRun.id).where(BackupRun.kind == "media")
            .where(BackupRun.status == "running")
        ).scalars())
        restores = [
            r.id for r in session.execute(
                select(BackupRun).where(BackupRun.kind == "restore")
                .where(BackupRun.status == "restoring")
            ).scalars()
            if (r.manifest or {}).get("phase") == "files"
        ]
    try:
        tests = end_cut_off_tests()
    except Exception:  # noqa: BLE001 — resuming the runs below matters more
        log.exception("couldn't end a test restore cut off by the restart")
        tests = 0
    for run_id in media:
        backup_media_task.delay(run_id=run_id, resume=True)
    for run_id in restores:
        restore_offsite_task.delay(run_id=run_id, resume=True)
    if media or restores:
        log.info("resumed off-site runs after a restart: media %s, restore %s",
                 media, restores)
    return {"media": media, "restore": restores, "tests_ended": tests}


@celery.task(
    name="backend.app.tasks.backup.backup_media_task",
    bind=True,
    autoretry_for=(OperationalError, DBAPIError),
    retry_backoff=30, retry_backoff_max=300, max_retries=2,
    # A chunk stops starting files after MEDIA_CHUNK_SECONDS; the limits leave
    # room for the file in flight, a multi-gigabyte archive at a low upload
    # limit included.
    soft_time_limit=21600, time_limit=23400,
)
def backup_media_task(self, *, run_id: int | None = None, triggered_by: str = "manual",
                      after: str | None = None, resume: bool = False) -> dict:
    """Send the library's images and attachments that aren't in the off-site
    bucket yet, one chunk of a run. Called without `run_id` it starts a run
    (unless one is already going); each chunk queues the next until nothing is
    left, and a restart resumes the run where it stopped. A file that can't be
    sent is listed on the run and tried again by the next run. Returns
    {'backup_run_id': N, ...}."""
    SessionLocal = _sync_session_factory()
    now = datetime.now(UTC)
    token = _chunk_token(self)
    with SessionLocal() as session:
        if run_id is None:
            going = session.execute(
                select(BackupRun.id).where(BackupRun.kind == "media")
                .where(BackupRun.status == "running").limit(1)
            ).scalar_one_or_none()
            if going is not None:
                return {"backup_run_id": going, "skipped": "a media backup is already running"}
            # During a restore the library's files are still arriving.
            if restore_going(session):
                return {"skipped": "a restore is going"}
            row = BackupRun(kind="media", status="running", triggered_by=triggered_by,
                            started_at=now, last_progress_at=now, size_bytes=0,
                            manifest={"sent": 0, "failed": 0, "failures": [], "cursor": "",
                                      "chunk": token})
            session.add(row)
            session.commit()
            run_id = row.id
        else:
            why = _claim_run(session, run_id, "running", token, after=after, resume=resume)
            if why:
                return {"backup_run_id": run_id, "skipped": why}

    try:
        with SessionLocal() as session:
            settings = ImportSettings.load_sync(session)
            target = backup_destination.saved_target(settings)
            limit_mbps = settings.backup_upload_limit_mbps
        if target is None:
            raise backup_s3.S3Error("No off-site destination is set up.")
        more = _send_media_chunk(SessionLocal, run_id, target, backup_s3.Pacer(limit_mbps),
                                 token)
    except (SoftTimeLimitExceeded, Exception) as exc:
        with SessionLocal() as session:
            row = session.get(BackupRun, run_id)
            if row is not None:
                _mark_failed(session, row, exc)
        raise

    if more is None:
        return {"backup_run_id": run_id, "skipped": "another chunk took the run over"}
    if more:
        backup_media_task.delay(run_id=run_id, after=token)
        return {"backup_run_id": run_id, "continues": True}
    with SessionLocal() as session:
        row = session.get(BackupRun, run_id)
        m = dict(row.manifest or {})
        row.status = "ok"
        row.finished_at = datetime.now(UTC)
        if m.get("failed"):
            first = (m.get("failures") or [{}])[0]
            row.error = (f"{m['failed']} file(s) weren't sent; the next run tries them again. "
                         f"First: {first.get('path')}: {first.get('error')}")[:2000]
        session.commit()
    return {"backup_run_id": run_id, "continues": False}


def _send_media_chunk(SessionLocal, run_id: int, target, pacer, token: str) -> bool | None:
    """Send files until the chunk's time is up or none are left. Progress is
    committed after every file, so a chunk that dies loses at most the file
    in flight. Returns True when files remain for another chunk, and None when
    another chunk has taken the run over."""
    tid = backup_media.target_id(target)
    deadline = time.monotonic() + MEDIA_CHUNK_SECONDS
    with SessionLocal() as session:
        row = session.get(BackupRun, run_id)
        m = dict(row.manifest or {})
        m["offsite"] = _offsite_manifest(target)
        cursor = m.get("cursor") or ""
        while True:
            batch = backup_media.pending(session, tid, after=cursor)
            # End the read before any upload starts: a transaction left open
            # across a long transfer holds back vacuum for its whole length.
            session.commit()
            if not batch:
                if not _still_ours(session, run_id, token):
                    return None
                m["cursor"] = cursor
                row.manifest = dict(m)
                row.last_progress_at = datetime.now(UTC)
                session.commit()
                return False
            for blob in batch:
                if time.monotonic() > deadline:
                    if not _still_ours(session, run_id, token):
                        return None
                    m["cursor"] = cursor
                    row.manifest = dict(m)
                    session.commit()
                    return True
                key = None
                try:
                    key = backup_media.send(target, blob, IMAGES_ROOT, pacer)
                except (backup_s3.S3Error, OSError) as exc:
                    m["failed"] = m.get("failed", 0) + 1
                    if len(m.get("failures") or []) < MEDIA_FAILURES_LISTED:
                        m["failures"] = [*(m.get("failures") or []),
                                         {"path": blob.path, "error": str(exc)[:300]}]
                if not _still_ours(session, run_id, token):
                    return None
                if key is not None:
                    backup_media.record(session, tid, [(blob, key)])
                    m["sent"] = m.get("sent", 0) + 1
                    row.size_bytes = (row.size_bytes or 0) + blob.size_bytes
                cursor = blob.sha256
                m["cursor"] = cursor
                row.manifest = dict(m)
                row.last_progress_at = datetime.now(UTC)
                session.commit()


@celery.task(
    name="backend.app.tasks.backup.reconcile_media_backup",
    soft_time_limit=1800, time_limit=2100,
)
def reconcile_media_backup() -> dict:
    """Weekly Beat. Compare what the off-site bucket holds under `media/` with
    what FabledCurator believes it sent: a missing object is sent again by
    the next run, and an object already there is not sent twice. Skipped
    while a media run is going, and when no destination is set up.

    Each check is kept as a run of kind "reconcile" with its counts
    ({listed, dropped, adopted}), so the Backups card lists it and the System
    view can warn when files were missing from the bucket."""
    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        target = backup_destination.saved_target(ImportSettings.load_sync(session))
        if target is None:
            return {"skipped": "no destination"}
        going = session.execute(
            select(BackupRun.id).where(BackupRun.kind == "media")
            .where(BackupRun.status == "running").limit(1)
        ).scalar_one_or_none()
        if going is not None:
            return {"skipped": "a media backup is running"}
        now = datetime.now(UTC)
        row = BackupRun(kind="reconcile", status="running", triggered_by="weekly",
                        started_at=now, last_progress_at=now,
                        manifest={"offsite": _offsite_manifest(target)})
        session.add(row)
        session.commit()
        run_id = row.id
        try:
            counts = backup_media.reconcile(session, target)
        except (SoftTimeLimitExceeded, Exception) as exc:
            session.rollback()
            row = session.get(BackupRun, run_id)
            if row is not None:
                _mark_failed(session, row, exc)
            raise
        row = session.get(BackupRun, run_id)
        row.status = "ok"
        row.finished_at = datetime.now(UTC)
        row.manifest = {**(row.manifest or {}), **counts}
        session.commit()
        return {"backup_run_id": run_id, **counts}


# ── Restoring from the off-site bucket ───────────────────────────────────────

# The file phase fetches for this long per task, then queues its next chunk,
# like a media run.
RESTORE_CHUNK_SECONDS = 45 * 60
RESTORE_RUNS_KEPT = 20
RESTORE_FAILURES_LISTED = 20
# A run over files already on disk notes its place this often, not per file.
_RESTORE_NOTE_SECONDS = 5
_RESTORE_COUNTS = ("restored", "copied", "present", "missing", "failed")


def backup_busy(session) -> str | None:
    """Why a backup or a restore can't start right now, or None."""
    row = session.execute(
        select(BackupRun.kind).where(BackupRun.status.in_(["running", "restoring"])).limit(1)
    ).first()
    if row is None:
        return None
    if row.kind == "restore":
        return "a restore is going"
    if row.kind == "test":
        return "a test restore is running"
    return f"a {row.kind} backup is running"


def restore_going(session) -> bool:
    return session.execute(
        select(BackupRun.id).where(BackupRun.kind == "restore")
        .where(BackupRun.status == "restoring").limit(1)
    ).first() is not None


# ── Test restore ─────────────────────────────────────────────────────────────

TEST_RUNS_KEPT = 20


@celery.task(
    name="backend.app.tasks.backup.restore_test_task",
    bind=True,
    max_retries=0,
    # Downloading a dump, loading it and rebuilding its indexes; a restore of
    # the database alone, with room to spare.
    soft_time_limit=4 * 3600, time_limit=4 * 3600 + 900,
)
def restore_test_task(self, *, dump_key: str | None = None,
                      source_backup_run_id: int | None = None) -> dict:
    """Test a database dump without touching this install: load it into a
    temporary database, check every table against the counts recorded with
    the dump, bring it up to this version, and drop it again
    (services/backup_verify.py). Tests the bucket's dump at `dump_key` (the
    newest when None), or, given `source_backup_run_id`, a dump on this
    server. The run ends "ok" when every table matched, "error" with what
    didn't otherwise. Returns {'backup_run_id': N, ...}."""
    SessionLocal = _sync_session_factory()
    cfg = get_config()
    now = datetime.now(UTC)
    with SessionLocal() as session:
        if why := restore_test_blocked(session):
            return {"skipped": why}
        row = BackupRun(kind="test", status="running", triggered_by="manual",
                        started_at=now, last_progress_at=now, manifest={"phase": "fetching"})
        session.add(row)
        session.commit()
        run_id = row.id

    def note(phase: str, **more) -> None:
        with SessionLocal() as session:
            r = session.get(BackupRun, run_id)
            _note(session, r, {**r.manifest, "phase": phase, **more})

    staged = offsite = None
    try:
        if source_backup_run_id is not None:
            dump_path, recorded, source = _local_dump(SessionLocal, source_backup_run_id)
        else:
            with SessionLocal() as session:
                settings = ImportSettings.load_sync(session)
                target = backup_destination.saved_target(settings)
                limit_mbps = settings.backup_upload_limit_mbps
            if target is None:
                raise backup_s3.S3Error("No off-site destination is set up.")
            dump = backup_restore.find_dump(target, dump_key)
            source = {"key": dump.key, "ts": dump.ts, "size_bytes": dump.size_bytes}
            offsite = {**_offsite_manifest(target), "key": dump.key}
            note("fetching", dump=source, offsite=offsite)
            recorded = backup_verify.fetch_counts(target, dump.key)
            staged = backup_restore.fetch_dump(
                target, dump, IMAGES_ROOT / "_backups", backup_s3.Pacer(limit_mbps))
            dump_path = staged
        note("loading", dump=source)
        result = backup_verify.trial_restore(
            db_url=cfg.database_url_sync, dump_path=dump_path, recorded=recorded,
            note=note)
    except (SoftTimeLimitExceeded, Exception) as exc:
        with SessionLocal() as session:
            r = session.get(BackupRun, run_id)
            if r is not None:
                _mark_failed(session, r, exc)
        raise
    finally:
        if staged is not None:
            staged.unlink(missing_ok=True)

    with SessionLocal() as session:
        r = session.get(BackupRun, run_id)
        r.status = "error" if result["mismatched"] else "ok"
        r.error = backup_verify.summary(result) if result["mismatched"] else None
        r.finished_at = datetime.now(UTC)
        r.size_bytes = source.get("size_bytes")
        r.manifest = {"phase": "done", "dump": source, "offsite": offsite,
                      "summary": backup_verify.summary(result), **result}
        session.commit()
    return {"backup_run_id": run_id, "mismatched": result["mismatched"]}


def _local_dump(SessionLocal, run_id: int) -> tuple[Path, dict | None, dict]:
    """A dump on this server: its path, its recorded counts, and what the
    test run says it tested."""
    with SessionLocal() as session:
        src = session.get(BackupRun, run_id)
        if src is None or src.kind != "db" or src.status != "ok" or not src.sql_path:
            raise backup_s3.S3Error(f"Backup {run_id} isn't a database dump on this server.")
        path = Path(src.sql_path)
        if not path.is_file():
            raise backup_s3.S3Error(f"{path} is no longer on this server.")
        recorded = backup_verify.local_counts((src.manifest or {}).get("manifest_path"))
        ts = (src.manifest or {}).get("ts")
        return path, recorded, {"run_id": run_id, "ts": ts, "size_bytes": src.size_bytes}


def restore_test_blocked(session) -> str | None:
    """Why a test restore can't start right now, or None. A backup going
    doesn't stop one: the test only reads this install's database."""
    if restore_going(session):
        return "a restore is going"
    going = session.execute(
        select(BackupRun.id).where(BackupRun.kind == "test")
        .where(BackupRun.status == "running").limit(1)
    ).first()
    return "a test restore is already running" if going else None


def end_cut_off_tests() -> int:
    """End every test restore a restart cut short, and drop the temporary
    database it left. Called when the long-maintenance lane starts, where the
    tests run, so none is alive then."""
    with _sync_session_factory()() as session:
        ended = session.execute(
            update(BackupRun).where(BackupRun.kind == "test")
            .where(BackupRun.status == "running")
            .values(status="error", finished_at=datetime.now(UTC),
                    error="Interrupted: the server restarted during the test. Test again.")
        ).rowcount or 0
        session.commit()
    if ended:
        backup_verify.drop_scratch(get_config().database_url_sync)
    return ended


def _library_is_empty(session) -> bool:
    return session.execute(text(
        "SELECT NOT EXISTS (SELECT 1 FROM image_record) "
        "AND NOT EXISTS (SELECT 1 FROM post_attachment)"
    )).scalar()


def _note(session, row: BackupRun, m: dict, token: str | None = None) -> bool:
    """Write the run's progress. Given a chunk's token, writes only while that
    chunk still carries the run, and returns False when it doesn't."""
    if token is not None and not _still_ours(session, row.id, token):
        return False
    row.manifest = dict(m)
    row.last_progress_at = datetime.now(UTC)
    session.commit()
    return True


@celery.task(
    name="backend.app.tasks.backup.restore_offsite_task",
    bind=True,
    max_retries=0,  # a half-applied restore is never retried on its own
    soft_time_limit=21600, time_limit=23400,
)
def restore_offsite_task(self, *, dump_key: str | None = None, files_only: bool = False,
                         run_id: int | None = None, after: str | None = None,
                         resume: bool = False) -> dict:
    """Restore this install from the off-site bucket, one chunk at a time.

    Started without `run_id`, it replaces the database with the dump at
    `dump_key` (the newest when None), then fetches every library file the
    restored rows name. With `files_only`, the database stays and only the
    missing files are fetched, which also repairs a library that lost files.
    Each chunk queues the next, and a restart resumes the file phase where it
    stopped; the run's row counts files restored from the bucket, copied from
    a duplicate already on disk, already present, missing from the bucket,
    and failed. Returns {'backup_run_id': N, ...}."""
    SessionLocal = _sync_session_factory()
    token = _chunk_token(self)
    if run_id is not None:
        return _restore_files(SessionLocal, run_id, token, after=after, resume=resume)

    now = datetime.now(UTC)
    with SessionLocal() as session:
        busy = backup_busy(session)
        if busy:
            return {"skipped": busy}
        settings = ImportSettings.load_sync(session)
        limit_mbps = settings.backup_upload_limit_mbps
        row = BackupRun(kind="restore", status="restoring", triggered_by="restore",
                        started_at=now, last_progress_at=now, size_bytes=0,
                        manifest={"phase": "files" if files_only else "database",
                                  "cursor": "", "failures": [],
                                  **dict.fromkeys(_RESTORE_COUNTS, 0)})
        session.add(row)
        session.commit()
        run_id = row.id
        try:
            target = backup_destination.saved_target(settings)
            if target is None:
                raise backup_s3.S3Error("No off-site destination is set up.")
        except backup_s3.S3Error as exc:
            _mark_failed(session, row, exc)
            raise
        row.manifest = {**row.manifest, "offsite": _offsite_manifest(target)}
        session.commit()

    if files_only:
        return _restore_files(SessionLocal, run_id, token, after=None, resume=False)
    run_id = _restore_database(SessionLocal, run_id, target, limit_mbps, dump_key)
    # The files follow as chunks of their own, on the worker.
    restore_offsite_task.delay(run_id=run_id)
    return {"backup_run_id": run_id, "continues": True}


def _restore_database(SessionLocal, run_id: int, target, limit_mbps, dump_key) -> int:
    """Replace the database with a dump from the bucket. Returns the run's id
    in the restored database, where its row is written again."""
    cfg = get_config()
    pacer = backup_s3.Pacer(limit_mbps)
    staged = None
    try:
        dump = backup_restore.find_dump(target, dump_key)
        with SessionLocal() as session:
            row = session.get(BackupRun, run_id)
            m = dict(row.manifest)
            m["dump"] = {"key": dump.key, "ts": dump.ts, "size_bytes": dump.size_bytes}
            # The database about to be replaced goes to the bucket first, so
            # the restore can itself be undone. An empty one isn't worth a
            # dump, and would sit in the list as the newest.
            if not _library_is_empty(session):
                session.commit()
                saved = _dump_offsite(cfg, target, limit_mbps)
                m["saved_before"] = saved["key"]
            _note(session, row, m)
            started_at, manifest = row.started_at, dict(m)
        staged = backup_restore.fetch_dump(target, dump, IMAGES_ROOT / "_backups", pacer)
    except (SoftTimeLimitExceeded, Exception) as exc:
        if staged is not None:
            staged.unlink(missing_ok=True)
        with SessionLocal() as session:
            row = session.get(BackupRun, run_id)
            if row is not None:
                _mark_failed(session, row, exc)
        raise

    # From here on the database is being replaced; every pooled connection
    # still points at the tables being dropped.
    get_sync_engine().dispose()
    try:
        backup_service.restore_db(db_url=cfg.database_url_sync, sql_path=staged)
        backup_restore.migrate()
        backup_service.end_other_connections(db_url=cfg.database_url_sync)
    except (SoftTimeLimitExceeded, Exception) as exc:
        _leave_empty_install(SessionLocal, cfg, target, limit_mbps, started_at, manifest, exc)
        raise
    finally:
        staged.unlink(missing_ok=True)
    get_sync_engine().dispose()

    with SessionLocal() as session:
        backup_restore.keep_destination(session, target, limit_mbps)
        # The dump holds its own run as still going: it is the dump just
        # restored, so it finished. Any other run going then is over.
        dump_at = datetime.strptime(manifest["dump"]["ts"], "%Y%m%dT%H%M%SZ").replace(tzinfo=UTC)
        own = session.execute(
            select(BackupRun).where(BackupRun.kind == "db", BackupRun.status == "running")
            # The name's time is truncated to the second, after the row began.
            .where(BackupRun.started_at.between(dump_at - timedelta(minutes=2),
                                                dump_at + timedelta(seconds=1)))
            .order_by(BackupRun.started_at.desc()).limit(1)
        ).scalar_one_or_none()
        if own is not None:
            own.status = "ok"
            own.manifest = {"ts": manifest["dump"]["ts"], "offsite": {
                **manifest["offsite"], "key": manifest["dump"]["key"]}}
            own.size_bytes = manifest["dump"]["size_bytes"]
        session.execute(
            update(BackupRun).where(BackupRun.status.in_(["running", "restoring"]))
            .values(status="error", finished_at=datetime.now(UTC),
                    error="Interrupted: this run was going when the restored dump was taken.")
        )
        row = BackupRun(kind="restore", status="restoring", triggered_by="restore",
                        started_at=started_at, last_progress_at=datetime.now(UTC),
                        size_bytes=0, manifest={**manifest, "phase": "files"})
        session.add(row)
        session.commit()
        return row.id


def _leave_empty_install(SessionLocal, cfg, target, limit_mbps, started_at, manifest,
                         exc: BaseException) -> None:
    """After a failed load, put back an empty database at this version, so the
    install starts and can try again, and record why it stopped."""
    try:
        backup_service.wipe_db(db_url=cfg.database_url_sync)
        backup_restore.migrate()
        backup_service.end_other_connections(db_url=cfg.database_url_sync)
        get_sync_engine().dispose()
        with SessionLocal() as session:
            backup_restore.keep_destination(session, target, limit_mbps)
            session.add(BackupRun(
                kind="restore", status="error", triggered_by="restore",
                started_at=started_at, finished_at=datetime.now(UTC), manifest=manifest,
                error=(f"The restore stopped while loading the database: {exc}. "
                       "The database is empty again; choose a dump and restore again.")[:2000],
            ))
            session.commit()
    except Exception:  # noqa: BLE001 — the original failure is the one raised
        log.exception("restore: the database couldn't be reset after a failed load")


def _restore_files(SessionLocal, run_id: int, token: str, *, after: str | None,
                   resume: bool) -> dict:
    """One chunk of a restore's file phase; queues the next while files remain."""
    with SessionLocal() as session:
        why = _claim_run(session, run_id, "restoring", token, after=after, resume=resume)
        if why:
            return {"backup_run_id": run_id, "skipped": why}
        row = session.get(BackupRun, run_id)
        settings = ImportSettings.load_sync(session)
        limit_mbps = settings.backup_upload_limit_mbps
        try:
            target = backup_destination.saved_target(settings)
            if target is None:
                raise backup_s3.S3Error("No off-site destination is set up.")
        except backup_s3.S3Error as exc:
            _mark_failed(session, row, exc)
            raise
    try:
        more = _fetch_files_chunk(SessionLocal, run_id, target, backup_s3.Pacer(limit_mbps),
                                  token)
    except (SoftTimeLimitExceeded, Exception) as exc:
        with SessionLocal() as session:
            row = session.get(BackupRun, run_id)
            if row is not None:
                _mark_failed(session, row, exc)
        raise

    if more is None:
        return {"backup_run_id": run_id, "skipped": "another chunk took the run over"}
    if more:
        restore_offsite_task.delay(run_id=run_id, after=token)
        return {"backup_run_id": run_id, "continues": True}
    with SessionLocal() as session:
        row = session.get(BackupRun, run_id)
        m = dict(row.manifest or {})
        row.status = "restored"
        row.finished_at = datetime.now(UTC)
        missed = m.get("missing", 0) + m.get("failed", 0)
        if missed:
            first = (m.get("failures") or [{}])[0]
            row.error = (f"{m.get('missing', 0)} file(s) aren't in the bucket and "
                         f"{m.get('failed', 0)} couldn't be restored. "
                         f"First: {first.get('path')}: {first.get('error')}")[:2000]
        session.commit()
    # Thumbnails aren't backed up: they are made again from the restored files.
    from .thumbnail import backfill_thumbnails
    backfill_thumbnails.delay()
    return {"backup_run_id": run_id, "continues": False}


def _fetch_files_chunk(SessionLocal, run_id: int, target, pacer, token: str) -> bool | None:
    """Put library files back until the chunk's time is up or none are left.
    Returns True when files remain for another chunk, and None when another
    chunk has taken the run over."""
    tid = backup_media.target_id(target)
    deadline = time.monotonic() + RESTORE_CHUNK_SECONDS
    with SessionLocal() as session:
        row = session.get(BackupRun, run_id)
        m = dict(row.manifest or {})
        cursor = m.get("cursor") or ""
        noted = time.monotonic()
        while True:
            batch = backup_restore.library_files(session, after=cursor)
            keys = backup_restore.sent_keys(session, tid, [f.sha256 for f in batch])
            # End the read before any download starts.
            session.commit()
            if not batch:
                m["cursor"] = cursor
                return False if _note(session, row, m, token) else None
            for f in batch:
                if time.monotonic() > deadline:
                    m["cursor"] = cursor
                    return True if _note(session, row, m, token) else None
                try:
                    outcome = backup_restore.restore_file(
                        target, f, keys.get(f.sha256), IMAGES_ROOT, pacer)
                except (backup_s3.S3Error, OSError) as exc:
                    outcome = "missing" if isinstance(exc, backup_s3.ObjectMissing) else "failed"
                    if len(m.get("failures") or []) < RESTORE_FAILURES_LISTED:
                        m["failures"] = [*(m.get("failures") or []),
                                         {"path": f.paths[0], "error": str(exc)[:300]}]
                m[outcome] = m.get(outcome, 0) + 1
                if outcome in ("restored", "copied"):
                    row.size_bytes = (row.size_bytes or 0) + f.size_bytes
                cursor = f.sha256
                m["cursor"] = cursor
                if outcome != "present" or time.monotonic() - noted > _RESTORE_NOTE_SECONDS:
                    if not _note(session, row, m, token):
                        return None
                    noted = time.monotonic()