//! ML Model Fallback and Failover Manager for Production Resilience

use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use std::time::{Instant, SystemTime};
use tokio::sync::{broadcast, RwLock};
use tracing::{debug, error, info, warn};

/// Model health status
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ModelHealth {
    Healthy,
    Degraded,
    Unhealthy,
    Failed,
    Offline,
}

/// Model availability status
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelStatus {
    /// Model identifier
    pub model_id: String,
    /// Current health status
    pub health: ModelHealth,
    /// Last successful prediction timestamp
    pub last_success: Option<SystemTime>,
    /// Consecutive failure count
    pub consecutive_failures: u32,
    /// Total prediction count
    pub total_predictions: u64,
    /// Success rate (0.0-1.0)
    pub success_rate: f64,
    /// Average latency in microseconds
    pub avg_latency_us: f64,
    /// Current accuracy score
    pub accuracy_score: f64,
    /// Model priority (higher = preferred)
    pub priority: i32,
    /// Whether model is enabled for predictions
    pub enabled: bool,
    /// Last health check timestamp
    pub last_health_check: SystemTime,
    /// Circuit breaker state
    pub circuit_breaker_state: CircuitBreakerState,
}

/// Circuit breaker states
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CircuitBreakerState {
    Closed,   // Normal operation
    Open,     // Blocking requests
    HalfOpen, // Testing recovery
}

/// Fallback configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FallbackConfig {
    /// Minimum number of healthy models required
    pub min_healthy_models: usize,
    /// Maximum consecutive failures before marking unhealthy
    pub max_consecutive_failures: u32,
    /// Minimum success rate threshold
    pub min_success_rate: f64,
    /// Maximum acceptable latency in microseconds
    pub max_latency_us: u64,
    /// Minimum accuracy threshold
    pub min_accuracy: f64,
    /// Health check interval in seconds
    pub health_check_interval_seconds: u64,
    /// Circuit breaker failure threshold
    pub circuit_breaker_failure_threshold: u32,
    /// Circuit breaker timeout in seconds
    pub circuit_breaker_timeout_seconds: u64,
    /// Enable automatic model switching
    pub enable_auto_switching: bool,
    /// Fallback timeout in milliseconds
    pub fallback_timeout_ms: u64,
}

impl Default for FallbackConfig {
    fn default() -> Self {
        Self {
            min_healthy_models: 1,
            max_consecutive_failures: 5,
            min_success_rate: 0.7,
            max_latency_us: 5000, // 5ms
            min_accuracy: 0.6,
            health_check_interval_seconds: 30,
            circuit_breaker_failure_threshold: 10,
            circuit_breaker_timeout_seconds: 60,
            enable_auto_switching: true,
            fallback_timeout_ms: 100,
        }
    }
}

/// Fallback strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FallbackStrategy {
    /// Use highest priority healthy model
    PriorityBased,
    /// Use model with best recent performance
    PerformanceBased,
    /// Use weighted ensemble of available models
    EnsembleBased,
    /// Use simple rule-based prediction
    RuleBasedFallback,
    /// Return neutral prediction
    NeutralFallback,
}

/// Prediction result with fallback information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FallbackPrediction {
    /// Prediction value
    pub prediction_value: f64,
    /// Confidence score
    pub confidence: f64,
    /// Model(s) used for prediction
    pub models_used: Vec<String>,
    /// Fallback strategy applied
    pub strategy_used: FallbackStrategy,
    /// Whether fallback was triggered
    pub fallback_triggered: bool,
    /// Inference latency
    pub latency_us: u64,
    /// Warning messages
    pub warnings: Vec<String>,
}

/// Model failover event
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FailoverEvent {
    /// Event timestamp
    pub timestamp: SystemTime,
    /// Event type
    pub event_type: FailoverEventType,
    /// Primary model that failed
    pub failed_model: Option<String>,
    /// Fallback model used
    pub fallback_model: Option<String>,
    /// Fallback strategy applied
    pub strategy: FallbackStrategy,
    /// Event message
    pub message: String,
    /// Impact assessment
    pub impact: FailoverImpact,
}

/// Failover event types
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FailoverEventType {
    ModelFailure,
    ModelDegraded,
    CircuitBreakerOpen,
    AutoSwitching,
    ManualSwitching,
    Recovery,
}

/// Failover impact assessment
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FailoverImpact {
    None,     // No impact, seamless fallback
    Low,      // Slight performance degradation
    Medium,   // Noticeable impact on predictions
    High,     // Significant impact, limited functionality
    Critical, // Major service degradation
}

/// ML Model Fallback Manager
#[derive(Debug)]
pub struct MLFallbackManager {
    /// Fallback configuration
    config: Arc<RwLock<FallbackConfig>>,
    /// Model status tracking
    model_status: Arc<RwLock<HashMap<String, ModelStatus>>>,
    /// Model priorities (higher = more preferred)
    model_priorities: Arc<RwLock<BTreeMap<i32, Vec<String>>>>,
    /// Current primary model
    current_primary: Arc<RwLock<Option<String>>>,
    /// Failover event history
    failover_events: Arc<RwLock<Vec<FailoverEvent>>>,
    /// Event broadcaster
    event_broadcaster: Arc<broadcast::Sender<FailoverEvent>>,
}

impl MLFallbackManager {
    /// Create new fallback manager
    pub fn new() -> Self {
        let (_event_sender, _) = broadcast::channel(100);

        Self {
            config: Arc::new(RwLock::new(FallbackConfig::default())),
            model_status: Arc::new(RwLock::new(HashMap::new())),
            model_priorities: Arc::new(RwLock::new(BTreeMap::new())),
            current_primary: Arc::new(RwLock::new(None)),
            failover_events: Arc::new(RwLock::new(Vec::new())),
            event_broadcaster: Arc::new(_event_sender),
        }
    }

    /// Register a new model
    pub async fn register_model(&self, model_id: String, priority: i32) {
        info!("Registering model {} with priority {}", model_id, priority);

        let status = ModelStatus {
            model_id: model_id.clone(),
            health: ModelHealth::Healthy,
            last_success: None,
            consecutive_failures: 0,
            total_predictions: 0,
            success_rate: 1.0,
            avg_latency_us: 0.0,
            accuracy_score: 1.0,
            priority,
            enabled: true,
            last_health_check: SystemTime::now(),
            circuit_breaker_state: CircuitBreakerState::Closed,
        };

        // Add to status tracking
        {
            let mut status_map = self.model_status.write().await;
            status_map.insert(model_id.clone(), status);
        }

        // Add to priority mapping
        {
            let mut priorities = self.model_priorities.write().await;
            priorities
                .entry(priority)
                .or_insert_with(Vec::new)
                .push(model_id.clone());
        }

        // Set as primary if this is the highest priority model
        {
            let mut current_primary = self.current_primary.write().await;
            if current_primary.is_none() {
                *current_primary = Some(model_id);
            }
        }
    }

    /// Record prediction result for model health tracking
    pub async fn record_prediction_result(
        &self,
        model_id: &str,
        success: bool,
        latency_us: u64,
        accuracy: Option<f64>,
    ) {
        let mut status_map = self.model_status.write().await;

        if let Some(status) = status_map.get_mut(model_id) {
            status.total_predictions += 1;

            if success {
                status.consecutive_failures = 0;
                status.last_success = Some(SystemTime::now());

                // Update success rate
                let total = status.total_predictions as f64;
                let current_successes = status.success_rate * (total - 1.0);
                status.success_rate = (current_successes + 1.0) / total;

                // Update latency
                status.avg_latency_us =
                    (status.avg_latency_us * (total - 1.0) + latency_us as f64) / total;

                // Update accuracy if provided
                if let Some(acc) = accuracy {
                    status.accuracy_score = (status.accuracy_score * (total - 1.0) + acc) / total;
                }
            } else {
                status.consecutive_failures += 1;

                // Update success rate
                let total = status.total_predictions as f64;
                let current_successes = status.success_rate * (total - 1.0);
                status.success_rate = current_successes / total;
            }

            // Update health status
            self.update_model_health(model_id, status).await;
        }
    }

    /// Update model health based on current metrics
    async fn update_model_health(&self, model_id: &str, status: &mut ModelStatus) {
        let config = self.config.read().await;

        let previous_health = status.health;

        // Check circuit breaker
        if status.circuit_breaker_state == CircuitBreakerState::Open {
            status.health = ModelHealth::Failed;
        } else if status.consecutive_failures >= config.max_consecutive_failures {
            status.health = ModelHealth::Failed;
        } else if status.success_rate < config.min_success_rate {
            status.health = ModelHealth::Unhealthy;
        } else if status.avg_latency_us > config.max_latency_us as f64 {
            status.health = ModelHealth::Degraded;
        } else if status.accuracy_score < config.min_accuracy {
            status.health = ModelHealth::Degraded;
        } else {
            status.health = ModelHealth::Healthy;
        }

        // Trigger failover if health changed significantly
        if previous_health != status.health && status.health == ModelHealth::Failed {
            self.trigger_failover(model_id, FailoverEventType::ModelFailure)
                .await;
        } else if previous_health == ModelHealth::Healthy && status.health == ModelHealth::Degraded
        {
            self.trigger_failover(model_id, FailoverEventType::ModelDegraded)
                .await;
        }
    }

    /// Get best available model for prediction
    pub async fn get_best_available_model(&self) -> Option<String> {
        let status_map = self.model_status.read().await;
        let priorities = self.model_priorities.read().await;

        // Find highest priority healthy model
        for (_, model_ids) in priorities.iter().rev() {
            for model_id in model_ids {
                if let Some(status) = status_map.get(model_id) {
                    if status.enabled && status.health == ModelHealth::Healthy {
                        return Some(model_id.clone());
                    }
                }
            }
        }

        // Fall back to degraded models if no healthy ones available
        for (_, model_ids) in priorities.iter().rev() {
            for model_id in model_ids {
                if let Some(status) = status_map.get(model_id) {
                    if status.enabled && status.health == ModelHealth::Degraded {
                        return Some(model_id.clone());
                    }
                }
            }
        }

        None
    }

    /// Get ensemble of available models
    pub async fn get_ensemble_models(&self, max_models: usize) -> Vec<String> {
        let status_map = self.model_status.read().await;
        let priorities = self.model_priorities.read().await;

        let mut ensemble = Vec::new();

        // Collect healthy models in priority order
        for (_, model_ids) in priorities.iter().rev() {
            for model_id in model_ids {
                if ensemble.len() >= max_models {
                    break;
                }

                if let Some(status) = status_map.get(model_id) {
                    if status.enabled
                        && (status.health == ModelHealth::Healthy
                            || status.health == ModelHealth::Degraded)
                    {
                        ensemble.push(model_id.clone());
                    }
                }
            }

            if ensemble.len() >= max_models {
                break;
            }
        }

        ensemble
    }

    /// Execute prediction with fallback strategy
    pub async fn predict_with_fallback(
        &self,
        features: &[f64],
        preferred_model: Option<String>,
    ) -> FallbackPrediction {
        let start_time = Instant::now();
        let mut warnings = Vec::new();

        // Try preferred model first
        if let Some(model_id) = preferred_model {
            if self.is_model_available(&model_id).await {
                match self.execute_model_prediction(&model_id, features).await {
                    Ok(prediction) => {
                        return FallbackPrediction {
                            prediction_value: prediction,
                            confidence: 0.9,
                            models_used: vec![model_id],
                            strategy_used: FallbackStrategy::PriorityBased,
                            fallback_triggered: false,
                            latency_us: start_time.elapsed().as_micros() as u64,
                            warnings,
                        };
                    },
                    Err(e) => {
                        warnings.push(format!("Preferred model {} failed: {}", model_id, e));
                    },
                }
            }
        }

        // Try best available model
        if let Some(model_id) = self.get_best_available_model().await {
            match self.execute_model_prediction(&model_id, features).await {
                Ok(prediction) => {
                    return FallbackPrediction {
                        prediction_value: prediction,
                        confidence: 0.85,
                        models_used: vec![model_id],
                        strategy_used: FallbackStrategy::PriorityBased,
                        fallback_triggered: true,
                        latency_us: start_time.elapsed().as_micros() as u64,
                        warnings,
                    };
                },
                Err(e) => {
                    warnings.push(format!("Best available model failed: {}", e));
                },
            }
        }

        // Try ensemble prediction
        let ensemble_models = self.get_ensemble_models(3).await;
        if !ensemble_models.is_empty() {
            let mut predictions = Vec::new();
            let mut successful_models = Vec::new();

            for model_id in &ensemble_models {
                if let Ok(prediction) = self.execute_model_prediction(model_id, features).await {
                    predictions.push(prediction);
                    successful_models.push(model_id.clone());
                }
            }

            if !predictions.is_empty() {
                let ensemble_prediction =
                    predictions.iter().sum::<f64>() / predictions.len() as f64;
                return FallbackPrediction {
                    prediction_value: ensemble_prediction,
                    confidence: 0.75,
                    models_used: successful_models,
                    strategy_used: FallbackStrategy::EnsembleBased,
                    fallback_triggered: true,
                    latency_us: start_time.elapsed().as_micros() as u64,
                    warnings,
                };
            }
        }

        warnings.push("All ML models failed, using rule-based fallback".to_string());

        // Rule-based fallback
        let rule_prediction = self.rule_based_prediction(features);
        FallbackPrediction {
            prediction_value: rule_prediction,
            confidence: 0.5,
            models_used: vec!["rule_based".to_string()],
            strategy_used: FallbackStrategy::RuleBasedFallback,
            fallback_triggered: true,
            latency_us: start_time.elapsed().as_micros() as u64,
            warnings,
        }
    }

    /// Check if model is available for predictions
    async fn is_model_available(&self, model_id: &str) -> bool {
        let status_map = self.model_status.read().await;

        if let Some(status) = status_map.get(model_id) {
            status.enabled
                && status.health != ModelHealth::Failed
                && status.circuit_breaker_state != CircuitBreakerState::Open
        } else {
            false
        }
    }

    /// Execute prediction for a specific model (simulation)
    async fn execute_model_prediction(
        &self,
        model_id: &str,
        features: &[f64],
    ) -> Result<f64, String> {
        // Simulate model execution (in production, this would call actual models)
        if features.is_empty() {
            return Err("Empty features".to_string());
        }

        let prediction = match model_id {
            id if id.contains("dqn") => {
                // DQN-style prediction
                let momentum = features.first().copied().unwrap_or(0.0);
                let volume = features.get(1).copied().unwrap_or(0.0);
                0.5 + (momentum * 0.3) + (volume * 0.1).tanh() * 0.2
            },
            id if id.contains("transformer") => {
                // Transformer-style prediction
                let feature_sum: f64 = features.iter().sum();
                let normalized = feature_sum / features.len() as f64;
                0.5 + normalized.tanh() * 0.35
            },
            _ => {
                // Default linear model
                let avg = features.iter().sum::<f64>() / features.len() as f64;
                0.5 + avg.tanh() * 0.25
            },
        };

        Ok(prediction.clamp(0.0, 1.0))
    }

    /// Rule-based fallback prediction
    fn rule_based_prediction(&self, features: &[f64]) -> f64 {
        if features.is_empty() {
            return 0.5; // Neutral prediction
        }

        // Simple rule: positive momentum = buy signal
        let momentum = features.first().copied().unwrap_or(0.0);
        let volume = features.get(1).copied().unwrap_or(0.0);

        let base_prediction = 0.5;
        let momentum_signal = momentum.clamp(-0.1, 0.1) * 2.0; // Scale momentum
        let volume_signal = if volume > 0.0 { 0.05 } else { -0.02 }; // Volume bias

        (base_prediction + momentum_signal + volume_signal).clamp(0.0, 1.0)
    }

    /// Trigger failover event
    async fn trigger_failover(&self, failed_model: &str, event_type: FailoverEventType) {
        let fallback_model = self.get_best_available_model().await;

        let strategy = if fallback_model.is_some() {
            FallbackStrategy::PriorityBased
        } else {
            FallbackStrategy::RuleBasedFallback
        };

        let impact = match event_type {
            FailoverEventType::ModelFailure => FailoverImpact::Medium,
            FailoverEventType::ModelDegraded => FailoverImpact::Low,
            FailoverEventType::CircuitBreakerOpen => FailoverImpact::High,
            _ => FailoverImpact::Low,
        };

        let event = FailoverEvent {
            timestamp: SystemTime::now(),
            event_type,
            failed_model: Some(failed_model.to_string()),
            fallback_model,
            strategy,
            message: format!(
                "Failover triggered for model {} due to {:?}",
                failed_model, event_type
            ),
            impact,
        };

        // Store event
        {
            let mut events = self.failover_events.write().await;
            events.push(event.clone());

            // Keep only recent events (last 100)
            if events.len() > 100 {
                events.remove(0);
            }
        }

        // Broadcast event
        if let Err(e) = self.event_broadcaster.send(event.clone()) {
            debug!("Failed to broadcast failover event: {}", e);
        }

        match impact {
            FailoverImpact::Critical | FailoverImpact::High => {
                error!("ML Failover: {}", event.message);
            },
            FailoverImpact::Medium => {
                warn!("ML Failover: {}", event.message);
            },
            _ => {
                info!("ML Failover: {}", event.message);
            },
        }
    }

    /// Get model status
    pub async fn get_model_status(&self, model_id: &str) -> Option<ModelStatus> {
        let status_map = self.model_status.read().await;
        status_map.get(model_id).cloned()
    }

    /// Get all model statuses
    pub async fn get_all_model_statuses(&self) -> HashMap<String, ModelStatus> {
        self.model_status.read().await.clone()
    }

    /// Get recent failover events
    pub async fn get_recent_failover_events(&self, limit: usize) -> Vec<FailoverEvent> {
        let events = self.failover_events.read().await;
        events.iter().rev().take(limit).cloned().collect()
    }

    /// Subscribe to failover events
    pub fn subscribe_failover_events(&self) -> broadcast::Receiver<FailoverEvent> {
        self.event_broadcaster.subscribe()
    }

    /// Update configuration
    pub async fn update_config(&self, new_config: FallbackConfig) {
        let mut config = self.config.write().await;
        *config = new_config;
        info!("Fallback configuration updated");
    }

    /// Manually trigger model switch
    pub async fn switch_primary_model(&self, new_primary: String) -> Result<(), String> {
        if !self.is_model_available(&new_primary).await {
            return Err(format!("Model {} is not available", new_primary));
        }

        let old_primary = {
            let mut current_primary = self.current_primary.write().await;
            let old = current_primary.clone();
            *current_primary = Some(new_primary.clone());
            old
        };

        let event = FailoverEvent {
            timestamp: SystemTime::now(),
            event_type: FailoverEventType::ManualSwitching,
            failed_model: old_primary,
            fallback_model: Some(new_primary),
            strategy: FallbackStrategy::PriorityBased,
            message: "Manual model switch performed".to_string(),
            impact: FailoverImpact::None,
        };

        // Store and broadcast event
        {
            let mut events = self.failover_events.write().await;
            events.push(event.clone());
        }

        let _ = self.event_broadcaster.send(event);
        Ok(())
    }

    /// Generate fallback system health report
    pub async fn generate_health_report(&self) -> String {
        let all_statuses = self.get_all_model_statuses().await;
        let recent_events = self.get_recent_failover_events(10).await;

        let mut report = String::new();
        report.push_str("=== ML Fallback System Health Report ===\n\n");

        // Model status summary
        let healthy_count = all_statuses
            .values()
            .filter(|s| s.health == ModelHealth::Healthy)
            .count();
        let total_count = all_statuses.len();

        report.push_str(&format!(
            "Models: {}/{} healthy\n",
            healthy_count, total_count
        ));

        // Individual model status
        for (model_id, status) in &all_statuses {
            report.push_str(&format!(
                "  {}: {:?} (Success Rate: {:.2}%, Latency: {:.1}μs)\n",
                model_id,
                status.health,
                status.success_rate * 100.0,
                status.avg_latency_us
            ));
        }

        // Recent failover events
        report.push_str(&format!(
            "\nRecent Failover Events: {}\n",
            recent_events.len()
        ));
        for event in recent_events.into_iter().take(5) {
            report.push_str(&format!("  {:?}: {}\n", event.event_type, event.message));
        }

        report
    }
}

impl Default for MLFallbackManager {
    fn default() -> Self {
        Self::new()
    }
}
