backend/core/backup_manager.py

"""
Database Backup Manager for AegisLM SaaS Backend.

Production-ready backup system with integrity verification,
automated scheduling, and restore capabilities.
"""

import asyncio
import hashlib
import json
import subprocess
import tempfile
from datetime import datetime, timedelta
from pathlib import Path
from typing import List, Dict, Optional, Any, Tuple
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession
import logging
import gzip
import shutil
import os

from .database import async_engine, get_redis
from .config import settings

logger = logging.getLogger(__name__)


class BackupMetadata:
    """Backup metadata information."""
    
    def __init__(self, backup_id: str, timestamp: datetime, checksum: str, 
                 size_bytes: int, tables_count: int, records_count: int):
        self.backup_id = backup_id
        self.timestamp = timestamp
        self.checksum = checksum
        self.size_bytes = size_bytes
        self.tables_count = tables_count
        self.records_count = records_count
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "backup_id": self.backup_id,
            "timestamp": self.timestamp.isoformat(),
            "checksum": self.checksum,
            "size_bytes": self.size_bytes,
            "tables_count": self.tables_count,
            "records_count": self.records_count
        }


class BackupManager:
    """Database backup and restore manager."""
    
    def __init__(self):
        self.backup_dir = Path(settings.BACKUP_DIR if hasattr(settings, 'BACKUP_DIR') else "./backups")
        self.backup_dir.mkdir(exist_ok=True)
        self.retention_days = getattr(settings, 'BACKUP_RETENTION_DAYS', 30)
        self.compression_enabled = getattr(settings, 'BACKUP_COMPRESSION', True)
        
    def _generate_backup_id(self) -> str:
        """Generate unique backup ID."""
        timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        return f"backup_{timestamp}"
    
    def _calculate_file_checksum(self, file_path: Path) -> str:
        """Calculate SHA-256 checksum of file."""
        hash_sha256 = hashlib.sha256()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_sha256.update(chunk)
        return hash_sha256.hexdigest()
    
    async def _get_database_stats(self) -> Tuple[int, int]:
        """Get database statistics (tables count, records count)."""
        async with async_engine.begin() as conn:
            # Get tables count
            tables_result = await conn.execute(text("""
                SELECT COUNT(*) FROM information_schema.tables 
                WHERE table_schema = 'public' AND table_type = 'BASE TABLE'
            """))
            tables_count = tables_result.scalar()
            
            # Get total records count
            records_result = await conn.execute(text("""
                SELECT SUM(n_tup_ins - n_tup_del) as total_records
                FROM pg_stat_user_tables
            """))
            records_count = records_result.scalar() or 0
            
            return tables_count, records_count
    
    async def create_backup(self, description: Optional[str] = None) -> Optional[BackupMetadata]:
        """Create a full database backup."""
        backup_id = self._generate_backup_id()
        temp_file = None
        
        try:
            # Get database statistics before backup
            tables_count, records_count = await self._get_database_stats()
            
            # Create temporary file for backup
            temp_file = tempfile.NamedTemporaryFile(suffix='.sql', delete=False)
            temp_path = Path(temp_file.name)
            
            # Extract database connection details
            db_url = settings.DATABASE_URL
            if not db_url.startswith('postgresql'):
                logger.error("Only PostgreSQL backups are supported")
                return None
            
            # Parse connection string
            # Format: postgresql+asyncpg://user:pass@host:port/dbname
            clean_url = db_url.replace('postgresql+asyncpg://', 'postgresql://')
            
            # Use pg_dump for backup
            cmd = [
                'pg_dump',
                '--verbose',
                '--no-password',
                '--format=custom',
                '--compress=9',
                clean_url
            ]
            
            logger.info(f"Creating backup {backup_id}...")
            
            # Run pg_dump
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=temp_file,
                stderr=asyncio.subprocess.PIPE
            )
            
            stdout, stderr = await process.communicate()
            
            if process.returncode != 0:
                logger.error(f"Backup failed: {stderr.decode()}")
                return None
            
            # Get file size
            file_size = temp_path.stat().st_size
            
            # Calculate checksum
            checksum = self._calculate_file_checksum(temp_path)
            
            # Move to backup directory
            backup_filename = f"{backup_id}.dump"
            backup_path = self.backup_dir / backup_filename
            
            if self.compression_enabled:
                # Compress the backup
                with open(temp_path, 'rb') as f_in:
                    with gzip.open(f"{backup_path}.gz", 'wb') as f_out:
                        shutil.copyfileobj(f_in, f_out)
                backup_path = Path(f"{backup_path}.gz")
                temp_path.unlink()  # Remove uncompressed file
            else:
                shutil.move(temp_path, backup_path)
            
            # Create metadata
            metadata = BackupMetadata(
                backup_id=backup_id,
                timestamp=datetime.utcnow(),
                checksum=checksum,
                size_bytes=backup_path.stat().st_size,
                tables_count=tables_count,
                records_count=records_count
            )
            
            # Save metadata
            metadata_path = self.backup_dir / f"{backup_id}_metadata.json"
            with open(metadata_path, 'w') as f:
                json.dump(metadata.to_dict(), f, indent=2)
            
            # Store backup info in Redis for quick access
            redis_client = await get_redis()
            await redis_client.setex(
                f"backup:{backup_id}",
                timedelta(days=self.retention_days),
                json.dumps(metadata.to_dict())
            )
            
            logger.info(f"Backup {backup_id} created successfully")
            return metadata
            
        except Exception as e:
            logger.error(f"Backup creation failed: {e}")
            if temp_file and temp_file.name:
                Path(temp_file.name).unlink(missing_ok=True)
            return None
    
    async def restore_backup(self, backup_id: str) -> bool:
        """Restore database from backup."""
        try:
            # Find backup file
            backup_path = None
            for path in self.backup_dir.glob(f"{backup_id}*"):
                if path.suffix in ['.dump', '.gz']:
                    backup_path = path
                    break
            
            if not backup_path:
                logger.error(f"Backup file not found: {backup_id}")
                return False
            
            # Verify backup integrity
            metadata = await self.get_backup_metadata(backup_id)
            if not metadata:
                logger.error(f"Backup metadata not found: {backup_id}")
                return False
            
            # Check checksum
            if backup_path.suffix == '.gz':
                # Decompress for checksum verification
                with gzip.open(backup_path, 'rb') as f:
                    temp_content = f.read()
                temp_checksum = hashlib.sha256(temp_content).hexdigest()
            else:
                temp_checksum = self._calculate_file_checksum(backup_path)
            
            if temp_checksum != metadata.checksum:
                logger.error(f"Backup integrity check failed for {backup_id}")
                return False
            
            # Create temporary file for restore
            temp_file = tempfile.NamedTemporaryFile(suffix='.dump', delete=False)
            temp_path = Path(temp_file.name)
            
            # Decompress if needed
            if backup_path.suffix == '.gz':
                with gzip.open(backup_path, 'rb') as f_in:
                    with open(temp_path, 'wb') as f_out:
                        shutil.copyfileobj(f_in, f_out)
            else:
                shutil.copy2(backup_path, temp_path)
            
            # Extract database connection details
            db_url = settings.DATABASE_URL
            clean_url = db_url.replace('postgresql+asyncpg://', 'postgresql://')
            
            # Use pg_restore for restore
            cmd = [
                'pg_restore',
                '--verbose',
                '--no-password',
                '--clean',
                '--if-exists',
                '--dbname', clean_url,
                str(temp_path)
            ]
            
            logger.info(f"Restoring backup {backup_id}...")
            
            # Run pg_restore
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            stdout, stderr = await process.communicate()
            
            # Clean up temp file
            temp_path.unlink()
            
            if process.returncode != 0:
                logger.error(f"Restore failed: {stderr.decode()}")
                return False
            
            logger.info(f"Backup {backup_id} restored successfully")
            return True
            
        except Exception as e:
            logger.error(f"Backup restore failed: {e}")
            return False
    
    async def get_backup_metadata(self, backup_id: str) -> Optional[BackupMetadata]:
        """Get backup metadata."""
        # Try Redis first
        try:
            redis_client = await get_redis()
            metadata_json = await redis_client.get(f"backup:{backup_id}")
            if metadata_json:
                data = json.loads(metadata_json)
                return BackupMetadata(
                    backup_id=data['backup_id'],
                    timestamp=datetime.fromisoformat(data['timestamp']),
                    checksum=data['checksum'],
                    size_bytes=data['size_bytes'],
                    tables_count=data['tables_count'],
                    records_count=data['records_count']
                )
        except Exception:
            pass
        
        # Try file system
        metadata_path = self.backup_dir / f"{backup_id}_metadata.json"
        if metadata_path.exists():
            try:
                with open(metadata_path, 'r') as f:
                    data = json.load(f)
                return BackupMetadata(
                    backup_id=data['backup_id'],
                    timestamp=datetime.fromisoformat(data['timestamp']),
                    checksum=data['checksum'],
                    size_bytes=data['size_bytes'],
                    tables_count=data['tables_count'],
                    records_count=data['records_count']
                )
            except Exception as e:
                logger.error(f"Failed to load metadata: {e}")
        
        return None
    
    async def list_backups(self) -> List[BackupMetadata]:
        """List all available backups."""
        backups = []
        
        # Scan backup directory
        for metadata_file in self.backup_dir.glob("*_metadata.json"):
            backup_id = metadata_file.stem.replace('_metadata', '')
            metadata = await self.get_backup_metadata(backup_id)
            if metadata:
                backups.append(metadata)
        
        # Sort by timestamp (newest first)
        backups.sort(key=lambda x: x.timestamp, reverse=True)
        return backups
    
    async def verify_backup_integrity(self, backup_id: str) -> Dict[str, Any]:
        """Verify backup file integrity."""
        metadata = await self.get_backup_metadata(backup_id)
        if not metadata:
            return {"verified": False, "error": "Backup metadata not found"}
        
        # Find backup file
        backup_path = None
        for path in self.backup_dir.glob(f"{backup_id}*"):
            if path.suffix in ['.dump', '.gz']:
                backup_path = path
                break
        
        if not backup_path:
            return {"verified": False, "error": "Backup file not found"}
        
        try:
            # Calculate current checksum
            if backup_path.suffix == '.gz':
                with gzip.open(backup_path, 'rb') as f:
                    content = f.read()
                current_checksum = hashlib.sha256(content).hexdigest()
            else:
                current_checksum = self._calculate_file_checksum(backup_path)
            
            # Compare with stored checksum
            is_valid = current_checksum == metadata.checksum
            
            return {
                "verified": is_valid,
                "stored_checksum": metadata.checksum,
                "current_checksum": current_checksum,
                "file_size": backup_path.stat().st_size,
                "expected_size": metadata.size_bytes,
                "size_match": backup_path.stat().st_size == metadata.size_bytes
            }
            
        except Exception as e:
            return {"verified": False, "error": str(e)}
    
    async def cleanup_old_backups(self) -> int:
        """Clean up old backups beyond retention period."""
        cutoff_date = datetime.utcnow() - timedelta(days=self.retention_days)
        cleaned_count = 0
        
        try:
            backups = await self.list_backups()
            
            for backup in backups:
                if backup.timestamp < cutoff_date:
                    # Remove backup files
                    for path in self.backup_dir.glob(f"{backup.backup_id}*"):
                        path.unlink()
                    
                    # Remove from Redis
                    redis_client = await get_redis()
                    await redis_client.delete(f"backup:{backup.backup_id}")
                    
                    cleaned_count += 1
                    logger.info(f"Cleaned up old backup: {backup.backup_id}")
            
            if cleaned_count > 0:
                logger.info(f"Cleaned up {cleaned_count} old backups")
            
        except Exception as e:
            logger.error(f"Backup cleanup failed: {e}")
        
        return cleaned_count
    
    async def get_backup_statistics(self) -> Dict[str, Any]:
        """Get backup statistics."""
        backups = await self.list_backups()
        
        if not backups:
            return {
                "total_backups": 0,
                "total_size_bytes": 0,
                "oldest_backup": None,
                "newest_backup": None,
                "average_size_bytes": 0
            }
        
        total_size = sum(b.size_bytes for b in backups)
        oldest = min(b.timestamp for b in backups)
        newest = max(b.timestamp for b in backups)
        
        return {
            "total_backups": len(backups),
            "total_size_bytes": total_size,
            "total_size_mb": round(total_size / (1024 * 1024), 2),
            "oldest_backup": oldest.isoformat(),
            "newest_backup": newest.isoformat(),
            "average_size_bytes": total_size // len(backups),
            "average_size_mb": round(total_size / len(backups) / (1024 * 1024), 2),
            "retention_days": self.retention_days
        }


# Global backup manager instance
backup_manager = BackupManager()


# Scheduled backup task
async def scheduled_backup():
    """Run scheduled backup."""
    logger.info("Starting scheduled backup...")
    
    metadata = await backup_manager.create_backup("Scheduled backup")
    if metadata:
        logger.info(f"Scheduled backup completed: {metadata.backup_id}")
        
        # Clean up old backups
        cleaned = await backup_manager.cleanup_old_backups()
        if cleaned > 0:
            logger.info(f"Cleaned up {cleaned} old backups")
    else:
        logger.error("Scheduled backup failed")


# Health check for backup system
async def backup_health_check() -> Dict[str, Any]:
    """Check backup system health."""
    try:
        # Check if we can create a small test backup
        test_backup = await backup_manager.create_backup("Health check test")
        if test_backup:
            # Verify the backup
            verification = await backup_manager.verify_backup_integrity(test_backup.backup_id)
            
            # Clean up test backup
            for path in backup_manager.backup_dir.glob(f"{test_backup.backup_id}*"):
                path.unlink()
            
            redis_client = await get_redis()
            await redis_client.delete(f"backup:{test_backup.backup_id}")
            
            return {
                "healthy": verification["verified"],
                "test_backup_id": test_backup.backup_id,
                "verification": verification
            }
        else:
            return {"healthy": False, "error": "Failed to create test backup"}
            
    except Exception as e:
        return {"healthy": False, "error": str(e)}


if __name__ == "__main__":
    import sys
    
    async def main():
        command = sys.argv[1] if len(sys.argv) > 1 else "help"
        
        if command == "create":
            metadata = await backup_manager.create_backup()
            if metadata:
                print(f"✅ Backup created: {metadata.backup_id}")
            else:
                print("❌ Backup creation failed")
        
        elif command == "list":
            backups = await backup_manager.list_backups()
            print(f"Total backups: {len(backups)}")
            for backup in backups:
                print(f"  - {backup.backup_id} ({backup.timestamp.isoformat()})")
        
        elif command == "verify":
            backup_id = sys.argv[2] if len(sys.argv) > 2 else None
            if not backup_id:
                print("Error: verify requires backup ID")
                sys.exit(1)
            
            verification = await backup_manager.verify_backup_integrity(backup_id)
            if verification["verified"]:
                print(f"✅ Backup {backup_id} is valid")
            else:
                print(f"❌ Backup {backup_id} verification failed")
                print(f"Error: {verification.get('error', 'Unknown error')}")
        
        elif command == "restore":
            backup_id = sys.argv[2] if len(sys.argv) > 2 else None
            if not backup_id:
                print("Error: restore requires backup ID")
                sys.exit(1)
            
            success = await backup_manager.restore_backup(backup_id)
            if success:
                print(f"✅ Backup {backup_id} restored successfully")
            else:
                print(f"❌ Backup {backup_id} restore failed")
        
        elif command == "stats":
            stats = await backup_manager.get_backup_statistics()
            print(f"Total backups: {stats['total_backups']}")
            print(f"Total size: {stats['total_size_mb']} MB")
            print(f"Average size: {stats['average_size_mb']} MB")
            print(f"Oldest backup: {stats['oldest_backup']}")
            print(f"Newest backup: {stats['newest_backup']}")
            print(f"Retention days: {stats['retention_days']}")
        
        elif command == "cleanup":
            cleaned = await backup_manager.cleanup_old_backups()
            print(f"✅ Cleaned up {cleaned} old backups")
        
        elif command == "health":
            health = await backup_health_check()
            if health["healthy"]:
                print("✅ Backup system is healthy")
            else:
                print("❌ Backup system health check failed")
                print(f"Error: {health.get('error', 'Unknown error')}")
        
        else:
            print("Usage: python backup_manager.py <command> [args]")
            print("Commands: create, list, verify <id>, restore <id>, stats, cleanup, health")
    
    asyncio.run(main())