backend/core/data_archiver.py

"""
Data Archiving System for AegisLM SaaS Backend.

Production-ready data archiving with automated cleanup,
compression, and retrieval capabilities.
"""

import asyncio
import json
import gzip
import shutil
from datetime import datetime, timedelta
from pathlib import Path
from typing import List, Dict, Optional, Any, Tuple
from sqlalchemy import text, select, and_, func
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
import logging
import csv
import io

from .database import async_engine, get_sync_db
from .config import settings
from .backup_manager import backup_manager

logger = logging.getLogger(__name__)


class ArchivePolicy:
    """Data archiving policy configuration."""
    
    def __init__(self, table_name: str, retention_days: int, 
                 archive_condition: Optional[str] = None,
                 compress: bool = True,
                 batch_size: int = 1000):
        self.table_name = table_name
        self.retention_days = retention_days
        self.archive_condition = archive_condition
        self.compress = compress
        self.batch_size = batch_size
    
    def get_cutoff_date(self) -> datetime:
        """Get cutoff date for archiving."""
        return datetime.utcnow() - timedelta(days=self.retention_days)


class ArchivedRecord:
    """Archived record metadata."""
    
    def __init__(self, archive_id: str, table_name: str, record_count: int,
                 archive_date: datetime, file_path: str, compressed: bool):
        self.archive_id = archive_id
        self.table_name = table_name
        self.record_count = record_count
        self.archive_date = archive_date
        self.file_path = file_path
        self.compressed = compressed
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "archive_id": self.archive_id,
            "table_name": self.table_name,
            "record_count": self.record_count,
            "archive_date": self.archive_date.isoformat(),
            "file_path": str(self.file_path),
            "compressed": self.compressed
        }


class DataArchiver:
    """Data archiving management system."""
    
    def __init__(self):
        self.archive_dir = Path(getattr(settings, 'ARCHIVE_DIR', './archives'))
        self.archive_dir.mkdir(exist_ok=True)
        
        # Default archiving policies
        self.policies = {
            'evaluations': ArchivePolicy(
                table_name='evaluations',
                retention_days=getattr(settings, 'EVALUATION_RETENTION_DAYS', 365),
                archive_condition="status IN ('completed', 'failed', 'cancelled')",
                compress=True,
                batch_size=500
            ),
            'api_keys': ArchivePolicy(
                table_name='api_keys',
                retention_days=getattr(settings, 'API_KEY_RETENTION_DAYS', 730),
                archive_condition="is_active = false AND expires_at < NOW()",
                compress=True,
                batch_size=1000
            ),
        }
        
        self.redis_client = None
    
    async def get_redis(self):
        """Get Redis client."""
        if not self.redis_client:
            from .database import get_redis
            self.redis_client = await get_redis()
        return self.redis_client
    
    def _generate_archive_id(self, table_name: str) -> str:
        """Generate unique archive ID."""
        timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        return f"archive_{table_name}_{timestamp}"
    
    async def archive_table_data(self, table_name: str, policy: Optional[ArchivePolicy] = None) -> Optional[ArchivedRecord]:
        """Archive data from a specific table."""
        try:
            # Use provided policy or default
            if not policy:
                policy = self.policies.get(table_name)
                if not policy:
                    logger.error(f"No archiving policy found for table: {table_name}")
                    return None
            
            # Get cutoff date
            cutoff_date = policy.get_cutoff_date()
            
            # Build archive query
            base_query = f"FROM {table_name} WHERE created_at < :cutoff_date"
            if policy.archive_condition:
                base_query += f" AND {policy.archive_condition}"
            
            # Get count of records to archive
            count_query = f"SELECT COUNT(*) {base_query}"
            
            async with async_engine.begin() as conn:
                result = await conn.execute(text(count_query), {"cutoff_date": cutoff_date})
                record_count = result.scalar()
                
                if record_count == 0:
                    logger.info(f"No records to archive in {table_name}")
                    return None
                
                logger.info(f"Archiving {record_count} records from {table_name}")
                
                # Generate archive ID and file path
                archive_id = self._generate_archive_id(table_name)
                file_extension = '.csv.gz' if policy.compress else '.csv'
                file_path = self.archive_dir / f"{archive_id}{file_extension}"
                
                # Export data in batches
                await self._export_data_in_batches(
                    conn, table_name, base_query, cutoff_date, 
                    file_path, policy.batch_size, policy.compress
                )
                
                # Delete archived records
                delete_query = f"DELETE {base_query}"
                await conn.execute(text(delete_query), {"cutoff_date": cutoff_date})
                
                # Create archive record
                archived_record = ArchivedRecord(
                    archive_id=archive_id,
                    table_name=table_name,
                    record_count=record_count,
                    archive_date=datetime.utcnow(),
                    file_path=file_path,
                    compressed=policy.compress
                )
                
                # Record archive in Redis
                await self._record_archive(archived_record)
                
                logger.info(f"Successfully archived {record_count} records from {table_name}")
                return archived_record
        
        except Exception as e:
            logger.error(f"Failed to archive data from {table_name}: {e}")
            return None
    
    async def _export_data_in_batches(self, conn, table_name: str, base_query: str,
                                    cutoff_date: datetime, file_path: Path,
                                    batch_size: int, compress: bool):
        """Export data in batches to avoid memory issues."""
        offset = 0
        
        # Get column names
        columns_query = f"""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = :table_name 
            ORDER BY ordinal_position
        """
        result = await conn.execute(text(columns_query), {"table_name": table_name})
        columns = [row.column_name for row in result.fetchall()]
        
        # Open file for writing
        if compress:
            output_file = gzip.open(file_path, 'wt', newline='', encoding='utf-8')
        else:
            output_file = open(file_path, 'w', newline='', encoding='utf-8')
        
        try:
            writer = csv.DictWriter(output_file, fieldnames=columns)
            writer.writeheader()
            
            while True:
                # Get batch of data
                select_query = f"SELECT * {base_query} ORDER BY id LIMIT :batch_size OFFSET :offset"
                result = await conn.execute(
                    text(select_query), 
                    {"cutoff_date": cutoff_date, "batch_size": batch_size, "offset": offset}
                )
                
                rows = result.fetchall()
                if not rows:
                    break
                
                # Write batch to CSV
                for row in rows:
                    row_dict = dict(row._mapping)
                    # Handle datetime serialization
                    for key, value in row_dict.items():
                        if isinstance(value, datetime):
                            row_dict[key] = value.isoformat()
                        elif value is None:
                            row_dict[key] = ''
                    writer.writerow(row_dict)
                
                offset += batch_size
                logger.info(f"Exported batch {offset // batch_size} for {table_name}")
        
        finally:
            output_file.close()
    
    async def restore_archived_data(self, archive_id: str) -> bool:
        """Restore data from archive."""
        try:
            # Get archive metadata
            archive_record = await self._get_archive_record(archive_id)
            if not archive_record:
                logger.error(f"Archive record not found: {archive_id}")
                return False
            
            file_path = Path(archive_record['file_path'])
            if not file_path.exists():
                logger.error(f"Archive file not found: {file_path}")
                return False
            
            # Open archive file
            if archive_record['compressed']:
                input_file = gzip.open(file_path, 'rt', newline='', encoding='utf-8')
            else:
                input_file = open(file_path, 'r', newline='', encoding='utf-8')
            
            try:
                # Read CSV data
                reader = csv.DictReader(input_file)
                rows = list(reader)
                
                if not rows:
                    logger.info(f"No data to restore from archive {archive_id}")
                    return True
                
                # Restore data in batches
                table_name = archive_record['table_name']
                await self._restore_data_in_batches(table_name, rows, batch_size=500)
                
                logger.info(f"Restored {len(rows)} records from archive {archive_id}")
                return True
            
            finally:
                input_file.close()
        
        except Exception as e:
            logger.error(f"Failed to restore archive {archive_id}: {e}")
            return False
    
    async def _restore_data_in_batches(self, table_name: str, rows: List[Dict], batch_size: int):
        """Restore data in batches."""
        async with async_engine.begin() as conn:
            for i in range(0, len(rows), batch_size):
                batch = rows[i:i + batch_size]
                
                for row in batch:
                    # Convert empty strings back to None
                    for key, value in row.items():
                        if value == '':
                            row[key] = None
                        elif key in ['created_at', 'updated_at', 'started_at', 'completed_at', 
                                   'last_login_at', 'expires_at', 'last_used_at'] and value:
                            row[key] = datetime.fromisoformat(value)
                
                    # Build INSERT query
                    columns = list(batch[0].keys())
                    placeholders = [f":{col}" for col in columns]
                    
                    insert_query = f"""
                        INSERT INTO {table_name} ({', '.join(columns)})
                        VALUES ({', '.join(placeholders)})
                    """
                    
                    await conn.execute(text(insert_query), batch)
                
                logger.info(f"Restored batch {i // batch_size + 1} for {table_name}")
    
    async def _record_archive(self, archive: ArchivedRecord):
        """Record archive metadata in Redis."""
        try:
            redis_client = await self.get_redis()
            await redis_client.setex(
                f"archive:{archive.archive_id}",
                86400 * 30,  # 30 days
                json.dumps(archive.to_dict())
            )
        
        except Exception as e:
            logger.error(f"Failed to record archive: {e}")
    
    async def _get_archive_record(self, archive_id: str) -> Optional[Dict[str, Any]]:
        """Get archive record from Redis."""
        try:
            redis_client = await self.get_redis()
            record = await redis_client.get(f"archive:{archive_id}")
            return json.loads(record) if record else None
        
        except Exception as e:
            logger.error(f"Failed to get archive record: {e}")
            return None
    
    async def list_archives(self, table_name: Optional[str] = None) -> List[Dict[str, Any]]:
        """List available archives."""
        try:
            redis_client = await self.get_redis()
            
            # Get all archive keys
            cursor = 0
            archives = []
            
            while True:
                cursor, keys = await redis_client.scan(cursor, match="archive:*", count=100)
                
                for key in keys:
                    archive_id = key.decode().split(':')[1]
                    record = await self._get_archive_record(archive_id)
                    
                    if record and (not table_name or record['table_name'] == table_name):
                        archives.append(record)
                
                if cursor == 0:
                    break
            
            # Sort by archive date (newest first)
            archives.sort(key=lambda x: x['archive_date'], reverse=True)
            return archives
        
        except Exception as e:
            logger.error(f"Failed to list archives: {e}")
            return []
    
    async def get_archive_statistics(self) -> Dict[str, Any]:
        """Get archiving statistics."""
        try:
            archives = await self.list_archives()
            
            # Calculate statistics
            total_archives = len(archives)
            total_records = sum(archive['record_count'] for archive in archives)
            
            # Group by table
            table_stats = {}
            for archive in archives:
                table = archive['table_name']
                if table not in table_stats:
                    table_stats[table] = {
                        'archive_count': 0,
                        'total_records': 0,
                        'total_size_mb': 0
                    }
                
                table_stats[table]['archive_count'] += 1
                table_stats[table]['total_records'] += archive['record_count']
                
                # Calculate file size
                file_path = Path(archive['file_path'])
                if file_path.exists():
                    size_mb = file_path.stat().st_size / (1024 * 1024)
                    table_stats[table]['total_size_mb'] += size_mb
            
            # Get current table sizes
            current_sizes = await self._get_current_table_sizes()
            
            return {
                "total_archives": total_archives,
                "total_records_archived": total_records,
                "table_statistics": table_stats,
                "current_table_sizes": current_sizes,
                "archive_directory": str(self.archive_dir),
                "policies": {name: {
                    'retention_days': policy.retention_days,
                    'archive_condition': policy.archive_condition
                } for name, policy in self.policies.items()}
            }
        
        except Exception as e:
            logger.error(f"Failed to get archive statistics: {e}")
            return {"error": str(e)}
    
    async def _get_current_table_sizes(self) -> Dict[str, Any]:
        """Get current table sizes."""
        try:
            async with async_engine.begin() as conn:
                result = await conn.execute(text("""
                    SELECT 
                        tablename,
                        pg_size_pretty(pg_total_relation_size(schemaname||'.'||tablename)) as size,
                        pg_total_relation_size(schemaname||'.'||tablename) as size_bytes
                    FROM pg_tables 
                    WHERE schemaname = 'public'
                    ORDER BY pg_total_relation_size(schemaname||'.'||tablename) DESC
                """))
                
                sizes = {}
                for row in result.fetchall():
                    sizes[row.tablename] = {
                        "size_pretty": row.size,
                        "size_bytes": row.size_bytes
                    }
                
                return sizes
        
        except Exception as e:
            logger.error(f"Failed to get table sizes: {e}")
            return {}
    
    async def run_scheduled_archiving(self) -> Dict[str, Any]:
        """Run scheduled data archiving for all policies."""
        results = {
            "archived": [],
            "failed": [],
            "skipped": []
        }
        
        try:
            for table_name, policy in self.policies.items():
                logger.info(f"Running scheduled archiving for {table_name}")
                
                archive_record = await self.archive_table_data(table_name, policy)
                
                if archive_record:
                    results["archived"].append({
                        "table": table_name,
                        "archive_id": archive_record.archive_id,
                        "record_count": archive_record.record_count
                    })
                else:
                    results["skipped"].append({
                        "table": table_name,
                        "reason": "No data to archive"
                    })
            
            # Clean up old archive records
            await self._cleanup_old_archive_records()
            
            logger.info(f"Scheduled archiving completed: {len(results['archived'])} archived, {len(results['failed'])} failed, {len(results['skipped'])} skipped")
            
        except Exception as e:
            logger.error(f"Scheduled archiving failed: {e}")
            results["error"] = str(e)
        
        return results
    
    async def _cleanup_old_archive_records(self):
        """Clean up old archive records from Redis."""
        try:
            redis_client = await self.get_redis()
            
            # Get all archive keys
            cursor = 0
            old_keys = []
            
            while True:
                cursor, keys = await redis_client.scan(cursor, match="archive:*", count=100)
                
                for key in keys:
                    archive_id = key.decode().split(':')[1]
                    record = await self._get_archive_record(archive_id)
                    
                    if record:
                        archive_date = datetime.fromisoformat(record['archive_date'])
                        if archive_date < datetime.utcnow() - timedelta(days=90):
                            old_keys.append(key)
                
                if cursor == 0:
                    break
            
            # Delete old keys
            if old_keys:
                await redis_client.delete(*old_keys)
                logger.info(f"Cleaned up {len(old_keys)} old archive records")
        
        except Exception as e:
            logger.error(f"Failed to cleanup old archive records: {e}")
    
    async def verify_archive_integrity(self, archive_id: str) -> Dict[str, Any]:
        """Verify archive file integrity."""
        try:
            archive_record = await self._get_archive_record(archive_id)
            if not archive_record:
                return {"valid": False, "error": "Archive record not found"}
            
            file_path = Path(archive_record['file_path'])
            if not file_path.exists():
                return {"valid": False, "error": "Archive file not found"}
            
            # Check file size
            file_size = file_path.stat().st_size
            
            # For compressed files, verify gzip integrity
            if archive_record['compressed']:
                try:
                    with gzip.open(file_path, 'rb') as f:
                        f.read(1024)  # Try to read first 1KB
                except Exception as e:
                    return {"valid": False, "error": f"Corrupted gzip file: {e}"}
            
            # Count records in file
            if archive_record['compressed']:
                input_file = gzip.open(file_path, 'rt', newline='', encoding='utf-8')
            else:
                input_file = open(file_path, 'r', newline='', encoding='utf-8')
            
            try:
                reader = csv.DictReader(input_file)
                record_count = sum(1 for _ in reader)
            finally:
                input_file.close()
            
            # Verify record count matches
            count_matches = record_count == archive_record['record_count']
            
            return {
                "valid": True,
                "file_size_bytes": file_size,
                "actual_record_count": record_count,
                "expected_record_count": archive_record['record_count'],
                "count_matches": count_matches,
                "archive_date": archive_record['archive_date']
            }
        
        except Exception as e:
            return {"valid": False, "error": str(e)}


# Global data archiver instance
data_archiver = DataArchiver()


# Scheduled archiving task
async def scheduled_archiving_task():
    """Run scheduled data archiving."""
    logger.info("Starting scheduled data archiving...")
    
    try:
        results = await data_archiver.run_scheduled_archiving()
        
        # Store results in Redis
        redis_client = await data_archiver.get_redis()
        await redis_client.setex(
            "archiving_results",
            3600,  # 1 hour
            json.dumps(results, default=str)
        )
        
        logger.info("Scheduled data archiving completed")
        
    except Exception as e:
        logger.error(f"Scheduled data archiving failed: {e}")


if __name__ == "__main__":
    import sys
    
    async def main():
        command = sys.argv[1] if len(sys.argv) > 1 else "help"
        
        if command == "archive":
            table_name = sys.argv[2] if len(sys.argv) > 2 else None
            if not table_name:
                print("Error: archive requires table name")
                sys.exit(1)
            
            archive_record = await data_archiver.archive_table_data(table_name)
            if archive_record:
                print(f"✅ Archived {archive_record.record_count} records: {archive_record.archive_id}")
            else:
                print("❌ Archiving failed")
        
        elif command == "restore":
            archive_id = sys.argv[2] if len(sys.argv) > 2 else None
            if not archive_id:
                print("Error: restore requires archive ID")
                sys.exit(1)
            
            success = await data_archiver.restore_archived_data(archive_id)
            if success:
                print(f"✅ Archive {archive_id} restored successfully")
            else:
                print(f"❌ Archive {archive_id} restore failed")
        
        elif command == "list":
            table_name = sys.argv[2] if len(sys.argv) > 2 else None
            archives = await data_archiver.list_archives(table_name)
            print(f"Archives: {len(archives)}")
            for archive in archives:
                print(f"  - {archive['archive_id']}: {archive['table_name']} ({archive['record_count']} records, {archive['archive_date']})")
        
        elif command == "stats":
            stats = await data_archiver.get_archive_statistics()
            print(json.dumps(stats, indent=2, default=str))
        
        elif command == "verify":
            archive_id = sys.argv[2] if len(sys.argv) > 2 else None
            if not archive_id:
                print("Error: verify requires archive ID")
                sys.exit(1)
            
            verification = await data_archiver.verify_archive_integrity(archive_id)
            if verification["valid"]:
                print(f"✅ Archive {archive_id} is valid")
            else:
                print(f"❌ Archive {archive_id} verification failed")
                print(f"Error: {verification.get('error', 'Unknown error')}")
        
        elif command == "scheduled":
            results = await data_archiver.run_scheduled_archiving()
            print(json.dumps(results, indent=2, default=str))
        
        else:
            print("Usage: python data_archiver.py <command> [args]")
            print("Commands: archive <table>, restore <archive_id>, list [table], stats, verify <archive_id>, scheduled")
    
    asyncio.run(main())