app/backup.py

"""
Backup Module — 3-Layer Backup Strategy (FIXED)
────────────────────────────────────────────────
Layer 1: Litestream → local file replica (every 1s, automatic)
Layer 2: Python sqlite3.backup() → snapshot to /tmp (on-demand + scheduled)
Layer 3: Upload snapshot to HF Bucket (persistent, survives restarts)

Sources:
  - Python backup API: https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.backup
  - SQLite Backup API:  https://sqlite.org/backup.html
  - Litestream:         https://litestream.io/how-it-works/
  - HF Hub API:         https://huggingface.co/docs/huggingface_hub
  - SQLite VACUUM INTO: https://www.sqlite.org/lang_vacuum.html
"""

import sqlite3
import os
import gzip
import shutil
import logging
from datetime import datetime, timedelta
from pathlib import Path
from huggingface_hub import HfApi, hf_hub_download, create_repo, repo_exists
from contextlib import contextmanager

logger = logging.getLogger(__name__)

DB_PATH = "/tmp/data/app.db"
BACKUP_DIR = "/tmp/data/backups"
SNAPSHOT_DIR = "/tmp/data/snapshots"
RESTORE_DIR = "/tmp/data/restore"

HF_TOKEN = os.environ.get("HF_TOKEN", "")
HF_BUCKET_REPO = os.environ.get("HF_BUCKET_REPO", "")  # "username/repo-name"

# Ensure all directories exist
os.makedirs(BACKUP_DIR, exist_ok=True)
os.makedirs(SNAPSHOT_DIR, exist_ok=True)
os.makedirs(RESTORE_DIR, exist_ok=True)
os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)


# ═══════════════════════════════════════════════════════
#  LAYER 2: Python sqlite3.backup() — On-Demand Snapshot
# ═══════════════════════════════════════════════════════

def create_snapshot(label: str = "manual") -> dict:
    """
    Create an atomic snapshot of the database using Python's backup API.
    
    This uses SQLite's Online Backup API under the hood — it's safe to call
    while the database is being written to. The backup is atomic and consistent.
    
    Source: https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.backup
    """
    if not os.path.exists(DB_PATH):
        logger.warning(f"Database not found at {DB_PATH}, creating empty database first")
        # Create empty database with schema
        from app.database import init_db
        init_db()
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    snapshot_name = f"app_{label}_{timestamp}.db"
    snapshot_path = os.path.join(SNAPSHOT_DIR, snapshot_name)
    
    try:
        # ── Atomic backup using Python's built-in backup API ──
        source = sqlite3.connect(DB_PATH)
        dest = sqlite3.connect(snapshot_path)
        
        source.backup(dest)
        
        dest.close()
        source.close()
        
        file_size = os.path.getsize(snapshot_path)
        
        # Log the backup
        _log_backup("snapshot", snapshot_path, file_size, "success")
        
        logger.info(f"✅ Snapshot created: {snapshot_name} ({file_size} bytes)")
        
        return {
            "status": "success",
            "file": snapshot_name,
            "path": snapshot_path,
            "size_bytes": file_size,
            "size_mb": round(file_size / (1024 * 1024), 2),
            "timestamp": timestamp,
        }
    except Exception as e:
        _log_backup("snapshot", snapshot_path, 0, "error", str(e))
        logger.error(f"❌ Snapshot failed: {e}")
        return {"status": "error", "message": str(e)}


def create_compressed_snapshot() -> dict:
    """
    Create a gzip-compressed snapshot.
    SQLite databases compress very well (often 60-80% reduction).
    
    Source: https://litestream.io/alternatives/cron/
    """
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    snapshot_name = f"app_compressed_{timestamp}.db.gz"
    gz_path = os.path.join(SNAPSHOT_DIR, snapshot_name)
    
    # Create uncompressed snapshot first (temp)
    result = create_snapshot(label="temp_for_compress")
    if result["status"] != "success":
        return result
    
    temp_snapshot_path = result["path"]
    original_size = result["size_bytes"]
    
    try:
        # Compress directly to final location
        with open(temp_snapshot_path, 'rb') as f_in:
            with gzip.open(gz_path, 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        
        # Remove temp uncompressed file
        if os.path.exists(temp_snapshot_path):
            os.remove(temp_snapshot_path)
        
        gz_size = os.path.getsize(gz_path)
        
        _log_backup("compressed_snapshot", gz_path, gz_size, "success")
        
        logger.info(f"✅ Compressed snapshot created: {snapshot_name} ({gz_size} bytes, {round((1 - gz_size/original_size)*100, 1)}% reduction)")
        
        return {
            "status": "success",
            "file": snapshot_name,
            "path": gz_path,  # This is now the .gz file
            "size_bytes": original_size,
            "size_mb": round(original_size / (1024 * 1024), 2),
            "compressed_size_bytes": gz_size,
            "compressed_size_mb": round(gz_size / (1024 * 1024), 2),
            "compression_ratio": round((1 - gz_size / original_size) * 100, 1),
        }
    except Exception as e:
        # Clean up on failure
        if os.path.exists(temp_snapshot_path):
            os.remove(temp_snapshot_path)
        if os.path.exists(gz_path):
            os.remove(gz_path)
        _log_backup("compressed_snapshot", gz_path, 0, "error", str(e))
        logger.error(f"❌ Compressed snapshot failed: {e}")
        return {"status": "error", "message": str(e)}


def create_vacuum_snapshot() -> dict:
    """
    Use VACUUM INTO for a fully compacted snapshot.
    
    This creates a snapshot from a single transaction — best for high-write scenarios.
    Also removes any free pages, giving you the smallest possible backup.
    
    Source: https://sqlite.org/lang_vacuum.html
    """
    if not os.path.exists(DB_PATH):
        return {"status": "error", "message": "Database not found"}
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    snapshot_name = f"app_vacuum_{timestamp}.db"
    snapshot_path = os.path.join(SNAPSHOT_DIR, snapshot_name)
    
    try:
        conn = sqlite3.connect(DB_PATH)
        conn.execute(f"VACUUM INTO '{snapshot_path}'")
        conn.close()
        
        file_size = os.path.getsize(snapshot_path)
        _log_backup("vacuum_snapshot", snapshot_path, file_size, "success")
        
        return {
            "status": "success",
            "file": snapshot_name,
            "path": snapshot_path,
            "size_bytes": file_size,
            "size_mb": round(file_size / (1024 * 1024), 2),
        }
    except Exception as e:
        return {"status": "error", "message": str(e)}


# ═══════════════════════════════════════════════════════
#  LAYER 3: HF Bucket Sync — Persistent Remote Storage
# ═══════════════════════════════════════════════════════

def ensure_hf_repo_exists() -> dict:
    """
    Ensure the HF dataset repository exists.
    Creates it if it doesn't exist.
    
    Returns status dict.
    """
    if not HF_TOKEN:
        return {"status": "error", "message": "HF_TOKEN not set"}
    
    if not HF_BUCKET_REPO:
        return {"status": "error", "message": "HF_BUCKET_REPO not set"}
    
    try:
        api = HfApi(token=HF_TOKEN)
        
        # Check if repo exists
        try:
            if repo_exists(repo_id=HF_BUCKET_REPO, repo_type="dataset", token=HF_TOKEN):
                logger.info(f"✅ HF repo exists: {HF_BUCKET_REPO}")
                return {"status": "success", "message": "Repo exists"}
        except Exception:
            pass
        
        # Create repo if it doesn't exist
        logger.info(f"Creating HF dataset repo: {HF_BUCKET_REPO}")
        create_repo(
            repo_id=HF_BUCKET_REPO,
            repo_type="dataset",
            private=True,
            exist_ok=True,
            token=HF_TOKEN
        )
        logger.info(f"✅ Created HF dataset repo: {HF_BUCKET_REPO}")
        return {"status": "success", "message": "Repo created"}
        
    except Exception as e:
        logger.error(f"❌ Failed to ensure HF repo exists: {e}")
        return {"status": "error", "message": str(e)}


def upload_to_hf_bucket(local_path: str = None) -> dict:
    """
    Upload a database snapshot to HuggingFace Bucket.
    
    If no local_path given, creates a fresh compressed snapshot first.
    
    Source: https://huggingface.co/docs/huggingface_hub/guides/upload
    """
    # Check credentials
    if not HF_TOKEN or not HF_BUCKET_REPO:
        msg = f"HF_TOKEN={'set' if HF_TOKEN else 'NOT SET'}, HF_BUCKET_REPO={'set' if HF_BUCKET_REPO else 'NOT SET'}"
        logger.error(f"❌ HF credentials missing: {msg}")
        return {"status": "error", "message": f"Credentials not configured: {msg}"}
    
    # Ensure repo exists
    repo_status = ensure_hf_repo_exists()
    if repo_status["status"] != "success":
        return repo_status
    
    # Create snapshot if not provided
    if local_path is None:
        result = create_compressed_snapshot()
        if result["status"] != "success":
            return result
        local_path = result["path"]  # Now correctly points to .gz file
    
    # Verify file exists
    if not os.path.exists(local_path):
        return {"status": "error", "message": f"Local file not found: {local_path}"}
    
    filename = os.path.basename(local_path)
    file_size = os.path.getsize(local_path)
    
    try:
        api = HfApi(token=HF_TOKEN)
        
        logger.info(f"📤 Uploading to HF Bucket: {filename} ({file_size} bytes)")
        
        # Upload latest (overwrites)
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo="db-backups/latest.db.gz",
            repo_id=HF_BUCKET_REPO,
            repo_type="dataset",
        )
        
        # Also keep a timestamped copy (rolling archive)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=f"db-backups/archive/{filename}",
            repo_id=HF_BUCKET_REPO,
            repo_type="dataset",
        )
        
        _log_backup("hf_bucket", f"db-backups/{filename}", file_size, "success")
        
        logger.info(f"✅ Uploaded to HF Bucket: {filename}")
        return {"status": "success", "remote_path": f"db-backups/{filename}", "size_bytes": file_size}
        
    except Exception as e:
        error_msg = str(e)
        _log_backup("hf_bucket", local_path, 0, "error", error_msg)
        logger.error(f"❌ HF Bucket upload failed: {error_msg}")
        return {"status": "error", "message": error_msg}


def restore_from_hf_bucket() -> dict:
    """
    Restore database from HF Bucket.
    Called on startup by entrypoint.sh if local DB doesn't exist.
    
    Source: https://huggingface.co/docs/huggingface_hub/guides/download
    """
    if not HF_TOKEN or not HF_BUCKET_REPO:
        logger.warning("⚠️ HF credentials not set, skipping restore")
        return {"status": "skipped", "message": "No HF credentials configured"}
    
    if os.path.exists(DB_PATH) and os.path.getsize(DB_PATH) > 0:
        logger.info("✅ Database already exists, skipping restore")
        return {"status": "skipped", "message": "DB already exists"}
    
    # Ensure restore directory exists
    os.makedirs(RESTORE_DIR, exist_ok=True)
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    
    try:
        logger.info(f"📥 Attempting to restore from HF Bucket: {HF_BUCKET_REPO}")
        
        # Download latest backup
        local_gz = hf_hub_download(
            repo_id=HF_BUCKET_REPO,
            filename="db-backups/latest.db.gz",
            repo_type="dataset",
            local_dir=RESTORE_DIR,
            token=HF_TOKEN,
        )
        
        logger.info(f"✅ Downloaded backup: {local_gz}")
        
        # Decompress
        with gzip.open(local_gz, 'rb') as f_in:
            with open(DB_PATH, 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
        
        # Verify integrity
        conn = sqlite3.connect(DB_PATH)
        result = conn.execute("PRAGMA integrity_check").fetchone()
        conn.close()
        
        if result[0] == "ok":
            logger.info("✅ Database restored and verified from HF Bucket")
            return {"status": "success", "integrity": "ok"}
        else:
            logger.error(f"❌ Restored database failed integrity: {result[0]}")
            os.remove(DB_PATH)  # Remove corrupted restore
            return {"status": "error", "message": f"Integrity check failed: {result[0]}"}
            
    except Exception as e:
        error_msg = str(e)
        # Check if it's a "file not found" error (first run)
        if "404" in error_msg or "not found" in error_msg.lower():
            logger.info("📦 No backup found in HF Bucket (this is normal for first run)")
            return {"status": "not_found", "message": "No backup exists yet - this is normal for first run"}
        else:
            logger.error(f"❌ Restore from HF Bucket failed: {error_msg}")
            return {"status": "error", "message": error_msg}


# ═══════════════════════════════════════════════════════
#  SCHEDULED BACKUP (runs inside FastAPI)
# ═══════════════════════════════════════════════════════

async def scheduled_backup_task():
    """
    Background task that runs every 30 minutes.
    Creates snapshot + uploads to HF Bucket.
    """
    import asyncio
    
    # Wait a bit before first backup (let app stabilize)
    await asyncio.sleep(60)
    
    while True:
        try:
            logger.info("⏰ Running scheduled backup...")
            result = upload_to_hf_bucket()
            if result["status"] == "success":
                logger.info("✅ Scheduled backup completed successfully")
            else:
                logger.warning(f"⚠️ Scheduled backup completed with issues: {result.get('message', 'unknown')}")
            
            cleanup_old_snapshots(max_age_hours=48)
        except Exception as e:
            logger.error(f"❌ Scheduled backup error: {e}")
        
        await asyncio.sleep(1800)  # 30 minutes


def cleanup_old_snapshots(max_age_hours: int = 48):
    """Remove local snapshots older than max_age_hours."""
    cutoff = datetime.now() - timedelta(hours=max_age_hours)
    cleaned = 0
    for f in Path(SNAPSHOT_DIR).glob("*.db*"):
        try:
            if datetime.fromtimestamp(f.stat().st_mtime) < cutoff:
                f.unlink()
                cleaned += 1
                logger.info(f"🗑️ Removed old snapshot: {f.name}")
        except Exception as e:
            logger.warning(f"Failed to remove {f.name}: {e}")
    
    if cleaned > 0:
        logger.info(f"🗑️ Cleaned up {cleaned} old snapshots")


def list_backups() -> dict:
    """List all available backups (local + remote info)."""
    local_files = []
    
    for f in sorted(Path(SNAPSHOT_DIR).glob("*.db*"), key=lambda x: x.stat().st_mtime, reverse=True):
        try:
            local_files.append({
                "name": f.name,
                "size_bytes": f.stat().st_size,
                "size_mb": round(f.stat().st_size / (1024 * 1024), 2),
                "modified": datetime.fromtimestamp(f.stat().st_mtime).isoformat(),
            })
        except Exception:
            pass
    
    return {
        "local_snapshots": local_files, 
        "snapshot_dir": SNAPSHOT_DIR,
        "hf_bucket_repo": HF_BUCKET_REPO if HF_BUCKET_REPO else "Not configured",
        "hf_token_set": bool(HF_TOKEN),
    }


# ── Internal helper ──
def _log_backup(backup_type, file_path, file_size, status, error=None):
    """Log backup event to database."""
    try:
        conn = sqlite3.connect(DB_PATH)
        conn.execute(
            "INSERT INTO backup_log (backup_type, file_path, file_size, status, error_message) VALUES (?, ?, ?, ?, ?)",
            (backup_type, file_path, file_size, status, error)
        )
        conn.commit()
        conn.close()
    except Exception as e:
        logger.warning(f"Failed to log backup event: {e}")


# ═══════════════════════════════════════════════════════
#  DIAGNOSTICS
# ═══════════════════════════════════════════════════════

def diagnose_hf_setup() -> dict:
    """
    Diagnose HF setup issues.
    Run this to check if everything is configured correctly.
    """
    results = {
        "checks": [],
        "overall_status": "unknown"
    }
    
    # Check HF_TOKEN
    if HF_TOKEN:
        results["checks"].append({
            "name": "HF_TOKEN",
            "status": "✅ SET",
            "value": f"{HF_TOKEN[:8]}...{HF_TOKEN[-4:]}" if len(HF_TOKEN) > 12 else "***"
        })
    else:
        results["checks"].append({
            "name": "HF_TOKEN",
            "status": "❌ NOT SET",
            "value": "Environment variable HF_TOKEN is missing"
        })
    
    # Check HF_BUCKET_REPO
    if HF_BUCKET_REPO:
        results["checks"].append({
            "name": "HF_BUCKET_REPO",
            "status": "✅ SET",
            "value": HF_BUCKET_REPO
        })
    else:
        results["checks"].append({
            "name": "HF_BUCKET_REPO",
            "status": "❌ NOT SET",
            "value": "Environment variable HF_BUCKET_REPO is missing"
        })
    
    # Try to verify token by accessing API
    if HF_TOKEN:
        try:
            api = HfApi(token=HF_TOKEN)
            user_info = api.whoami()
            results["checks"].append({
                "name": "Token Validation",
                "status": "✅ VALID",
                "value": f"Logged in as: {user_info.get('name', 'unknown')}"
            })
        except Exception as e:
            results["checks"].append({
                "name": "Token Validation",
                "status": "❌ INVALID",
                "value": str(e)
            })
    
    # Check if repo exists/is accessible
    if HF_TOKEN and HF_BUCKET_REPO:
        try:
            exists = repo_exists(repo_id=HF_BUCKET_REPO, repo_type="dataset", token=HF_TOKEN)
            results["checks"].append({
                "name": "Repository Check",
                "status": "✅ EXISTS" if exists else "⚠️ WILL BE CREATED",
                "value": HF_BUCKET_REPO
            })
        except Exception as e:
            results["checks"].append({
                "name": "Repository Check",
                "status": "❌ ERROR",
                "value": str(e)
            })
    
    # Determine overall status
    failed = [c for c in results["checks"] if "❌" in c["status"]]
    if not failed:
        results["overall_status"] = "✅ All checks passed"
    else:
        results["overall_status"] = f"❌ {len(failed)} check(s) failed"
    
    return results