app/tasks/backup_task.py

"""Author RAG Chatbot SaaS — Automated Backup Task.

Scheduled database and ChromaDB backups with retention policies.
Run via: python -m app.tasks.backup_task
Or schedule via cron: 0 2 * * * (daily at 2 AM)
"""

import hashlib
import os
import shutil
from datetime import datetime, timezone
from pathlib import Path

import structlog

from app.config import get_settings
from app.tasks.celery_app import celery_app

logger = structlog.get_logger(__name__)
cfg = get_settings()

BACKUP_DIR = Path("/data/backups")
RETENTION = {
    "daily": 7,    # Keep last 7 daily backups
    "weekly": 4,   # Keep last 4 weekly backups
    "monthly": 3,  # Keep last 3 monthly backups
}


def _sha256(path: Path) -> str:
    """Calculate SHA-256 checksum for integrity verification."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(8192):
            h.update(chunk)
    return h.hexdigest()


def backup_sqlite() -> Path | None:
    """Create a compressed backup of the SQLite database."""
    db_url = str(cfg.DATABASE_URL)
    if not db_url.startswith("sqlite"):
        logger.info("Not using SQLite, skipping SQLite backup")
        return None

    # Extract file path from URL
    db_path = db_url.split("///")[-1]
    if not os.path.exists(db_path):
        logger.warning("SQLite DB file not found", path=db_path)
        return None

    now = datetime.now(timezone.utc)
    backup_subdir = BACKUP_DIR / "sqlite" / now.strftime("%Y-%m-%d")
    backup_subdir.mkdir(parents=True, exist_ok=True)

    backup_file = backup_subdir / f"db_{now.strftime('%H%M%S')}.sqlite3"
    shutil.copy2(db_path, backup_file)

    # Write checksum
    checksum = _sha256(backup_file)
    (backup_subdir / f"db_{now.strftime('%H%M%S')}.sha256").write_text(checksum)

    size_mb = backup_file.stat().st_size / 1024 / 1024
    logger.info("SQLite backup created", path=str(backup_file), size_mb=f"{size_mb:.2f}", checksum=checksum[:16])
    return backup_file


def backup_chromadb() -> Path | None:
    """Backup ChromaDB persistent directory."""
    chroma_dir = Path(cfg.CHROMA_PERSIST_DIR)
    if not chroma_dir.exists():
        logger.warning("ChromaDB dir not found", path=str(chroma_dir))
        return None

    now = datetime.now(timezone.utc)
    backup_subdir = BACKUP_DIR / "chromadb" / now.strftime("%Y-%m-%d")
    backup_subdir.mkdir(parents=True, exist_ok=True)

    archive_path = backup_subdir / f"chroma_{now.strftime('%H%M%S')}"
    shutil.copytree(chroma_dir, archive_path, dirs_exist_ok=True)

    logger.info("ChromaDB backup created", path=str(archive_path))
    return archive_path


def _cleanup_old_backups(subdir: str) -> None:
    """Remove backups older than retention policy."""
    base = BACKUP_DIR / subdir
    if not base.exists():
        return

    dirs = sorted(base.iterdir(), reverse=True)
    max_keep = RETENTION["daily"]  # Keep last N daily backups

    for i, d in enumerate(dirs):
        if i >= max_keep and d.is_dir():
            shutil.rmtree(d)
            logger.info("Removed old backup", path=str(d))


@celery_app.task(name="app.tasks.backup_task.run_backup")
def run_backup() -> dict:
    """Run full backup cycle."""
    results = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "sqlite": None,
        "chromadb": None,
    }

    try:
        sqlite_path = backup_sqlite()
        results["sqlite"] = str(sqlite_path) if sqlite_path else "skipped"
    except Exception as e:
        logger.error("SQLite backup failed", error=str(e))
        results["sqlite"] = f"error: {e}"

    try:
        chroma_path = backup_chromadb()
        results["chromadb"] = str(chroma_path) if chroma_path else "skipped"
    except Exception as e:
        logger.error("ChromaDB backup failed", error=str(e))
        results["chromadb"] = f"error: {e}"

    # Cleanup old backups
    _cleanup_old_backups("sqlite")
    _cleanup_old_backups("chromadb")

    logger.info("Backup cycle complete", results=results)
    return results


if __name__ == "__main__":
    run_backup()