backend/tests/startup_validation.py

"""
Production startup validation for AegisLM SaaS Backend.

Validates all critical services are available before starting the application.
Fails fast if any critical service is unavailable.
"""

import asyncio
import sys
import os
import logging
from typing import Dict, Any, Optional
from datetime import datetime

# Add the backend directory to Python path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from core.config import settings
from core.database import check_db_health, check_redis_health


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class StartupValidator:
    """Validates all services are ready for production startup."""
    
    def __init__(self):
        self.validation_results: Dict[str, Dict[str, Any]] = {}
        self.critical_services = ["database", "redis"]
        self.optional_services = ["model_providers"]
    
    async def validate_all_services(self) -> bool:
        """
        Validate all required services.
        
        Returns:
            bool: True if all critical services are healthy
        """
        logger.info("🚀 Starting production validation...")
        
        # Validate critical services
        await self._validate_database()
        await self._validate_redis()
        
        # Validate optional services
        await self._validate_model_providers()
        
        # Validate configuration
        self._validate_configuration()
        
        # Generate validation report
        self._generate_validation_report()
        
        # Check if all critical services are healthy
        critical_healthy = all(
            self.validation_results[service]["healthy"] 
            for service in self.critical_services
        )
        
        if critical_healthy:
            logger.info("✅ All critical services validated successfully")
            return True
        else:
            logger.error("❌ Critical services validation failed")
            return False
    
    async def _validate_database(self):
        """Validate database connection."""
        logger.info("🔍 Validating database connection...")
        
        try:
            db_healthy = await check_db_health()
            
            self.validation_results["database"] = {
                "healthy": db_healthy,
                "message": "Database connection successful" if db_healthy else "Database connection failed",
                "timestamp": datetime.now().isoformat(),
                "critical": True
            }
            
            if db_healthy:
                logger.info("✅ Database validation passed")
            else:
                logger.error("❌ Database validation failed")
                
        except Exception as e:
            self.validation_results["database"] = {
                "healthy": False,
                "message": f"Database validation error: {str(e)}",
                "timestamp": datetime.now().isoformat(),
                "critical": True
            }
            logger.error(f"❌ Database validation error: {e}")
    
    async def _validate_redis(self):
        """Validate Redis connection."""
        logger.info("🔍 Validating Redis connection...")
        
        try:
            redis_healthy = await check_redis_health()
            
            self.validation_results["redis"] = {
                "healthy": redis_healthy,
                "message": "Redis connection successful" if redis_healthy else "Redis connection failed",
                "timestamp": datetime.now().isoformat(),
                "critical": True
            }
            
            if redis_healthy:
                logger.info("✅ Redis validation passed")
            else:
                logger.error("❌ Redis validation failed")
                
        except Exception as e:
            self.validation_results["redis"] = {
                "healthy": False,
                "message": f"Redis validation error: {str(e)}",
                "timestamp": datetime.now().isoformat(),
                "critical": True
            }
            logger.error(f"❌ Redis validation error: {e}")
    
    async def _validate_model_providers(self):
        """Validate model provider configurations."""
        logger.info("🔍 Validating model provider configurations...")
        
        providers = {
            "openai": settings.OPENAI_API_KEY,
            "anthropic": settings.ANTHROPIC_API_KEY,
            "google": settings.GOOGLE_API_KEY,
            "groq": getattr(settings, 'GROQ_API_KEY', None),
            "mistral": getattr(settings, 'MISTRAL_API_KEY', None),
            "huggingface": getattr(settings, 'HUGGINGFACE_API_KEY', None)
        }
        
        configured_providers = [name for name, key in providers.items() if key]
        
        self.validation_results["model_providers"] = {
            "healthy": len(configured_providers) > 0,
            "message": f"Configured providers: {configured_providers}",
            "configured_providers": configured_providers,
            "timestamp": datetime.now().isoformat(),
            "critical": False
        }
        
        if configured_providers:
            logger.info(f"✅ Model providers configured: {configured_providers}")
        else:
            logger.warning("⚠️ No model providers configured")
    
    def _validate_configuration(self):
        """Validate critical configuration."""
        logger.info("🔍 Validating configuration...")
        
        config_issues = []
        
        # Check critical configuration
        if not settings.SECRET_KEY or settings.SECRET_KEY == "your-production-secret-key-change-this":
            config_issues.append("SECRET_KEY must be set to a secure value")
        
        if not settings.DATABASE_URL:
            config_issues.append("DATABASE_URL must be set")
        
        if not settings.REDIS_URL:
            config_issues.append("REDIS_URL must be set")
        
        # Check learning system configuration
        if settings.ENABLE_LEARNING:
            if not settings.TARGET_MODELS:
                config_issues.append("TARGET_MODELS must be set when ENABLE_LEARNING=true")
        
        self.validation_results["configuration"] = {
            "healthy": len(config_issues) == 0,
            "message": "Configuration valid" if len(config_issues) == 0 else f"Configuration issues: {config_issues}",
            "issues": config_issues,
            "timestamp": datetime.now().isoformat(),
            "critical": True
        }
        
        if config_issues:
            logger.error(f"❌ Configuration issues: {config_issues}")
        else:
            logger.info("✅ Configuration validation passed")
    
    def _generate_validation_report(self):
        """Generate and log validation report."""
        logger.info("\n" + "="*80)
        logger.info("📊 PRODUCTION VALIDATION REPORT")
        logger.info("="*80)
        
        for service, result in self.validation_results.items():
            status = "✅ HEALTHY" if result["healthy"] else "❌ UNHEALTHY"
            critical = " (CRITICAL)" if result.get("critical", False) else " (OPTIONAL)"
            logger.info(f"{service.upper()}{critical}: {status}")
            logger.info(f"  Message: {result['message']}")
            logger.info(f"  Timestamp: {result['timestamp']}")
        
        logger.info("="*80)


async def main():
    """Main validation function."""
    validator = StartupValidator()
    
    # Run validation
    all_healthy = await validator.validate_all_services()
    
    # Exit with appropriate code
    if all_healthy:
        logger.info("🎉 Production validation completed successfully!")
        logger.info("🚀 Ready to start AegisLM SaaS Backend")
        sys.exit(0)
    else:
        logger.error("💥 Production validation failed!")
        logger.error("🛑 Cannot start AegisLM SaaS Backend")
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())