#!/usr/bin/env python3
"""
Restore script for StreamA2026
Restores a backup of PostgreSQL database, Redis, Caddy SSL, and configuration files.
"""
import os
import subprocess
import tarfile
import json
import sys
from pathlib import Path

# Configuration
BACKUP_DIR = os.getenv('BACKUP_DIR', './backups')
POSTGRES_HOST = os.getenv('POSTGRES_HOST', 'postgres')
POSTGRES_PORT = os.getenv('POSTGRES_PORT', '5432')
POSTGRES_DB = os.getenv('POSTGRES_DB', 'streama2026')
POSTGRES_USER = os.getenv('POSTGRES_USER', 'streama')
POSTGRES_PASSWORD = os.getenv('POSTGRES_PASSWORD', 'stream')
REDIS_HOST = os.getenv('REDIS_HOST', 'redis')
REDIS_PORT = os.getenv('REDIS_PORT', '6379')
REDIS_PASSWORD = os.getenv('REDIS_PASSWORD', 'redis')

# Create temporary directory in project
TEMP_DIR = Path('./temp_restore')
TEMP_DIR.mkdir(exist_ok=True)

def extract_backup(backup_filename):
    """Extract backup archive to temporary directory."""
    import shutil
    
    backup_path = Path(BACKUP_DIR) / backup_filename
    if not backup_path.exists():
        print(f"Error: Backup file not found: {backup_path}")
        return None
    
    # Clean and create temp directory
    if TEMP_DIR.exists():
        shutil.rmtree(TEMP_DIR)
    TEMP_DIR.mkdir(exist_ok=True)
    
    try:
        with tarfile.open(backup_path, 'r:gz') as tar:
            tar.extractall(path=TEMP_DIR)
        
        print(f"Backup extracted to {TEMP_DIR}")
        return str(TEMP_DIR)
    except Exception as e:
        print(f"Error extracting backup: {e}")
        return None

def read_metadata(temp_dir):
    """Read backup metadata."""
    metadata_path = Path(temp_dir) / 'backup_metadata.json'
    
    if not metadata_path.exists():
        print("Warning: backup_metadata.json not found, using default metadata")
        return {
            'backup_version': 'unknown',
            'flyway_version': 'unknown',
            'components': []
        }
    
    with open(metadata_path, 'r') as f:
        return json.load(f)

def restore_postgresql(temp_dir):
    """Restore PostgreSQL database from SQL dump."""
    try:
        # Find SQL file
        sql_files = list(Path(temp_dir).glob('postgres_backup_*.sql'))
        if not sql_files:
            print("Warning: No PostgreSQL backup file found")
            return False
        
        sql_file = sql_files[0]
        print(f"Restoring PostgreSQL from {sql_file}")
        
        # Copy SQL file to container and restore
        cmd = [
            'docker', 'cp',
            str(sql_file),
            'streama2026-postgres:/tmp/restore.sql'
        ]
        subprocess.run(cmd, check=True)
        
        # Restore database
        cmd = [
            'docker', 'exec', 'streama2026-postgres',
            'psql', '-U', POSTGRES_USER, '-d', POSTGRES_DB,
            '-f', '/tmp/restore.sql'
        ]
        subprocess.run(cmd, check=True)
        
        print("PostgreSQL restore completed successfully")
        return True
    except Exception as e:
        print(f"PostgreSQL restore failed: {e}")
        return False

def restore_redis(temp_dir):
    """Restore Redis database from RDB file."""
    try:
        # Find RDB file
        rdb_files = list(Path(temp_dir).glob('redis_backup_*.rdb'))
        if not rdb_files:
            print("Warning: No Redis backup file found")
            return False
        
        rdb_file = rdb_files[0]
        print(f"Restoring Redis from {rdb_file}")
        
        # Copy RDB file to container
        cmd = [
            'docker', 'cp',
            str(rdb_file),
            'streama2026-redis:/data/dump.rdb'
        ]
        subprocess.run(cmd, check=True)
        
        # Restart Redis to load the dump
        cmd = ['docker', 'restart', 'streama2026-redis']
        subprocess.run(cmd, check=True)
        
        print("Redis restore completed successfully")
        return True
    except Exception as e:
        print(f"Redis restore failed: {e}")
        return False

def restore_caddy(temp_dir):
    """Restore Caddy SSL certificates and configuration."""
    try:
        # Find Caddy backup file
        caddy_files = list(Path(temp_dir).glob('caddy_backup_*.tar.gz'))
        if not caddy_files:
            print("Warning: No Caddy backup file found")
            return False
        
        caddy_file = caddy_files[0]
        print(f"Restoring Caddy from {caddy_file}")
        
        # Extract Caddy backup
        caddy_temp = TEMP_DIR / 'caddy_temp'
        caddy_temp.mkdir(exist_ok=True)
        with tarfile.open(caddy_file, 'r:gz') as tar:
            tar.extractall(path=caddy_temp)
        
        # Restore Caddy data volume (copy to container, which writes to volume)
        if (caddy_temp / 'caddy_data').exists():
            cmd = [
                'docker', 'cp',
                str(caddy_temp / 'caddy_data'),
                'streama2026-caddy:/data'
            ]
            subprocess.run(cmd, check=True)
            print("Caddy data volume restored")
        
        # Restore Caddy config volume (copy to container, which writes to volume)
        if (caddy_temp / 'caddy_config').exists():
            cmd = [
                'docker', 'cp',
                str(caddy_temp / 'caddy_config'),
                'streama2026-caddy:/config'
            ]
            subprocess.run(cmd, check=True)
            print("Caddy config volume restored")
        
        # Restart Caddy
        cmd = ['docker', 'restart', 'streama2026-caddy']
        subprocess.run(cmd, check=True)
        
        print("Caddy restore completed successfully")
        return True
    except Exception as e:
        print(f"Caddy restore failed: {e}")
        return False

def restore_config(temp_dir):
    """Restore configuration files."""
    try:
        # Find config backup file
        config_files = list(Path(temp_dir).glob('config_backup_*.tar.gz'))
        if not config_files:
            print("Warning: No config backup file found")
            return False
        
        config_file = config_files[0]
        print(f"Restoring config from {config_file}")
        
        # Extract config backup
        config_temp = TEMP_DIR / 'config_temp'
        config_temp.mkdir(exist_ok=True)
        with tarfile.open(config_file, 'r:gz') as tar:
            tar.extractall(path=config_temp)
        
        # Restore .env file
        if (config_temp / '.env').exists():
            import shutil
            shutil.copy(config_temp / '.env', '.env')
            print(".env file restored")
        
        # Restore docker-compose.yml
        if (config_temp / 'docker-compose.yml').exists():
            import shutil
            shutil.copy(config_temp / 'docker-compose.yml', 'docker-compose.yml')
            print("docker-compose.yml restored")
        
        # Restore nginx config
        if (config_temp / 'nginx').exists():
            import shutil
            if Path('nginx').exists():
                shutil.rmtree('nginx')
            shutil.copytree(config_temp / 'nginx', 'nginx')
            print("nginx config restored")
        
        print("Config restore completed successfully")
        return True
    except Exception as e:
        print(f"Config restore failed: {e}")
        return False

def run_flyway_migrations():
    """Run Flyway migrations after restore."""
    try:
        print("Running Flyway migrations...")
        cmd = ['docker', 'exec', 'streama2026-flyway', 'flyway', 'migrate']
        subprocess.run(cmd, check=True)
        print("Flyway migrations completed successfully")
        return True
    except Exception as e:
        print(f"Flyway migrations failed: {e}")
        return False

def restore_backup(backup_filename):
    """Restore a complete backup."""
    print(f"Starting restore process for: {backup_filename}")
    
    # Extract backup
    temp_dir = extract_backup(backup_filename)
    if not temp_dir:
        print("Restore failed: Could not extract backup")
        return False
    
    # Read metadata
    metadata = read_metadata(temp_dir)
    print(f"Backup version: {metadata.get('backup_version', 'unknown')}")
    print(f"Flyway version: {metadata.get('flyway_version', 'unknown')}")
    print(f"Components: {metadata.get('components', [])}")
    
    # Restore components
    success = True
    
    # PostgreSQL
    if 'postgresql' in metadata.get('components', []):
        if not restore_postgresql(temp_dir):
            success = False
    
    # Redis
    if 'redis' in metadata.get('components', []):
        if not restore_redis(temp_dir):
            success = False
    
    # Caddy
    if 'caddy' in metadata.get('components', []):
        if not restore_caddy(temp_dir):
            success = False
    
    # Config
    if 'config' in metadata.get('components', []):
        if not restore_config(temp_dir):
            success = False
    
    # Run Flyway migrations
    if success:
        if not run_flyway_migrations():
            success = False
    
    # Cleanup
    import shutil
    if TEMP_DIR.exists():
        shutil.rmtree(TEMP_DIR)
    
    if success:
        print("Restore completed successfully")
        return True
    else:
        print("Restore completed with errors")
        return False

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: python restore.py <backup_filename>")
        sys.exit(1)
    
    backup_filename = sys.argv[1]
    result = restore_backup(backup_filename)
    
    if result:
        sys.exit(0)
    else:
        sys.exit(1)