IOT_application/Machine-Backend/backup_db.py

"""
SQLite Database Backup Script
Usage: python backup_db.py [backup|restore|list]
"""

import os
import shutil
import sqlite3
from datetime import datetime
from pathlib import Path

# Configuration
INSTANCE_DIR = Path(__file__).parent / 'instance'
DB_FILE = INSTANCE_DIR / 'machines.db'
BACKUP_DIR = INSTANCE_DIR / 'backups'


def ensure_backup_dir():
    """Create backup directory if it doesn't exist"""
    BACKUP_DIR.mkdir(parents=True, exist_ok=True)


def backup_database():
    """Create a timestamped backup of the database"""
    if not DB_FILE.exists():
        print(f"❌ Database not found: {DB_FILE}")
        return False
    
    ensure_backup_dir()
    
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_file = BACKUP_DIR / f'machines_backup_{timestamp}.db'
    
    try:
        # Method 1: Simple file copy (faster)
        shutil.copy2(DB_FILE, backup_file)
        
        # Get file size
        size = backup_file.stat().st_size
        size_mb = size / (1024 * 1024)
        
        print(f"✅ Backup created successfully!")
        print(f"   File: {backup_file.name}")
        print(f"   Size: {size_mb:.2f} MB")
        print(f"   Location: {backup_file}")
        
        return True
        
    except Exception as e:
        print(f"❌ Backup failed: {e}")
        return False


def backup_database_with_integrity():
    """Create a backup using SQLite's built-in backup API (slower but safer)"""
    if not DB_FILE.exists():
        print(f"❌ Database not found: {DB_FILE}")
        return False
    
    ensure_backup_dir()
    
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_file = BACKUP_DIR / f'machines_backup_{timestamp}.db'
    
    try:
        # Connect to source database
        source_conn = sqlite3.connect(str(DB_FILE))
        
        # Connect to backup database
        backup_conn = sqlite3.connect(str(backup_file))
        
        # Perform backup
        with backup_conn:
            source_conn.backup(backup_conn)
        
        source_conn.close()
        backup_conn.close()
        
        # Get file size
        size = backup_file.stat().st_size
        size_mb = size / (1024 * 1024)
        
        print(f"✅ Backup created successfully (with integrity check)!")
        print(f"   File: {backup_file.name}")
        print(f"   Size: {size_mb:.2f} MB")
        print(f"   Location: {backup_file}")
        
        return True
        
    except Exception as e:
        print(f"❌ Backup failed: {e}")
        if backup_file.exists():
            backup_file.unlink()  # Delete partial backup
        return False


def list_backups():
    """List all available backups"""
    ensure_backup_dir()
    
    backups = sorted(BACKUP_DIR.glob('machines_backup_*.db'), reverse=True)
    
    if not backups:
        print("📂 No backups found")
        return
    
    print(f"\n📂 Available Backups ({len(backups)} total):")
    print("=" * 80)
    
    for i, backup in enumerate(backups, 1):
        size = backup.stat().st_size / (1024 * 1024)
        mtime = datetime.fromtimestamp(backup.stat().st_mtime)
        
        # Parse timestamp from filename
        try:
            parts = backup.stem.split('_')
            date_str = parts[-2]
            time_str = parts[-1]
            backup_date = datetime.strptime(f"{date_str}_{time_str}", "%Y%m%d_%H%M%S")
            date_display = backup_date.strftime("%Y-%m-%d %H:%M:%S")
        except:
            date_display = mtime.strftime("%Y-%m-%d %H:%M:%S")
        
        print(f"{i:2d}. {backup.name}")
        print(f"    Date: {date_display}")
        print(f"    Size: {size:.2f} MB")
        print()


def restore_database(backup_name=None):
    """Restore database from a backup"""
    ensure_backup_dir()
    
    backups = sorted(BACKUP_DIR.glob('machines_backup_*.db'), reverse=True)
    
    if not backups:
        print("❌ No backups found")
        return False
    
    # If no backup specified, show list and ask
    if not backup_name:
        list_backups()
        print("=" * 80)
        choice = input("Enter backup number to restore (or 'q' to quit): ").strip()
        
        if choice.lower() == 'q':
            print("❌ Restore cancelled")
            return False
        
        try:
            index = int(choice) - 1
            if index < 0 or index >= len(backups):
                print("❌ Invalid backup number")
                return False
            backup_file = backups[index]
        except ValueError:
            print("❌ Invalid input")
            return False
    else:
        # Find backup by name
        backup_file = BACKUP_DIR / backup_name
        if not backup_file.exists():
            print(f"❌ Backup not found: {backup_name}")
            return False
    
    # Confirm restore
    print(f"\n⚠️  WARNING: This will replace your current database!")
    print(f"   Current: {DB_FILE}")
    print(f"   Backup:  {backup_file.name}")
    confirm = input("\nType 'yes' to confirm restore: ").strip().lower()
    
    if confirm != 'yes':
        print("❌ Restore cancelled")
        return False
    
    try:
        # Create a safety backup of current database
        if DB_FILE.exists():
            safety_backup = BACKUP_DIR / f'machines_before_restore_{datetime.now().strftime("%Y%m%d_%H%M%S")}.db'
            shutil.copy2(DB_FILE, safety_backup)
            print(f"✅ Safety backup created: {safety_backup.name}")
        
        # Restore from backup
        shutil.copy2(backup_file, DB_FILE)
        
        print(f"✅ Database restored successfully!")
        print(f"   Restored from: {backup_file.name}")
        print(f"\n⚠️  Remember to restart your Flask server!")
        
        return True
        
    except Exception as e:
        print(f"❌ Restore failed: {e}")
        return False


def cleanup_old_backups(keep_count=10):
    """Keep only the most recent N backups"""
    ensure_backup_dir()
    
    backups = sorted(BACKUP_DIR.glob('machines_backup_*.db'), reverse=True)
    
    if len(backups) <= keep_count:
        print(f"✅ No cleanup needed (found {len(backups)} backups, keeping {keep_count})")
        return
    
    to_delete = backups[keep_count:]
    
    print(f"\n🗑️  Cleanup: Keeping {keep_count} most recent backups")
    print(f"   Deleting {len(to_delete)} old backups...")
    
    for backup in to_delete:
        try:
            backup.unlink()
            print(f"   ✓ Deleted: {backup.name}")
        except Exception as e:
            print(f"   ✗ Failed to delete {backup.name}: {e}")
    
    print(f"✅ Cleanup complete!")


def main():
    """Main entry point"""
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python backup_db.py [backup|restore|list|cleanup]")
        print("\nCommands:")
        print("  backup  - Create a new backup")
        print("  restore - Restore from a backup")
        print("  list    - List all available backups")
        print("  cleanup - Delete old backups (keep 10 most recent)")
        return
    
    command = sys.argv[1].lower()
    
    if command == 'backup':
        backup_database()
    elif command == 'backup-safe':
        backup_database_with_integrity()
    elif command == 'restore':
        backup_name = sys.argv[2] if len(sys.argv) > 2 else None
        restore_database(backup_name)
    elif command == 'list':
        list_backups()
    elif command == 'cleanup':
        keep = int(sys.argv[2]) if len(sys.argv) > 2 else 10
        cleanup_old_backups(keep)
    else:
        print(f"❌ Unknown command: {command}")
        print("   Use: backup, restore, list, or cleanup")


if __name__ == '__main__':
    main()