//! Ensemble Risk Manager for ML Production Trading
//!
//! Integrates risk management system with ML ensemble to provide:
//! - Prediction confidence thresholds
//! - Per-model circuit breakers
//! - VaR integration
//! - Cascade failure detection
//! - Model health monitoring

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};

use common::types::{Price, Symbol};
use ml::ensemble::EnsembleDecision;
use ml::{MLError, MLResult};
use risk::circuit_breaker::{BrokerAccountService, CircuitBreakerConfig, RealCircuitBreaker};
use risk::var_calculator::var_engine::PositionInfo;

/// Ensemble risk manager configuration
#[derive(Debug, Clone)]
pub struct EnsembleRiskConfig {
    /// Minimum confidence threshold to accept predictions (0.0-1.0)
    pub min_confidence_threshold: f64,

    /// Maximum consecutive errors before disabling a model
    pub max_consecutive_errors: u32,

    /// Number of models that can fail before halting ensemble
    pub cascade_failure_threshold: usize,

    /// Cascade detection window (seconds)
    pub cascade_detection_window_secs: u64,

    /// Enable VaR validation for predictions
    pub enable_var_validation: bool,

    /// Maximum disagreement rate before rejecting prediction (0.0-1.0)
    pub max_disagreement_rate: f64,

    /// Cooldown period before re-enabling a failed model (seconds)
    pub model_cooldown_period_secs: u64,
}

impl Default for EnsembleRiskConfig {
    fn default() -> Self {
        Self {
            min_confidence_threshold: 0.60, // 60% confidence minimum
            max_consecutive_errors: 3,
            cascade_failure_threshold: 2, // 2+ models fail = cascade
            cascade_detection_window_secs: 60,
            enable_var_validation: true,
            max_disagreement_rate: 0.50,     // 50% disagreement max
            model_cooldown_period_secs: 300, // 5 minutes
        }
    }
}

/// Model health status
#[derive(Debug, Clone)]
pub struct ModelHealth {
    pub model_id: String,
    pub enabled: bool,
    pub consecutive_errors: u32,
    pub last_error_time: Option<Instant>,
    pub total_predictions: u64,
    pub successful_predictions: u64,
    pub failed_predictions: u64,
    pub disabled_at: Option<Instant>,
    pub cooldown_until: Option<Instant>,
}

impl ModelHealth {
    fn new(model_id: String) -> Self {
        Self {
            model_id,
            enabled: true,
            consecutive_errors: 0,
            last_error_time: None,
            total_predictions: 0,
            successful_predictions: 0,
            failed_predictions: 0,
            disabled_at: None,
            cooldown_until: None,
        }
    }

    fn is_in_cooldown(&self) -> bool {
        if let Some(cooldown_until) = self.cooldown_until {
            Instant::now() < cooldown_until
        } else {
            false
        }
    }

    fn error_rate(&self) -> f64 {
        if self.total_predictions == 0 {
            0.0
        } else {
            self.failed_predictions as f64 / self.total_predictions as f64
        }
    }
}

/// Cascade failure detection state
#[derive(Debug, Clone)]
pub struct CascadeState {
    pub failed_models: Vec<String>,
    pub failure_window_start: Instant,
    pub is_cascading: bool,
}

impl CascadeState {
    fn new() -> Self {
        Self {
            failed_models: Vec::new(),
            failure_window_start: Instant::now(),
            is_cascading: false,
        }
    }

    fn reset(&mut self) {
        self.failed_models.clear();
        self.failure_window_start = Instant::now();
        self.is_cascading = false;
    }
}

/// Risk validation result for ensemble predictions
#[derive(Debug, Clone)]
pub struct RiskValidationResult {
    pub approved: bool,
    pub rejection_reason: Option<String>,
    pub confidence: f64,
    pub disagreement_rate: f64,
    pub var_validated: bool,
    pub cascade_detected: bool,
    pub disabled_models: Vec<String>,
    pub validation_latency_us: u64,
}

impl RiskValidationResult {
    pub fn rejected(reason: String, confidence: f64, disagreement_rate: f64) -> Self {
        Self {
            approved: false,
            rejection_reason: Some(reason),
            confidence,
            disagreement_rate,
            var_validated: false,
            cascade_detected: false,
            disabled_models: Vec::new(),
            validation_latency_us: 0,
        }
    }

    pub fn approved(confidence: f64, disagreement_rate: f64) -> Self {
        Self {
            approved: true,
            rejection_reason: None,
            confidence,
            disagreement_rate,
            var_validated: true,
            cascade_detected: false,
            disabled_models: Vec::new(),
            validation_latency_us: 0,
        }
    }
}

/// Ensemble Risk Manager
pub struct EnsembleRiskManager {
    config: EnsembleRiskConfig,
    model_health: Arc<RwLock<HashMap<String, ModelHealth>>>,
    cascade_state: Arc<RwLock<CascadeState>>,
    circuit_breaker: Option<Arc<RealCircuitBreaker>>,
}

impl EnsembleRiskManager {
    /// Create new ensemble risk manager
    pub fn new(config: EnsembleRiskConfig) -> Self {
        Self {
            config,
            model_health: Arc::new(RwLock::new(HashMap::new())),
            cascade_state: Arc::new(RwLock::new(CascadeState::new())),
            circuit_breaker: None,
        }
    }

    /// Create with circuit breaker integration
    pub async fn with_circuit_breaker(
        config: EnsembleRiskConfig,
        circuit_breaker_config: CircuitBreakerConfig,
        broker_service: Arc<dyn BrokerAccountService>,
    ) -> MLResult<Self> {
        let circuit_breaker = RealCircuitBreaker::new(circuit_breaker_config, broker_service)
            .await
            .map_err(|e| {
                MLError::ConfigError(format!("Circuit breaker init failed: {}", e))
            })?;

        Ok(Self {
            config,
            model_health: Arc::new(RwLock::new(HashMap::new())),
            cascade_state: Arc::new(RwLock::new(CascadeState::new())),
            circuit_breaker: Some(Arc::new(circuit_breaker)),
        })
    }

    /// Register a model for health monitoring
    pub async fn register_model(&self, model_id: String) -> MLResult<()> {
        let mut health_map = self.model_health.write().await;

        if !health_map.contains_key(&model_id) {
            health_map.insert(model_id.clone(), ModelHealth::new(model_id.clone()));
            info!("Registered model {} for risk monitoring", model_id);
        }

        Ok(())
    }

    /// Validate ensemble prediction against risk thresholds
    pub async fn validate_prediction(
        &self,
        decision: &EnsembleDecision,
        account_id: &str,
    ) -> MLResult<RiskValidationResult> {
        let start_time = Instant::now();

        // Check confidence threshold
        if decision.confidence < self.config.min_confidence_threshold {
            warn!(
                "Prediction rejected: confidence {:.3} < threshold {:.3}",
                decision.confidence, self.config.min_confidence_threshold
            );
            return Ok(RiskValidationResult::rejected(
                format!(
                    "Low confidence: {:.3} < {:.3}",
                    decision.confidence, self.config.min_confidence_threshold
                ),
                decision.confidence,
                decision.disagreement_rate,
            ));
        }

        // Check disagreement rate
        if decision.disagreement_rate > self.config.max_disagreement_rate {
            warn!(
                "Prediction rejected: disagreement {:.3} > threshold {:.3}",
                decision.disagreement_rate, self.config.max_disagreement_rate
            );
            return Ok(RiskValidationResult::rejected(
                format!(
                    "High disagreement: {:.3} > {:.3}",
                    decision.disagreement_rate, self.config.max_disagreement_rate
                ),
                decision.confidence,
                decision.disagreement_rate,
            ));
        }

        // Check cascade failure state
        let cascade_state = self.cascade_state.read().await;
        if cascade_state.is_cascading {
            error!("Prediction rejected: cascade failure detected");
            return Ok(RiskValidationResult::rejected(
                "Cascade failure: 2+ models failed".to_string(),
                decision.confidence,
                decision.disagreement_rate,
            ));
        }
        drop(cascade_state);

        // Check circuit breaker if available
        if let Some(ref circuit_breaker) = self.circuit_breaker {
            let circuit_active = circuit_breaker.is_active(account_id).await;
            if circuit_active {
                warn!(
                    "Prediction rejected: circuit breaker active for account {}",
                    account_id
                );
                return Ok(RiskValidationResult::rejected(
                    "Circuit breaker active".to_string(),
                    decision.confidence,
                    decision.disagreement_rate,
                ));
            }
        }

        // Get disabled models
        let health_map = self.model_health.read().await;
        let disabled_models: Vec<String> = health_map
            .values()
            .filter(|h| !h.enabled)
            .map(|h| h.model_id.clone())
            .collect();

        let validation_latency_us = start_time.elapsed().as_micros() as u64;

        Ok(RiskValidationResult {
            approved: true,
            rejection_reason: None,
            confidence: decision.confidence,
            disagreement_rate: decision.disagreement_rate,
            var_validated: self.config.enable_var_validation,
            cascade_detected: false,
            disabled_models,
            validation_latency_us,
        })
    }

    /// Record model prediction result
    pub async fn record_prediction_result(&self, model_id: &str, success: bool) -> MLResult<()> {
        let mut health_map = self.model_health.write().await;

        let health = health_map
            .entry(model_id.to_string())
            .or_insert_with(|| ModelHealth::new(model_id.to_string()));

        health.total_predictions += 1;

        if success {
            health.successful_predictions += 1;
            health.consecutive_errors = 0;
            health.last_error_time = None;
        } else {
            health.failed_predictions += 1;
            health.consecutive_errors += 1;
            health.last_error_time = Some(Instant::now());

            // Check if model should be disabled
            if health.consecutive_errors >= self.config.max_consecutive_errors {
                self.disable_model_internal(health).await?;
            }
        }

        Ok(())
    }

    /// Disable a model due to consecutive errors
    async fn disable_model_internal(&self, health: &mut ModelHealth) -> MLResult<()> {
        if health.enabled {
            health.enabled = false;
            health.disabled_at = Some(Instant::now());
            health.cooldown_until =
                Some(Instant::now() + Duration::from_secs(self.config.model_cooldown_period_secs));

            error!(
                "Model {} disabled after {} consecutive errors (error rate: {:.1}%)",
                health.model_id,
                health.consecutive_errors,
                health.error_rate() * 100.0
            );

            // Update cascade state
            self.update_cascade_state(&health.model_id).await?;
        }

        Ok(())
    }

    /// Update cascade failure detection state
    async fn update_cascade_state(&self, failed_model_id: &str) -> MLResult<()> {
        let mut cascade_state = self.cascade_state.write().await;

        // Check if we're in detection window
        let elapsed = Instant::now().duration_since(cascade_state.failure_window_start);
        if elapsed.as_secs() > self.config.cascade_detection_window_secs {
            // Reset window
            cascade_state.reset();
        }

        // Add failed model
        if !cascade_state
            .failed_models
            .contains(&failed_model_id.to_string())
        {
            cascade_state
                .failed_models
                .push(failed_model_id.to_string());
        }

        // Check cascade threshold
        if cascade_state.failed_models.len() >= self.config.cascade_failure_threshold
            && !cascade_state.is_cascading {
                cascade_state.is_cascading = true;
                error!(
                    "CASCADE FAILURE DETECTED: {} models failed within {}s: {:?}",
                    cascade_state.failed_models.len(),
                    self.config.cascade_detection_window_secs,
                    cascade_state.failed_models
                );
            }

        Ok(())
    }

    /// Manually reset cascade failure state
    pub async fn reset_cascade_state(&self) -> MLResult<()> {
        let mut cascade_state = self.cascade_state.write().await;
        cascade_state.reset();
        info!("Cascade failure state manually reset");
        Ok(())
    }

    /// Re-enable a model after cooldown period
    pub async fn try_enable_model(&self, model_id: &str) -> MLResult<bool> {
        let mut health_map = self.model_health.write().await;

        if let Some(health) = health_map.get_mut(model_id) {
            if !health.enabled && !health.is_in_cooldown() {
                health.enabled = true;
                health.consecutive_errors = 0;
                health.disabled_at = None;
                health.cooldown_until = None;

                info!("Model {} re-enabled after cooldown period", model_id);
                return Ok(true);
            }
        }

        Ok(false)
    }

    /// Get model health status
    pub async fn get_model_health(&self, model_id: &str) -> Option<ModelHealth> {
        let health_map = self.model_health.read().await;
        health_map.get(model_id).cloned()
    }

    /// Get all model health statuses
    pub async fn get_all_model_health(&self) -> HashMap<String, ModelHealth> {
        self.model_health.read().await.clone()
    }

    /// Get cascade state
    pub async fn get_cascade_state(&self) -> CascadeState {
        self.cascade_state.read().await.clone()
    }

    /// Check if ensemble is operational
    pub async fn is_operational(&self) -> bool {
        let cascade_state = self.cascade_state.read().await;
        !cascade_state.is_cascading
    }

    /// Get count of enabled models
    pub async fn enabled_model_count(&self) -> usize {
        let health_map = self.model_health.read().await;
        health_map.values().filter(|h| h.enabled).count()
    }

    /// Validate VaR for portfolio positions
    pub async fn validate_var(
        &self,
        _portfolio_id: &str,
        _positions: &HashMap<Symbol, PositionInfo>,
        current_pnl: Price,
        portfolio_value: Price,
    ) -> MLResult<bool> {
        if !self.config.enable_var_validation {
            return Ok(true);
        }

        // For now, use simplified VaR check
        // In production, calculate comprehensive VaR
        let loss_percentage = if portfolio_value > Price::ZERO {
            current_pnl.to_f64() / portfolio_value.to_f64()
        } else {
            0.0
        };

        // 2% daily loss limit
        const MAX_LOSS_PCT: f64 = 0.02;

        if loss_percentage >= MAX_LOSS_PCT {
            warn!(
                "VaR validation failed: loss {:.2}% >= {:.2}%",
                loss_percentage * 100.0,
                MAX_LOSS_PCT * 100.0
            );
            return Ok(false);
        }

        debug!(
            "VaR validated: loss {:.2}% < {:.2}%",
            loss_percentage * 100.0,
            MAX_LOSS_PCT * 100.0
        );

        Ok(true)
    }
}

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

    #[tokio::test]
    async fn test_ensemble_risk_manager_creation() {
        let config = EnsembleRiskConfig::default();
        let manager = EnsembleRiskManager::new(config);

        assert!(manager.is_operational().await);
        assert_eq!(manager.enabled_model_count().await, 0);
    }

    #[tokio::test]
    async fn test_model_registration() {
        let config = EnsembleRiskConfig::default();
        let manager = EnsembleRiskManager::new(config);

        manager.register_model("DQN".to_string()).await.unwrap();
        manager.register_model("PPO".to_string()).await.unwrap();

        assert_eq!(manager.enabled_model_count().await, 2);
    }

    #[tokio::test]
    async fn test_low_confidence_rejection() {
        let config = EnsembleRiskConfig {
            min_confidence_threshold: 0.70,
            ..Default::default()
        };
        let manager = EnsembleRiskManager::new(config);

        let decision = EnsembleDecision::new(
            TradingAction::Buy,
            0.50, // Below threshold
            0.65,
            0.20,
            HashMap::new(),
        );

        let result = manager
            .validate_prediction(&decision, "TEST_ACCOUNT")
            .await
            .unwrap();

        assert!(!result.approved);
        assert!(result.rejection_reason.is_some());
        assert!(result.rejection_reason.unwrap().contains("Low confidence"));
    }

    #[tokio::test]
    async fn test_high_disagreement_rejection() {
        let config = EnsembleRiskConfig {
            max_disagreement_rate: 0.40,
            ..Default::default()
        };
        let manager = EnsembleRiskManager::new(config);

        let decision = EnsembleDecision::new(
            TradingAction::Buy,
            0.75,
            0.65,
            0.55, // Above threshold
            HashMap::new(),
        );

        let result = manager
            .validate_prediction(&decision, "TEST_ACCOUNT")
            .await
            .unwrap();

        assert!(!result.approved);
        assert!(result
            .rejection_reason
            .unwrap()
            .contains("High disagreement"));
    }

    #[tokio::test]
    async fn test_consecutive_errors_disable_model() {
        let config = EnsembleRiskConfig {
            max_consecutive_errors: 3,
            ..Default::default()
        };
        let manager = EnsembleRiskManager::new(config);

        manager.register_model("DQN".to_string()).await.unwrap();

        // Record 3 consecutive errors
        for _ in 0..3 {
            manager
                .record_prediction_result("DQN", false)
                .await
                .unwrap();
        }

        let health = manager.get_model_health("DQN").await.unwrap();
        assert!(!health.enabled);
        assert_eq!(health.consecutive_errors, 3);
    }

    #[tokio::test]
    async fn test_cascade_failure_detection() {
        let config = EnsembleRiskConfig {
            max_consecutive_errors: 2,
            cascade_failure_threshold: 2,
            ..Default::default()
        };
        let manager = EnsembleRiskManager::new(config);

        manager.register_model("DQN".to_string()).await.unwrap();
        manager.register_model("PPO".to_string()).await.unwrap();

        // Fail both models
        for _ in 0..2 {
            manager
                .record_prediction_result("DQN", false)
                .await
                .unwrap();
            manager
                .record_prediction_result("PPO", false)
                .await
                .unwrap();
        }

        let cascade_state = manager.get_cascade_state().await;
        assert!(cascade_state.is_cascading);
        assert_eq!(cascade_state.failed_models.len(), 2);
        assert!(!manager.is_operational().await);
    }

    #[tokio::test]
    async fn test_successful_predictions_reset_errors() {
        let config = EnsembleRiskConfig {
            max_consecutive_errors: 3,
            ..Default::default()
        };
        let manager = EnsembleRiskManager::new(config);

        manager.register_model("DQN".to_string()).await.unwrap();

        // Record 2 errors, then success
        manager
            .record_prediction_result("DQN", false)
            .await
            .unwrap();
        manager
            .record_prediction_result("DQN", false)
            .await
            .unwrap();
        manager.record_prediction_result("DQN", true).await.unwrap();

        let health = manager.get_model_health("DQN").await.unwrap();
        assert!(health.enabled);
        assert_eq!(health.consecutive_errors, 0);
    }

    #[tokio::test]
    async fn test_model_cooldown_period() {
        let config = EnsembleRiskConfig {
            max_consecutive_errors: 2,
            model_cooldown_period_secs: 1, // 1 second for testing
            ..Default::default()
        };
        let manager = EnsembleRiskManager::new(config);

        manager.register_model("DQN".to_string()).await.unwrap();

        // Disable model with 2 errors
        manager
            .record_prediction_result("DQN", false)
            .await
            .unwrap();
        manager
            .record_prediction_result("DQN", false)
            .await
            .unwrap();

        let health = manager.get_model_health("DQN").await.unwrap();
        assert!(!health.enabled);

        // Try to enable immediately (should fail - in cooldown)
        let enabled = manager.try_enable_model("DQN").await.unwrap();
        assert!(!enabled);

        // Wait for cooldown
        tokio::time::sleep(Duration::from_secs(2)).await;

        // Try again (should succeed)
        let enabled = manager.try_enable_model("DQN").await.unwrap();
        assert!(enabled);

        let health = manager.get_model_health("DQN").await.unwrap();
        assert!(health.enabled);
    }

    #[tokio::test]
    async fn test_approved_prediction() {
        let config = EnsembleRiskConfig::default();
        let manager = EnsembleRiskManager::new(config);

        let decision = EnsembleDecision::new(
            TradingAction::Buy,
            0.85, // Good confidence
            0.70,
            0.15, // Low disagreement
            HashMap::new(),
        );

        let result = manager
            .validate_prediction(&decision, "TEST_ACCOUNT")
            .await
            .unwrap();

        assert!(result.approved);
        assert!(result.rejection_reason.is_none());
    }
}
