pve-cloud-backups/backend/app/backup_service.py

from __future__ import annotations

import asyncio
import uuid
from datetime import datetime, timedelta, timezone
from pathlib import Path

from . import commands
from .db import db, row_to_dict, utc_now
from .jobs import advance_job, due_jobs
from .notifications import notify
from .settings_store import get_settings


ACTIVE_STATES = {
    "queued",
    "pve_running",
    "local_ready",
    "uploading",
    "remote_ready",
    "local_deleting",
    "deleting",
}
RUNNING_STATES = ACTIVE_STATES - {"queued"}
RECOVERABLE_STATES = (ACTIVE_STATES - {"queued"}) | {"local_downloading"}
TERMINAL_STATES = {"completed", "failed", "deleted"}


def _active_for_guest(conn, vmid: int) -> bool:
    row = conn.execute(
        f"SELECT 1 FROM backups WHERE guest_vmid = ? AND state IN ({','.join('?' for _ in ACTIVE_STATES)}) LIMIT 1",
        (vmid, *sorted(ACTIVE_STATES)),
    ).fetchone()
    return row is not None


def _backup_row(backup_id: str) -> dict:
    with db() as conn:
        row = conn.execute("SELECT * FROM backups WHERE id = ?", (backup_id,)).fetchone()
    if not row:
        raise KeyError(f"Backup {backup_id} not found")
    return dict(row)


def _job_row(job_id: int | None) -> dict | None:
    if job_id is None:
        return None
    with db() as conn:
        row = conn.execute("SELECT * FROM backup_jobs WHERE id = ?", (job_id,)).fetchone()
    return row_to_dict(row)


def update_backup(backup_id: str, **values) -> dict:
    values["updated_at"] = utc_now()
    assignments = ", ".join(f"{key} = ?" for key in values)
    with db() as conn:
        conn.execute(
            f"UPDATE backups SET {assignments} WHERE id = ?",
            (*values.values(), backup_id),
        )
        row = conn.execute("SELECT * FROM backups WHERE id = ?", (backup_id,)).fetchone()
    return dict(row)


def fail_backup(backup_id: str, error: str) -> dict:
    return update_backup(backup_id, state="failed", error_message=error, completed_at=utc_now())


def _claim_queued_backup(backup_id: str) -> dict | None:
    now = utc_now()
    with db() as conn:
        cur = conn.execute(
            "UPDATE backups SET state = 'pve_running', updated_at = ? WHERE id = ? AND state = 'queued'",
            (now, backup_id),
        )
        if cur.rowcount == 0:
            return None
        row = conn.execute("SELECT * FROM backups WHERE id = ?", (backup_id,)).fetchone()
    return dict(row)


def queue_backup_for_job(job: dict) -> dict | None:
    now = utc_now()
    backup_id = str(uuid.uuid4())
    settings = get_settings()
    max_active = max(1, int(settings["max_concurrent_backups"]))
    with db() as conn:
        active_count = conn.execute(
            f"SELECT COUNT(*) AS count FROM backups WHERE state IN ({','.join('?' for _ in ACTIVE_STATES)})",
            tuple(sorted(ACTIVE_STATES)),
        ).fetchone()["count"]
        if active_count >= max_active:
            return None
        if _active_for_guest(conn, int(job["guest_vmid"])):
            return None
        conn.execute(
            """
            INSERT INTO backups(
                id, job_id, guest_vmid, guest_name, guest_type, node, state,
                started_at, updated_at
            )
            VALUES (?, ?, ?, ?, ?, ?, 'queued', ?, ?)
            """,
            (
                backup_id,
                job["id"],
                job["guest_vmid"],
                job["guest_name"],
                job["guest_type"],
                job["node"],
                now,
                now,
            ),
        )
        row = conn.execute("SELECT * FROM backups WHERE id = ?", (backup_id,)).fetchone()
    return dict(row)


def list_backups(guest_vmid: int | None = None, state: str | None = None) -> list[dict]:
    query = "SELECT * FROM backups"
    args: list = []
    where = []
    if guest_vmid is not None:
        where.append("guest_vmid = ?")
        args.append(guest_vmid)
    if state:
        where.append("state = ?")
        args.append(state)
    if where:
        query += " WHERE " + " AND ".join(where)
    query += " ORDER BY started_at DESC LIMIT 500"
    with db() as conn:
        rows = [dict(row) for row in conn.execute(query, args).fetchall()]
    return [_reconcile_local_path(row) for row in rows]


def _reconcile_local_path(backup: dict) -> dict:
    """Self-heal a stale local_path left over from a file that no longer exists."""
    local_path = backup.get("local_path")
    if local_path and not Path(local_path).exists():
        return update_backup(backup["id"], local_path=None)
    return backup


def get_backup(backup_id: str) -> dict:
    return _reconcile_local_path(_backup_row(backup_id))


def purge_deleted_backup(backup_id: str) -> dict:
    backup = _backup_row(backup_id)
    if backup["state"] != "deleted":
        raise RuntimeError(f"Backup {backup_id} is in state {backup['state']} and is not ready to purge")
    with db() as conn:
        conn.execute("DELETE FROM backups WHERE id = ?", (backup_id,))
    return {"id": backup_id, "purged": True}


def dashboard() -> dict:
    with db() as conn:
        latest = row_to_dict(conn.execute("SELECT * FROM backups ORDER BY started_at DESC LIMIT 1").fetchone())
        active = [dict(row) for row in conn.execute(
            f"SELECT * FROM backups WHERE state IN ({','.join('?' for _ in ACTIVE_STATES)}) ORDER BY started_at DESC",
            tuple(sorted(ACTIVE_STATES)),
        ).fetchall()]
        failed = conn.execute("SELECT COUNT(*) AS count FROM backups WHERE state = 'failed'").fetchone()["count"]
        remote = conn.execute("SELECT COUNT(*) AS count FROM backups WHERE state = 'completed' AND remote_path IS NOT NULL").fetchone()["count"]
        next_jobs = [dict(row) for row in conn.execute(
            "SELECT * FROM backup_jobs WHERE enabled = 1 ORDER BY next_run_at ASC LIMIT 10"
        ).fetchall()]
    settings = get_settings()
    local_dir = Path(settings["local_backup_dir"])
    usage = {"path": str(local_dir), "bytes": 0}
    if local_dir.exists():
        usage["bytes"] = sum(p.stat().st_size for p in local_dir.glob("*") if p.is_file())
    return {
        "latest_backup": latest,
        "active_tasks": active,
        "failed_backups": failed,
        "remote_backup_count": remote,
        "local_staging_usage": usage,
        "next_scheduled_jobs": next_jobs,
    }


async def _wait_for_pve_task(backup: dict) -> dict:
    while True:
        status = await asyncio.to_thread(commands.get_task_status, backup["node"], backup["upid"])
        if status.get("status") == "stopped":
            exitstatus = status.get("exitstatus")
            if exitstatus == "OK":
                return status
            raise RuntimeError(f"Proxmox task failed with exitstatus={exitstatus}")
        await asyncio.sleep(10)


def _remote_object(settings: dict, backup: dict, archive: Path) -> str:
    remote = str(settings["rclone_remote"]).rstrip(":")
    base_path = str(settings.get("rclone_remote_path") or "pve-cloud-backup").strip("/")
    object_path = f"{backup['id']}/{archive.name}"
    if base_path:
        object_path = f"{base_path}/{object_path}"
    return f"{remote}:{object_path}"


def _remote_parent_dir(remote_object: str) -> str | None:
    if ":" not in remote_object:
        return None
    remote, object_path = remote_object.split(":", 1)
    object_path = object_path.strip("/")
    if "/" not in object_path:
        return None
    parent_path = object_path.rsplit("/", 1)[0]
    if not parent_path:
        return None
    return f"{remote}:{parent_path}"


async def _remove_empty_remote_parent(settings: dict, remote_object: str) -> None:
    parent = _remote_parent_dir(remote_object)
    if not parent:
        return
    try:
        await asyncio.to_thread(commands.rclone_rmdir, settings["rclone_path"], parent)
    except commands.CommandError:
        # Non-empty folders and backend-specific "directory not found" behavior
        # should not turn an already successful exact file delete into a failed
        # backup deletion.
        return


async def _delete_remote_file_and_empty_parent(settings: dict, remote_object: str) -> None:
    await asyncio.to_thread(commands.rclone_delete_file, settings["rclone_path"], remote_object)
    await _remove_empty_remote_parent(settings, remote_object)


async def process_backup(backup_id: str) -> None:
    backup = _backup_row(backup_id)
    settings = get_settings()
    job = _job_row(backup.get("job_id"))
    try:
        if backup["state"] == "queued":
            await notify("backup_started", backup=backup)
            backup = _claim_queued_backup(backup_id)
            if backup is None:
                return
            upid = await asyncio.to_thread(
                commands.start_proxmox_backup,
                node=backup["node"],
                vmid=int(backup["guest_vmid"]),
                storage=settings["proxmox_storage"],
                mode=(job or {}).get("backup_mode") or settings["default_backup_mode"],
                compression=(job or {}).get("compression") or settings["default_compression"],
            )
            backup = update_backup(backup_id, upid=upid)

        if backup["state"] == "pve_running":
            if not backup.get("upid"):
                raise RuntimeError("Backup is pve_running but has no Proxmox UPID recorded; refusing to poll task status")
            await _wait_for_pve_task(backup)
            backup = update_backup(backup_id, state="local_ready", pve_completed_at=utc_now())
            await notify("pve_completed", backup=backup)

        if backup["state"] == "local_ready":
            archive = await asyncio.to_thread(
                commands.discover_archive,
                settings["local_backup_dir"],
                int(backup["guest_vmid"]),
                backup["guest_type"],
                backup["started_at"],
            )
            size = archive.stat().st_size
            remote_object = _remote_object(settings, backup, archive)
            backup = update_backup(
                backup_id,
                local_path=str(archive),
                remote_path=remote_object,
                size_bytes=size,
                state="uploading",
                upload_started_at=utc_now(),
            )
            await notify("upload_started", backup=backup)

        if backup["state"] == "uploading":
            await asyncio.to_thread(
                commands.rclone_copyto,
                settings["rclone_path"],
                backup["local_path"],
                backup["remote_path"],
            )
            exists = await asyncio.to_thread(
                commands.rclone_object_exists,
                settings["rclone_path"],
                backup["remote_path"],
            )
            if not exists:
                raise RuntimeError("Upload finished but remote object was not found")
            backup = update_backup(backup_id, state="remote_ready")

        if backup["state"] == "remote_ready":
            backup = update_backup(backup_id, state="local_deleting")

        if backup["state"] == "local_deleting":
            if not settings.get("keep_local_backups") and backup.get("local_path"):
                local_path = Path(backup["local_path"])
                if local_path.exists():
                    local_path.unlink()
                backup = update_backup(backup_id, local_path=None)
            backup = update_backup(backup_id, state="completed", completed_at=utc_now())
            await notify("backup_completed", backup=backup)
            await apply_retention_for_job(backup["job_id"])
    except Exception as exc:
        failed = fail_backup(backup_id, str(exc))
        await notify("backup_failed", backup=failed)


async def delete_remote_backup(backup: dict, reason: str) -> dict:
    settings = get_settings()
    update_backup(backup["id"], state="deleting")
    if backup.get("remote_path"):
        exists = await asyncio.to_thread(
            commands.rclone_object_exists,
            settings["rclone_path"],
            backup["remote_path"],
        )
        if exists:
            await _delete_remote_file_and_empty_parent(settings, backup["remote_path"])
        else:
            await _remove_empty_remote_parent(settings, backup["remote_path"])
    if backup.get("local_path"):
        local_path = Path(backup["local_path"])
        if local_path.exists() and local_path.is_file():
            local_path.unlink()
    deleted = update_backup(backup["id"], state="deleted", deleted_at=utc_now(), local_path=None)
    await notify("remote_deleted", backup=deleted, message=reason)
    return deleted


async def delete_backup_manually(backup_id: str, reason: str = "manual deletion") -> dict:
    backup = _backup_row(backup_id)
    if backup["state"] == "deleted":
        return purge_deleted_backup(backup_id)
    if backup["state"] in ACTIVE_STATES:
        raise RuntimeError(f"Backup {backup_id} is active in state {backup['state']} and cannot be deleted")
    return await delete_remote_backup(backup, reason)


def _archive_name_from_remote(remote_path: str) -> str:
    return Path(remote_path.split(":", 1)[1]).name


async def download_backup_locally(backup_id: str) -> dict:
    backup = _backup_row(backup_id)
    if backup["state"] != "completed" or not backup.get("remote_path"):
        raise RuntimeError(f"Backup {backup_id} has no completed remote archive to download")
    if backup.get("local_path") and Path(backup["local_path"]).exists():
        return backup

    now = utc_now()
    with db() as conn:
        cur = conn.execute(
            "UPDATE backups SET state = 'local_downloading', updated_at = ? WHERE id = ? AND state = 'completed'",
            (now, backup_id),
        )
        if cur.rowcount == 0:
            raise RuntimeError(f"Backup {backup_id} is busy and cannot be downloaded right now")

    settings = get_settings()
    target = Path(settings["local_backup_dir"]) / _archive_name_from_remote(backup["remote_path"])
    try:
        await asyncio.to_thread(
            commands.rclone_download,
            settings["rclone_path"],
            backup["remote_path"],
            str(target),
        )
        if not target.exists():
            raise RuntimeError("Download finished but local archive was not found")
    except Exception:
        update_backup(backup_id, state="completed")
        raise

    downloaded = update_backup(backup_id, state="completed", local_path=str(target))
    await notify("local_downloaded", backup=downloaded)
    return downloaded


async def delete_local_copy(backup_id: str) -> dict:
    backup = _backup_row(backup_id)
    if backup["state"] != "completed":
        raise RuntimeError(f"Backup {backup_id} is in state {backup['state']} and has no local copy to delete")
    if not backup.get("local_path"):
        return backup
    local_path = Path(backup["local_path"])
    if local_path.exists() and local_path.is_file():
        local_path.unlink()
    updated = update_backup(backup_id, local_path=None)
    await notify("local_deleted", backup=updated)
    return updated


async def delete_job_with_backups(job_id: int) -> dict:
    with db() as conn:
        job = row_to_dict(conn.execute("SELECT * FROM backup_jobs WHERE id = ?", (job_id,)).fetchone())
        if not job:
            raise KeyError(f"Backup job {job_id} not found")
        active = conn.execute(
            f"SELECT COUNT(*) AS count FROM backups WHERE job_id = ? AND state IN ({','.join('?' for _ in ACTIVE_STATES)})",
            (job_id, *sorted(ACTIVE_STATES)),
        ).fetchone()["count"]
        if active:
            raise RuntimeError(f"Backup job {job_id} has {active} active backup(s) and cannot be deleted")
        backups = [
            dict(row)
            for row in conn.execute(
                "SELECT * FROM backups WHERE job_id = ? AND state != 'deleted' ORDER BY started_at ASC",
                (job_id,),
            ).fetchall()
        ]

    deleted_backup_ids: list[str] = []
    for backup in backups:
        deleted = await delete_backup_manually(backup["id"], reason=f"job {job_id} deleted")
        deleted_backup_ids.append(deleted["id"])

    with db() as conn:
        conn.execute("DELETE FROM backup_jobs WHERE id = ?", (job_id,))
    return {"job_id": job_id, "deleted_backups": deleted_backup_ids}


async def apply_retention_for_job(job_id: int | None) -> None:
    if job_id is None:
        return
    with db() as conn:
        job = row_to_dict(conn.execute("SELECT * FROM backup_jobs WHERE id = ?", (job_id,)).fetchone())
        backups = [dict(row) for row in conn.execute(
            """
            SELECT * FROM backups
            WHERE job_id = ? AND state = 'completed' AND remote_path IS NOT NULL
            ORDER BY completed_at DESC
            """,
            (job_id,),
        ).fetchall()]
    if not job:
        return
    delete: list[dict] = []
    if job["retention_type"] == "latest":
        delete = backups[int(job["retention_value"]):]
    else:
        cutoff = datetime.now(timezone.utc) - timedelta(days=int(job["retention_value"]))
        for backup in backups:
            completed_at = backup.get("completed_at")
            if completed_at and datetime.fromisoformat(completed_at) < cutoff:
                delete.append(backup)
    for backup in delete:
        try:
            await delete_remote_backup(backup, f"retention {job['retention_type']}={job['retention_value']}")
        except Exception as exc:
            failed = fail_backup(backup["id"], f"Retention delete failed: {exc}")
            await notify("backup_failed", backup=failed)


async def retention_sweep() -> None:
    with db() as conn:
        job_ids = [
            row["id"]
            for row in conn.execute(
                """
                SELECT DISTINCT backup_jobs.id
                FROM backup_jobs
                JOIN backups ON backups.job_id = backup_jobs.id
                WHERE backups.state = 'completed' AND backups.remote_path IS NOT NULL
                """
            ).fetchall()
        ]
    for job_id in job_ids:
        await apply_retention_for_job(job_id)


async def recover_unfinished() -> None:
    with db() as conn:
        rows = conn.execute(
            f"SELECT * FROM backups WHERE state IN ({','.join('?' for _ in RECOVERABLE_STATES)})",
            tuple(sorted(RECOVERABLE_STATES)),
        ).fetchall()
    for row in rows:
        backup = dict(row)
        await notify("recovery_action", backup=backup, message=f"resuming state {backup['state']}")
        if backup["state"] == "deleting":
            await delete_remote_backup(backup, "recovery delete resume")
        elif backup["state"] == "local_downloading":
            update_backup(backup["id"], state="completed")
            try:
                await download_backup_locally(backup["id"])
            except Exception as exc:
                await notify("recovery_action", backup=backup, message=f"resume download failed: {exc}")
        else:
            await process_backup(backup["id"])


async def scheduler_tick() -> None:
    settings = get_settings()
    limit = max(1, int(settings["max_concurrent_backups"]))
    with db() as conn:
        active_count = conn.execute(
            f"SELECT COUNT(*) AS count FROM backups WHERE state IN ({','.join('?' for _ in ACTIVE_STATES)})",
            tuple(sorted(ACTIVE_STATES)),
        ).fetchone()["count"]
    free_slots = max(0, limit - active_count)
    if free_slots <= 0:
        return
    for job in due_jobs(free_slots):
        queued = queue_backup_for_job(job)
        advance_job(job)
        if queued:
            asyncio.create_task(process_backup(queued["id"]))


async def worker_loop() -> None:
    await recover_unfinished()
    retention_sweep_interval_seconds = 3600
    last_retention_sweep = 0.0
    while True:
        await scheduler_tick()
        now = asyncio.get_running_loop().time()
        if now - last_retention_sweep >= retention_sweep_interval_seconds:
            await retention_sweep()
            last_retention_sweep = now
        settings = get_settings()
        max_active = max(1, int(settings["max_concurrent_backups"]))
        with db() as conn:
            active_count = conn.execute(
                f"SELECT COUNT(*) AS count FROM backups WHERE state IN ({','.join('?' for _ in RUNNING_STATES)})",
                tuple(sorted(RUNNING_STATES)),
            ).fetchone()["count"]
            free_slots = max(0, max_active - active_count)
            queued = [dict(row) for row in conn.execute(
                "SELECT * FROM backups WHERE state = 'queued' ORDER BY started_at ASC LIMIT ?",
                (free_slots,),
            ).fetchall()]
        for backup in queued:
            asyncio.create_task(process_backup(backup["id"]))
        await asyncio.sleep(30)