main.py

"""
Nebula AI Framework - Main Entry Point
=====================================

This script initializes and manages the lifecycle of the Nebula AI Framework,
acting as the central entry point for execution. It orchestrates vital components
and workflows, ensuring smooth and reliable operation.

Core Features:
    - Configuration Loader: Flexible and validated YAML-based setup
    - Logging: JSON-configurable logging for debugging and runtime insights
    - Data Pipeline: Integration with preprocessing, training, and monitoring modules
    - Orchestration: Centralized workflow management for AI processes
    - Error Handling: Graceful error catching and logging
"""

import logging
import yaml
import json
import os
import sys
from typing import Dict, Any, Optional
from datetime import datetime

# Add the project root to Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from core.base import ConfigManager, NebulaException
from modules.data_pipeline import DataPipeline
from modules.model_trainer import ModelTrainer
from modules.monitoring import ModelMonitor
from modules.inference_service import InferenceService
from modules.feedback_loop import FeedbackLoop
from modules.error_tracker import ErrorTracker, ErrorSeverity, ErrorCategory
from modules.data_validation import DataValidator
from modules.security import SecurityManager
from modules.monitoring import NumpyJSONEncoder


def setup_logging(config: Dict[str, Any]) -> None:
    """
    Initialize logging configuration based on the provided config.
    
    :param config: Configuration dictionary containing logging settings
    """
    try:
        log_config = config.get('logging', {})
        log_level = getattr(logging, log_config.get('level', 'INFO').upper())
        log_format = log_config.get('format', '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        log_file = log_config.get('log_file', 'logs/app.log')
        error_log_file = log_config.get('error_log_file', 'logs/errors.log')
        
        # Create logs directory
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        
        # Configure root logger
        logging.basicConfig(
            level=log_level,
            format=log_format,
            handlers=[
                logging.FileHandler(log_file),
                logging.FileHandler(error_log_file) if error_log_file != log_file else logging.NullHandler(),
                logging.StreamHandler()
            ]
        )
        
        logger = logging.getLogger(__name__)
        logger.info("Logging configuration initialized successfully")
        
    except Exception as e:
        # Fallback to basic logging
        logging.basicConfig(level=logging.INFO)
        logging.error(f"Failed to configure logging: {e}. Using basic logging.")


def load_config(config_path: str = "config/config.yaml") -> Dict[str, Any]:
    """
    Load and validate the main configuration file.
    
    :param config_path: Path to the configuration file
    :return: Parsed and validated configuration dictionary
    """
    try:
        config = ConfigManager.load_config(config_path)
        
        # Validate required keys
        required_keys = ["data_pipeline", "model", "monitoring"]
        for key in required_keys:
            if key not in config:
                raise NebulaException(f"Missing required configuration key: {key}")
        
        return config
        
    except Exception as e:
        raise NebulaException(f"Failed to load configuration: {e}")


def initialize_components(config: Dict[str, Any]) -> Dict[str, Any]:
    """
    Initialize all framework components.
    
    :param config: Configuration dictionary
    :return: Dictionary of initialized components
    """
    components = {}
    logger = logging.getLogger(__name__)
    
    try:
        logger.info("Initializing framework components...")
        
        # Initialize Security Manager first
        security_config = config.get('security', {})
        security_manager = SecurityManager(security_config)
        if not security_manager.initialize():
            raise NebulaException("Failed to initialize Security Manager")
        components['security_manager'] = security_manager
        
        # Initialize Error Tracker
        error_tracker = ErrorTracker(config.get('error_tracking', {}))
        if not error_tracker.initialize():
            raise NebulaException("Failed to initialize Error Tracker")
        components['error_tracker'] = error_tracker
        
        # Initialize Data Validator
        data_validator = DataValidator(config.get('data_validation', {}))
        if not data_validator.initialize():
            raise NebulaException("Failed to initialize Data Validator")
        components['data_validator'] = data_validator
        
        # Initialize Data Pipeline
        data_pipeline = DataPipeline(config.get('data_pipeline', {}))
        if not data_pipeline.initialize():
            raise NebulaException("Failed to initialize Data Pipeline")
        components['data_pipeline'] = data_pipeline
        
        # Initialize Model Trainer
        model_trainer = ModelTrainer(config.get('model', {}))
        if not model_trainer.initialize():
            raise NebulaException("Failed to initialize Model Trainer")
        components['model_trainer'] = model_trainer
        
        # Initialize Feedback Loop
        feedback_loop = FeedbackLoop(config.get('feedback_loop', {}))
        if not feedback_loop.initialize():
            raise NebulaException("Failed to initialize Feedback Loop")
        components['feedback_loop'] = feedback_loop
        
        # Initialize Model Monitor
        monitor = ModelMonitor(config.get('monitoring', {}))
        if not monitor.initialize():
            raise NebulaException("Failed to initialize Model Monitor")
        components['monitor'] = monitor
        
        # Initialize Inference Service
        inference_config = config.get('inference', {})
        inference_config.update(config.get('api_server', {}))
        inference_service = InferenceService(inference_config)
        if not inference_service.initialize():
            raise NebulaException("Failed to initialize Inference Service")
        components['inference_service'] = inference_service
        
        logger.info("All components initialized successfully")
        return components
        
    except Exception as e:
        logger.error(f"Component initialization failed: {e}")
        # Track the error
        if 'error_tracker' in components:
            components['error_tracker'].track_error(e, ErrorSeverity.HIGH, ErrorCategory.INFRASTRUCTURE)
        raise NebulaException(f"Component initialization failed: {e}")


def run_training_pipeline(components: Dict[str, Any], config: Dict[str, Any]) -> Dict[str, Any]:
    """
    Execute the complete training pipeline with comprehensive validation and monitoring.
    
    :param components: Dictionary of initialized components
    :param config: Configuration dictionary
    :return: Training results
    """
    logger = logging.getLogger(__name__)
    error_tracker = components.get('error_tracker')
    security_manager = components.get('security_manager')
    
    try:
        logger.info("Starting training pipeline...")
        
        # Log pipeline start
        if security_manager:
            security_manager.log_access('system', 'training_pipeline', 'start')
        
        # Data processing with validation
        data_pipeline = components['data_pipeline']
        data_validator = components['data_validator']
        
        # Load and validate data
        features, target = data_pipeline.process()
        
        # Validate data quality
        validation_results = data_validator.validate_data(features)
        if not validation_results['validation_passed']:
            logger.warning("Data validation failed - proceeding with caution")
        
        # Split data
        X_train, X_test, y_train, y_test = data_pipeline.split_data(features, target)
        
        # Model training
        model_trainer = components['model_trainer']
        training_results = model_trainer.train(X_train, y_train)
        
        # Model evaluation
        predictions = model_trainer.predict(X_test)
        performance_metrics = components['monitor'].record_model_performance(
            y_test, predictions, 
            model_type=config['model'].get('type', 'classification')
        )
        
        # Collect feedback on training performance
        feedback_loop = components['feedback_loop']
        for i, (input_data, pred, actual) in enumerate(zip(X_test.values[:10], predictions[:10], y_test.values[:10])):
            feedback_loop.collect_feedback(
                input_data=input_data.tolist(),
                prediction=pred,
                ground_truth=actual,
                model_version="1.0",
                user_rating=5 if pred == actual else 1
            )
        
        # Save model
        model_path = model_trainer.save_model()
        
        results = {
            'training_results': training_results,
            'performance_metrics': performance_metrics,
            'model_path': model_path,
            'data_summary': data_pipeline.get_data_summary(),
            'validation_results': validation_results,
            'timestamp': datetime.now().isoformat()
        }
        
        logger.info("Training pipeline completed successfully")
        
        # Log pipeline completion
        if security_manager:
            security_manager.log_access('system', 'training_pipeline', 'complete')
        
        return results
        
    except Exception as e:
        logger.error(f"Training pipeline failed: {e}")
        
        # Track the error
        if error_tracker:
            error_tracker.track_error(e, ErrorSeverity.HIGH, ErrorCategory.MODEL, component='training_pipeline')
        
        # Log security event
        if security_manager:
            security_manager.log_access('system', 'training_pipeline', 'error')
        
        raise NebulaException(f"Training pipeline failed: {e}")


def start_inference_service(components: Dict[str, Any]) -> bool:
    """
    Start the inference service.
    
    :param components: Dictionary of initialized components
    :return: True if service started successfully
    """
    logger = logging.getLogger(__name__)
    
    try:
        inference_service = components['inference_service']
        
        if inference_service.start_service():
            logger.info("Inference service started successfully")
            return True
        else:
            logger.error("Failed to start inference service")
            return False
            
    except Exception as e:
        logger.error(f"Failed to start inference service: {e}")
        return False


def start_monitoring(components: Dict[str, Any]) -> bool:
    """
    Start the monitoring system.
    
    :param components: Dictionary of initialized components
    :return: True if monitoring started successfully
    """
    logger = logging.getLogger(__name__)
    
    try:
        monitor = components['monitor']
        
        if monitor.start_monitoring():
            logger.info("Monitoring system started successfully")
            return True
        else:
            logger.error("Failed to start monitoring system")
            return False
            
    except Exception as e:
        logger.error(f"Failed to start monitoring system: {e}")
        return False


def cleanup_components(components: Dict[str, Any]) -> None:
    """
    Clean up all components gracefully.
    
    :param components: Dictionary of components to clean up
    """
    logger = logging.getLogger(__name__)
    
    try:
        logger.info("Cleaning up components...")
        
        for component_name, component in components.items():
            try:
                if hasattr(component, 'cleanup'):
                    component.cleanup()
                logger.info(f"Cleaned up {component_name}")
            except Exception as e:
                logger.error(f"Failed to clean up {component_name}: {e}")
        
        logger.info("Component cleanup completed")
        
    except Exception as e:
        logger.error(f"Cleanup failed: {e}")


def generate_report(components: Dict[str, Any], training_results: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """
    Generate a comprehensive framework report.
    
    :param components: Dictionary of components
    :param training_results: Optional training results
    :return: Generated report
    """
    logger = logging.getLogger(__name__)
    
    try:
        # Get monitoring report
        monitor = components.get('monitor')
        monitor_report = monitor.generate_report() if monitor else {}
        
        # Get inference service info
        inference_service = components.get('inference_service')
        inference_info = inference_service.get_service_info() if inference_service else {}
        
        # Compile complete report
        report = {
            'generated_at': datetime.now().isoformat(),
            'framework_status': {
                'components_initialized': len(components),
                'component_status': {name: getattr(comp, 'initialized', False) for name, comp in components.items()}
            },
            'monitoring_report': monitor_report,
            'inference_service': inference_info
        }
        
        if training_results:
            report['training_results'] = training_results
        
        # Save report
        os.makedirs('reports', exist_ok=True)
        report_path = f"reports/nebula_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        with open(report_path, 'w') as f:
            json.dump(report, f, indent=2, cls=NumpyJSONEncoder)
        
        logger.info(f"Report generated and saved to {report_path}")
        return report
        
    except Exception as e:
        logger.error(f"Failed to generate report: {e}")
        return {'error': str(e)}


def main():
    """
    Main entry point for the Nebula AI Framework.
    
    This function orchestrates the complete lifecycle of the framework,
    including initialization, training, monitoring, and cleanup.
    """
    logger = None
    components = None
    
    try:
        # Load configuration
        config = load_config()
        
        # Setup logging
        setup_logging(config)
        logger = logging.getLogger(__name__)
        logger.info("Nebula AI Framework starting...")
        
        # Initialize components
        components = initialize_components(config)
        
        # Run training pipeline
        training_results = run_training_pipeline(components, config)
        
        # Start monitoring
        start_monitoring(components)
        
        # Start inference service
        start_inference_service(components)
        
        # Generate and save report
        report = generate_report(components, training_results)
        
        logger.info("Nebula AI Framework started successfully")
        logger.info(f"Inference service available at: http://{config.get('api_server', {}).get('host', '0.0.0.0')}:{config.get('api_server', {}).get('port', 8080)}")
        logger.info("Press Ctrl+C to stop the framework")
        
        # Keep the main thread alive
        try:
            import time
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            logger.info("Shutdown signal received")
        
    except NebulaException as e:
        if logger:
            logger.error(f"Nebula Framework error: {e}")
        else:
            print(f"Nebula Framework error: {e}")
        sys.exit(1)
        
    except Exception as e:
        if logger:
            logger.error(f"Unexpected error: {e}")
        else:
            print(f"Unexpected error: {e}")
        sys.exit(1)
        
    finally:
        # Cleanup
        if components:
            cleanup_components(components)
        
        if logger:
            logger.info("Nebula AI Framework shutdown complete")


if __name__ == "__main__":
    main()