//! # Ensemble Coordinator
//!
//! Coordinates multiple ML models for different serving modes.
//! Implements realistic ensemble strategies based on latency constraints.

#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic)]

use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};

use crate::integration::inference_engine::InferenceEngine;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{debug, info};

use super::{IntegrationHubConfig, ServingMode};
use crate::{InferenceResult, MLError, ModelMetadata, ModelType};

/// Configuration for ensemble coordination
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnsembleConfig {
    /// Maximum number of models in ensemble
    pub max_models: usize,
    /// Default ensemble strategy
    pub default_strategy: EnsembleStrategy,
    /// Timeout for ensemble coordination
    pub coordination_timeout_us: u64,
    /// Enable parallel execution
    pub enable_parallel: bool,
    /// Memory limit for ensemble
    pub memory_limit_mb: usize,
}

impl Default for EnsembleConfig {
    fn default() -> Self {
        Self {
            max_models: 5,
            default_strategy: EnsembleStrategy::WeightedAverage,
            coordination_timeout_us: 1000,
            enable_parallel: true,
            memory_limit_mb: 1024,
        }
    }
}

/// Ensemble strategies for model coordination
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub enum EnsembleStrategy {
    /// Single best model selection
    SingleModel,
    /// Weighted average of predictions
    WeightedAverage,
    /// Majority voting for classification
    MajorityVoting,
    /// Dynamic weighted average based on recent performance
    DynamicWeighting,
    /// Adaptive ensemble based on market conditions
    AdaptiveEnsemble,
}

/// Model context for ensemble coordination
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelContext {
    /// Unique model identifier
    pub model_id: String,
    /// Model type
    pub model_type: ModelType,
    /// Model weight in ensemble
    pub weight: f64,
    /// Model priority (higher = more important)
    pub priority: u32,
    /// Maximum latency allowed for this model
    pub max_latency_us: u64,
    /// Memory requirement in MB
    pub memory_mb: usize,
}

/// Execution plan for ensemble coordination
#[derive(Debug, Clone)]
pub struct ExecutionPlan {
    /// Selected models for execution
    pub models: Vec<ModelContext>,
    /// Ensemble strategy to use
    pub strategy: EnsembleStrategy,
    /// Total timeout for execution
    pub timeout_us: u64,
    /// Whether to execute models in parallel
    pub parallel_execution: bool,
    /// Expected memory usage
    pub expected_memory_mb: usize,
}

/// Execution statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ExecutionStats {
    /// Total number of executions
    pub total_executions: u64,
    /// Successful executions
    pub successful_executions: u64,
    /// Failed executions
    pub failed_executions: u64,
    /// Average execution time in microseconds
    pub avg_execution_time_us: f64,
    /// Total ensemble predictions made
    pub total_predictions: u64,
    /// Models currently registered
    pub registered_models: usize,
}

/// Ensemble Coordinator for managing multiple ML models
#[derive(Debug)]
pub struct EnsembleCoordinator {
    /// Configuration
    config: EnsembleConfig,
    /// Registered models
    models: Arc<RwLock<HashMap<String, ModelContext>>>,
    /// Execution statistics
    stats: Arc<RwLock<ExecutionStats>>,
    /// Model performance history
    performance_history: Arc<RwLock<HashMap<String, VecDeque<f64>>>>,
    /// Inference engine for real predictions
    inference_engine: Arc<InferenceEngine>,
}

impl EnsembleCoordinator {
    /// Create new ensemble coordinator
    pub async fn new() -> Result<Self, MLError> {
        Self::with_config(EnsembleConfig::default()).await
    }

    /// Create new ensemble coordinator with inference engine
    pub async fn with_engine(
        config: EnsembleConfig,
        inference_engine: Arc<InferenceEngine>,
    ) -> Self {
        Self {
            config,
            models: Arc::new(RwLock::new(HashMap::new())),
            stats: Arc::new(RwLock::new(ExecutionStats::default())),
            performance_history: Arc::new(RwLock::new(HashMap::new())),
            inference_engine,
        }
    }

    /// Create new ensemble coordinator with configuration
    pub async fn with_config(config: EnsembleConfig) -> Result<Self, MLError> {
        let hub_config = IntegrationHubConfig::default();
        let inference_engine = Arc::new(InferenceEngine::new(&hub_config).await.map_err(|e| {
            MLError::InitializationError {
                component: "InferenceEngine".to_owned(),
                message: format!("{:?}", e),
            }
        })?);
        Ok(Self {
            config,
            models: Arc::new(RwLock::new(HashMap::new())),
            stats: Arc::new(RwLock::new(ExecutionStats::default())),
            performance_history: Arc::new(RwLock::new(HashMap::new())),
            inference_engine,
        })
    }

    /// Register a model with the coordinator
    pub async fn register_model(&self, context: ModelContext) {
        let mut models = self.models.write().await;
        let mut stats = self.stats.write().await;

        let model_id = context.model_id.clone();
        models.insert(context.model_id.clone(), context);
        stats.registered_models = models.len();

        info!("Model registered: {} (total: {})", model_id, models.len());
    }

    /// Unregister a model
    pub async fn unregister_model(&self, model_id: &str) -> bool {
        let mut models = self.models.write().await;
        let mut stats = self.stats.write().await;
        let mut history = self.performance_history.write().await;

        let removed = models.remove(model_id).is_some();
        if removed {
            history.remove(model_id);
            stats.registered_models = models.len();
            info!("Model unregistered: {}", model_id);
        }

        removed
    }

    /// Create execution plan for given constraints
    pub async fn create_execution_plan(
        &self,
        serving_mode: &ServingMode,
        model_type: &ModelType,
        budget_us: u64,
    ) -> Result<ExecutionPlan, MLError> {
        let models = self.models.read().await;

        // Filter models by type and latency constraints
        let mut candidates: Vec<_> = models
            .values()
            .filter(|model| model.model_type == *model_type && model.max_latency_us <= budget_us)
            .cloned()
            .collect();

        if candidates.is_empty() {
            return Err(MLError::ModelNotFound(format!(
                "No models found for type {:?} within {}μs budget",
                model_type, budget_us
            )));
        }

        // Sort by priority (higher first)
        candidates.sort_by(|a, b| b.priority.cmp(&a.priority));

        // Select execution strategy based on serving mode
        let (strategy, models_to_use, parallel, timeout) = match serving_mode {
            ServingMode::UltraLowLatency => {
                // Use only the fastest, highest priority model
                let best_model =
                    candidates
                        .into_iter()
                        .next()
                        .ok_or_else(|| MLError::ConfigError("No candidate models available for UltraLowLatency mode"
                                .to_string()))?;
                let timeout = (budget_us / 2).min(50); // Conservative timeout
                (
                    EnsembleStrategy::SingleModel,
                    vec![best_model],
                    false,
                    timeout,
                )
            },
            ServingMode::LowLatency => {
                // Use top 2 models with weighted average
                let selected = candidates.into_iter().take(2).collect();
                let timeout = (budget_us * 3 / 4).min(200);
                (
                    EnsembleStrategy::WeightedAverage,
                    selected,
                    self.config.enable_parallel,
                    timeout,
                )
            },
            ServingMode::HighThroughput => {
                // Use all available models with dynamic weighting
                let timeout = budget_us;
                (
                    EnsembleStrategy::DynamicWeighting,
                    candidates,
                    true,
                    timeout,
                )
            },
        };

        let expected_memory = models_to_use.iter().map(|m| m.memory_mb).sum();

        Ok(ExecutionPlan {
            models: models_to_use,
            strategy,
            timeout_us: timeout,
            parallel_execution: parallel,
            expected_memory_mb: expected_memory,
        })
    }

    /// Execute ensemble prediction with given plan
    pub async fn execute_ensemble(
        &self,
        plan: &ExecutionPlan,
        input_features: &[f32],
    ) -> Result<InferenceResult, MLError> {
        let start_time = Instant::now();

        // Update execution stats
        {
            let mut stats = self.stats.write().await;
            stats.total_executions += 1;
        }

        // Execute models according to plan
        let results = if plan.parallel_execution && plan.models.len() > 1 {
            self.execute_parallel(&plan.models, input_features, plan.timeout_us)
                .await?
        } else {
            self.execute_sequential(&plan.models, input_features, plan.timeout_us)
                .await?
        };

        // Combine results using ensemble strategy
        let final_result = self.combine_results(&results, &plan.strategy).await?;

        let execution_time = start_time.elapsed();

        // Update performance stats
        {
            let mut stats = self.stats.write().await;
            stats.successful_executions += 1;
            stats.total_predictions += 1;

            // Update rolling average
            let new_time_us = execution_time.as_micros() as f64;
            let total = stats.total_executions as f64;
            stats.avg_execution_time_us =
                (stats.avg_execution_time_us * (total - 1.0) + new_time_us) / total;
        }

        debug!(
            "Ensemble execution completed in {}μs using {} models",
            execution_time.as_micros(),
            plan.models.len()
        );

        Ok(final_result)
    }

    /// Execute models in parallel
    async fn execute_parallel(
        &self,
        models: &[ModelContext],
        input_features: &[f32],
        timeout_us: u64,
    ) -> Result<Vec<(String, InferenceResult)>, MLError> {
        let _timeout = Duration::from_micros(timeout_us);
        let mut results = Vec::new();

        for model in models {
            match self.execute_single_model(model, input_features).await {
                Ok(result) => results.push((model.model_id.clone(), result)),
                Err(e) => {
                    tracing::warn!(
                        model_id = %model.model_id,
                        error = %e,
                        "Model failed during ensemble execution -- skipping"
                    );
                },
            }
        }

        if results.is_empty() {
            return Err(MLError::ModelError(
                "All models failed during ensemble execution".to_owned(),
            ));
        }

        Ok(results)
    }

    /// Execute models sequentially
    async fn execute_sequential(
        &self,
        models: &[ModelContext],
        input_features: &[f32],
        timeout_us: u64,
    ) -> Result<Vec<(String, InferenceResult)>, MLError> {
        let mut results = Vec::new();
        let start_time = Instant::now();

        for model in models {
            if start_time.elapsed().as_micros() as u64 >= timeout_us {
                break; // Timeout reached
            }

            match self.execute_single_model(model, input_features).await {
                Ok(result) => results.push((model.model_id.clone(), result)),
                Err(e) => {
                    tracing::warn!(
                        model_id = %model.model_id,
                        error = %e,
                        "Model failed during ensemble execution -- skipping"
                    );
                },
            }
        }

        if results.is_empty() {
            return Err(MLError::ModelError(
                "All models failed during ensemble execution".to_owned(),
            ));
        }

        Ok(results)
    }

    /// Execute a single model using real inference engine
    async fn execute_single_model(
        &self,
        model: &ModelContext,
        input_features: &[f32],
    ) -> Result<InferenceResult, MLError> {
        let _start_time = Instant::now();

        // Use real inference engine for prediction
        let result = match model.model_type {
            ModelType::DistilledMicroNet | ModelType::CompactDQN => {
                // Use micro model inference for low-latency models
                self.inference_engine
                    .process_micro_inference(&model.model_id, input_features)
                    .await
            },
            ModelType::DQN | ModelType::MAMBA | ModelType::TFT => {
                // Use ONNX inference for complex models
                self.inference_engine
                    .process_onnx_inference(&model.model_id, input_features)
                    .await
            },
            ModelType::RainbowDQN
            | ModelType::TGGN
            | ModelType::LNN
            | ModelType::TLOB
            | ModelType::PPO
            | ModelType::Transformer
            | ModelType::Mamba
            | ModelType::LiquidNet
            | ModelType::TGNN
            | ModelType::Ensemble
            | ModelType::KAN
            | ModelType::XLSTM
            | ModelType::Diffusion => {
                // Fallback to intelligent prediction based on features
                self.generate_model_specific_prediction(model, input_features)
                    .await
            },
        };

        match result {
            Ok(inference_result) => Ok(inference_result),
            Err(e) => {
                tracing::warn!(
                    model_id = %model.model_id,
                    error = %e,
                    "Model inference failed -- no fallback available"
                );
                Err(e)
            },
        }
    }

    /// Fallback when no loaded model is available for a given model context.
    ///
    /// Returns an error unconditionally. Heuristic predictors have been removed
    /// to prevent the ensemble from trading on fake signals. The ensemble's
    /// per-model error handling will skip this model and continue with others.
    async fn generate_model_specific_prediction(
        &self,
        model: &ModelContext,
        _input_features: &[f32],
    ) -> Result<InferenceResult, MLError> {
        Err(MLError::ModelError(format!(
            "Model '{}' (type {:?}) has no loaded weights -- cannot generate real prediction. \
             Fallback heuristics have been removed to prevent trading on fake signals.",
            model.model_id, model.model_type
        )))
    }

    /// Combine results using ensemble strategy
    async fn combine_results(
        &self,
        results: &[(String, InferenceResult)],
        strategy: &EnsembleStrategy,
    ) -> Result<InferenceResult, MLError> {
        if results.is_empty() {
            return Err(MLError::InferenceError("No results to combine".to_owned()));
        }

        match strategy {
            EnsembleStrategy::SingleModel => {
                // Return the first (best) result
                results
                    .first()
                    .map(|(_, r)| r.clone())
                    .ok_or_else(|| MLError::InferenceError("No results available".to_owned()))
            },
            EnsembleStrategy::WeightedAverage => {
                // Weighted average of predictions
                let models = self.models.read().await;
                let mut weighted_sum = 0.0;
                let mut total_weight = 0.0;
                let mut total_confidence = 0.0;
                let mut max_latency = 0;

                for (model_id, result) in results {
                    if let Some(model) = models.get(model_id) {
                        weighted_sum += result.prediction_value * model.weight;
                        total_weight += model.weight;
                        total_confidence += result.confidence;
                        max_latency = max_latency.max(result.latency_us);
                    }
                }

                let final_prediction = if total_weight > 0.0 {
                    weighted_sum / total_weight
                } else {
                    0.0
                };

                Ok(InferenceResult {
                    model_id: "ensemble".to_owned(),
                    prediction_value: final_prediction,
                    confidence: total_confidence / results.len() as f64,
                    latency_us: max_latency,
                    timestamp: SystemTime::now()
                        .duration_since(UNIX_EPOCH)
                        .map_err(|e| MLError::ModelError(format!("Time error: {}", e)))?
                        .as_micros() as u64,
                    metadata: ModelMetadata {
                        model_type: ModelType::CompactDQN, // Default for ensemble
                        version: "ensemble-1.0".to_owned(),
                        features_used: results
                            .first()
                            .map(|(_, r)| r.metadata.features_used)
                            .unwrap_or(0),
                        memory_usage_mb: results
                            .iter()
                            .map(|(_, r)| r.metadata.memory_usage_mb)
                            .sum(),
                        additional_metadata: HashMap::new(),
                    },
                })
            },
            EnsembleStrategy::MajorityVoting
            | EnsembleStrategy::DynamicWeighting
            | EnsembleStrategy::AdaptiveEnsemble => {
                // For other strategies, use weighted average as fallback
                Box::pin(self.combine_results(results, &EnsembleStrategy::WeightedAverage)).await
            },
        }
    }

    /// Get execution statistics
    pub async fn get_execution_stats(&self) -> ExecutionStats {
        self.stats.read().await.clone()
    }

    /// Get list of registered models
    pub async fn get_registered_models(&self) -> Vec<String> {
        self.models.read().await.keys().cloned().collect()
    }

    /// Update model performance metrics
    pub async fn update_model_performance(&self, model_id: &str, performance_score: f64) {
        let mut history = self.performance_history.write().await;
        let scores = history
            .entry(model_id.to_string())
            .or_insert_with(VecDeque::new);

        scores.push_back(performance_score);

        // Keep only last 100 scores
        if scores.len() > 100 {
            scores.pop_front();
        }
    }

    /// Get model performance history
    pub async fn get_model_performance(&self, model_id: &str) -> Option<Vec<f64>> {
        let history = self.performance_history.read().await;
        history
            .get(model_id)
            .map(|scores| scores.iter().copied().collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_coordinator_creation() -> Result<(), Box<dyn std::error::Error>> {
        let coordinator = EnsembleCoordinator::with_config(EnsembleConfig::default()).await?;
        let stats = coordinator.get_execution_stats().await;
        assert_eq!(stats.total_executions, 0);
        assert_eq!(stats.registered_models, 0);
        Ok(())
    }

    #[tokio::test]
    async fn test_model_registration() -> Result<(), Box<dyn std::error::Error>> {
        let coordinator = EnsembleCoordinator::with_config(EnsembleConfig::default()).await?;

        let context = ModelContext {
            model_id: "test_model".to_owned(),
            model_type: ModelType::CompactDQN,
            weight: 1.0,
            priority: 10,
            max_latency_us: 100,
            memory_mb: 50,
        };

        coordinator.register_model(context).await;

        let plan = coordinator
            .create_execution_plan(&ServingMode::LowLatency, &ModelType::CompactDQN, 1000)
            .await?;

        assert_eq!(plan.models.len(), 1);
        assert_eq!(
            plan.models.first().map(|m| &m.model_id),
            Some(&"test_model".to_owned())
        );
        Ok(())
    }

    #[tokio::test]
    async fn test_execution_plan_ultra_low_latency() -> Result<(), Box<dyn std::error::Error>> {
        let coordinator = EnsembleCoordinator::with_config(EnsembleConfig::default()).await?;

        // Register a fast model
        let context = ModelContext {
            model_id: "ultra_fast_model".to_owned(),
            model_type: ModelType::DistilledMicroNet,
            weight: 1.0,
            priority: 100,
            max_latency_us: 20,
            memory_mb: 1,
        };

        coordinator.register_model(context).await;

        let plan = coordinator
            .create_execution_plan(
                &ServingMode::UltraLowLatency,
                &ModelType::DistilledMicroNet,
                100,
            )
            .await?;

        assert!(matches!(plan.strategy, EnsembleStrategy::SingleModel));
        assert_eq!(plan.models.len(), 1);
        assert_eq!(plan.timeout_us, 50); // Should be clamped
        assert!(!plan.parallel_execution);
        Ok(())
    }
}
