//! Shared Circuit Breaker for ML Training
//!
//! Provides dynamic throttling to prevent runaway losses during training.
//! Used by both DQN and PPO training pipelines. This is the single source
//! of truth -- DQN and PPO re-export from here.

use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use ::common::resilience::circuit_breaker::{
    CircuitBreakerState as CommonCBState, CircuitBreakerTrait,
};
use parking_lot::RwLock;
use tracing::{debug, info, warn};

/// Circuit breaker state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
    /// Circuit is closed - trading allowed
    Closed,
    /// Circuit is open - trading blocked (cooldown period)
    Open,
    /// Circuit is half-open - limited trading to test recovery
    HalfOpen,
}

/// Configuration for the circuit breaker
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CircuitBreakerConfig {
    /// Number of consecutive failures before opening
    pub failure_threshold: usize,
    /// Number of consecutive successes needed to close from half-open
    pub success_threshold: usize,
    /// Cooldown duration when circuit opens (in seconds)
    #[serde(with = "duration_serde")]
    pub timeout_duration: Duration,
    /// Maximum number of test calls in half-open state
    pub half_open_max_calls: usize,
}

// Helper module for Duration serialization
mod duration_serde {
    use serde::{Deserialize, Deserializer, Serialize, Serializer};
    use std::time::Duration;

    pub(super) fn serialize<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        duration.as_secs().serialize(serializer)
    }

    pub(super) fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
    where
        D: Deserializer<'de>,
    {
        let secs = u64::deserialize(deserializer)?;
        Ok(Duration::from_secs(secs))
    }
}

impl Default for CircuitBreakerConfig {
    fn default() -> Self {
        Self {
            failure_threshold: 5,
            success_threshold: 3,
            timeout_duration: Duration::from_secs(60),
            half_open_max_calls: 2,
        }
    }
}

/// Simplified circuit breaker for ML training
#[derive(Debug)]
pub struct CircuitBreaker {
    config: CircuitBreakerConfig,
    state: Arc<RwLock<CircuitState>>,
    consecutive_failures: AtomicU32,
    consecutive_successes: AtomicU32,
    half_open_calls: AtomicU32,
    open_timestamp: Arc<RwLock<Option<Instant>>>,
    total_failures: AtomicU64,
    total_successes: AtomicU64,
}

impl CircuitBreaker {
    /// Create new circuit breaker with configuration
    pub fn new(config: CircuitBreakerConfig) -> Self {
        info!(
            "Initializing Circuit Breaker - failure_threshold={}, success_threshold={}, timeout={}s",
            config.failure_threshold,
            config.success_threshold,
            config.timeout_duration.as_secs()
        );

        Self {
            config,
            state: Arc::new(RwLock::new(CircuitState::Closed)),
            consecutive_failures: AtomicU32::new(0),
            consecutive_successes: AtomicU32::new(0),
            half_open_calls: AtomicU32::new(0),
            open_timestamp: Arc::new(RwLock::new(None)),
            total_failures: AtomicU64::new(0),
            total_successes: AtomicU64::new(0),
        }
    }

    /// Check if request should be allowed through
    pub fn allow_request(&self) -> bool {
        let current_state = *self.state.read();

        match current_state {
            CircuitState::Closed => true,
            CircuitState::Open => {
                if let Some(open_time) = *self.open_timestamp.read() {
                    if open_time.elapsed() >= self.config.timeout_duration {
                        info!(
                            "Circuit breaker transitioning to HALF-OPEN after {}s cooldown",
                            self.config.timeout_duration.as_secs()
                        );
                        *self.state.write() = CircuitState::HalfOpen;
                        self.half_open_calls.store(0, Ordering::SeqCst);
                        true
                    } else {
                        debug!(
                            "Circuit breaker OPEN - blocking request ({}s remaining)",
                            (self.config.timeout_duration - open_time.elapsed()).as_secs()
                        );
                        false
                    }
                } else {
                    warn!("Circuit breaker OPEN but no timestamp - allowing request");
                    true
                }
            }
            CircuitState::HalfOpen => {
                let calls = self.half_open_calls.fetch_add(1, Ordering::SeqCst);
                if calls < self.config.half_open_max_calls as u32 {
                    debug!(
                        "Circuit breaker HALF-OPEN - allowing test call {}/{}",
                        calls + 1,
                        self.config.half_open_max_calls
                    );
                    true
                } else {
                    debug!("Circuit breaker HALF-OPEN - max test calls reached");
                    false
                }
            }
        }
    }

    /// Record a successful operation
    pub fn record_success(&self) {
        self.total_successes.fetch_add(1, Ordering::Relaxed);
        self.consecutive_failures.store(0, Ordering::SeqCst);
        let successes = self.consecutive_successes.fetch_add(1, Ordering::SeqCst) + 1;

        let current_state = *self.state.read();

        match current_state {
            CircuitState::Closed => {
                debug!(
                    "Circuit breaker CLOSED - success recorded ({} consecutive)",
                    successes
                );
            }
            CircuitState::HalfOpen => {
                if successes >= self.config.success_threshold as u32 {
                    info!(
                        "Circuit breaker transitioning to CLOSED after {} consecutive successes",
                        successes
                    );
                    *self.state.write() = CircuitState::Closed;
                    self.consecutive_successes.store(0, Ordering::SeqCst);
                    *self.open_timestamp.write() = None;
                } else {
                    debug!(
                        "Circuit breaker HALF-OPEN - success {}/{}",
                        successes, self.config.success_threshold
                    );
                }
            }
            CircuitState::Open => {
                warn!("Circuit breaker OPEN but received success - resetting");
                *self.state.write() = CircuitState::Closed;
                self.consecutive_successes.store(0, Ordering::SeqCst);
                *self.open_timestamp.write() = None;
            }
        }
    }

    /// Record a failed operation
    pub fn record_failure(&self) {
        self.total_failures.fetch_add(1, Ordering::Relaxed);
        self.consecutive_successes.store(0, Ordering::SeqCst);
        let failures = self.consecutive_failures.fetch_add(1, Ordering::SeqCst) + 1;

        let current_state = *self.state.read();

        match current_state {
            CircuitState::Closed => {
                if failures >= self.config.failure_threshold as u32 {
                    warn!(
                        "Circuit breaker transitioning to OPEN after {} consecutive failures",
                        failures
                    );
                    *self.state.write() = CircuitState::Open;
                    *self.open_timestamp.write() = Some(Instant::now());
                    self.consecutive_failures.store(0, Ordering::SeqCst);
                } else {
                    debug!(
                        "Circuit breaker CLOSED - failure {}/{}",
                        failures, self.config.failure_threshold
                    );
                }
            }
            CircuitState::HalfOpen => {
                warn!("Circuit breaker HALF-OPEN - failure detected, returning to OPEN");
                *self.state.write() = CircuitState::Open;
                *self.open_timestamp.write() = Some(Instant::now());
                self.consecutive_failures.store(0, Ordering::SeqCst);
                self.half_open_calls.store(0, Ordering::SeqCst);
            }
            CircuitState::Open => {
                debug!("Circuit breaker OPEN - additional failure recorded");
            }
        }
    }

    /// Get current circuit state
    pub fn current_state(&self) -> CircuitState {
        *self.state.read()
    }

    /// Get statistics
    pub fn stats(&self) -> CircuitBreakerStats {
        CircuitBreakerStats {
            state: self.current_state(),
            consecutive_failures: self.consecutive_failures.load(Ordering::Relaxed),
            consecutive_successes: self.consecutive_successes.load(Ordering::Relaxed),
            total_failures: self.total_failures.load(Ordering::Relaxed),
            total_successes: self.total_successes.load(Ordering::Relaxed),
        }
    }

    /// Reset the circuit breaker to closed state
    pub fn reset(&self) {
        info!("Manually resetting circuit breaker to CLOSED state");
        *self.state.write() = CircuitState::Closed;
        self.consecutive_failures.store(0, Ordering::SeqCst);
        self.consecutive_successes.store(0, Ordering::SeqCst);
        self.half_open_calls.store(0, Ordering::SeqCst);
        *self.open_timestamp.write() = None;
    }
}

#[async_trait::async_trait]
impl CircuitBreakerTrait for CircuitBreaker {
    async fn can_execute(&self) -> bool {
        self.allow_request()
    }

    async fn record_success(&self) {
        CircuitBreaker::record_success(self);
    }

    async fn record_failure(&self) {
        CircuitBreaker::record_failure(self);
    }

    async fn state(&self) -> CommonCBState {
        match self.current_state() {
            CircuitState::Closed => CommonCBState::Closed,
            CircuitState::Open => CommonCBState::Open,
            CircuitState::HalfOpen => CommonCBState::HalfOpen,
        }
    }

    async fn reset(&self) {
        CircuitBreaker::reset(self);
    }
}

/// Circuit breaker statistics
#[derive(Debug, Clone)]
pub struct CircuitBreakerStats {
    pub state: CircuitState,
    pub consecutive_failures: u32,
    pub consecutive_successes: u32,
    pub total_failures: u64,
    pub total_successes: u64,
}

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

    #[test]
    fn test_circuit_breaker_closed_state() {
        let config = CircuitBreakerConfig::default();
        let breaker = CircuitBreaker::new(config);

        assert_eq!(breaker.current_state(), CircuitState::Closed);
        assert!(breaker.allow_request());
    }

    #[test]
    fn test_circuit_breaker_opens_after_failures() {
        let config = CircuitBreakerConfig {
            failure_threshold: 3,
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        breaker.record_failure();
        breaker.record_failure();
        assert_eq!(breaker.current_state(), CircuitState::Closed);
        assert!(breaker.allow_request());

        breaker.record_failure();
        assert_eq!(breaker.current_state(), CircuitState::Open);
        assert!(!breaker.allow_request());
    }

    #[test]
    fn test_circuit_breaker_success_resets_failures() {
        let config = CircuitBreakerConfig {
            failure_threshold: 3,
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        breaker.record_failure();
        breaker.record_failure();
        breaker.record_success();
        breaker.record_failure();
        breaker.record_failure();
        assert_eq!(breaker.current_state(), CircuitState::Closed);
    }

    #[test]
    fn test_circuit_breaker_half_open_transition() {
        let config = CircuitBreakerConfig {
            failure_threshold: 2,
            timeout_duration: Duration::from_millis(100),
            half_open_max_calls: 2,
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        breaker.record_failure();
        breaker.record_failure();
        assert_eq!(breaker.current_state(), CircuitState::Open);

        thread::sleep(Duration::from_millis(150));

        assert!(breaker.allow_request());
        assert_eq!(breaker.current_state(), CircuitState::HalfOpen);
    }

    #[test]
    fn test_circuit_breaker_closes_from_half_open() {
        let config = CircuitBreakerConfig {
            failure_threshold: 2,
            success_threshold: 2,
            timeout_duration: Duration::from_millis(100),
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        breaker.record_failure();
        breaker.record_failure();

        thread::sleep(Duration::from_millis(150));
        breaker.allow_request();

        breaker.record_success();
        assert_eq!(breaker.current_state(), CircuitState::HalfOpen);

        breaker.record_success();
        assert_eq!(breaker.current_state(), CircuitState::Closed);
    }

    #[test]
    fn test_circuit_breaker_reopens_from_half_open() {
        let config = CircuitBreakerConfig {
            failure_threshold: 2,
            timeout_duration: Duration::from_millis(100),
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        breaker.record_failure();
        breaker.record_failure();

        thread::sleep(Duration::from_millis(150));
        breaker.allow_request();
        assert_eq!(breaker.current_state(), CircuitState::HalfOpen);

        breaker.record_failure();
        assert_eq!(breaker.current_state(), CircuitState::Open);
    }

    #[test]
    fn test_circuit_breaker_stats() {
        let config = CircuitBreakerConfig::default();
        let breaker = CircuitBreaker::new(config);

        breaker.record_success();
        breaker.record_success();
        breaker.record_failure();

        let stats = breaker.stats();
        assert_eq!(stats.total_successes, 2);
        assert_eq!(stats.total_failures, 1);
        assert_eq!(stats.consecutive_successes, 0);
        assert_eq!(stats.consecutive_failures, 1);
    }

    #[test]
    fn test_circuit_breaker_reset() {
        let config = CircuitBreakerConfig {
            failure_threshold: 2,
            ..Default::default()
        };
        let breaker = CircuitBreaker::new(config);

        breaker.record_failure();
        breaker.record_failure();
        assert_eq!(breaker.current_state(), CircuitState::Open);

        breaker.reset();
        assert_eq!(breaker.current_state(), CircuitState::Closed);
        assert!(breaker.allow_request());

        let stats = breaker.stats();
        assert_eq!(stats.consecutive_failures, 0);
    }
}
