//! Resilience Module - Circuit Breaker, Retry, and Timeout
//!
//! Provides three layers of defense for federation operations:
//! - **CircuitBreaker**: Prevents cascading failures by stopping requests to unhealthy stores
//! - **RetryExecutor**: Retries transient failures with exponential backoff
//! - **TimeoutController**: Enforces operation-level timeouts

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

use super::error::{FederationError, Result};

// =============================================================================
// Circuit Breaker
// =============================================================================

/// Circuit breaker state
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum CircuitState {
    /// Normal operation — all requests pass through
    Closed = 0,
    /// Breaker tripped — all requests are rejected
    Open = 1,
    /// Allowing a limited number of probe requests
    HalfOpen = 2,
}

impl From<u8> for CircuitState {
    fn from(v: u8) -> Self {
        match v {
            0 => CircuitState::Closed,
            1 => CircuitState::Open,
            2 => CircuitState::HalfOpen,
            _ => CircuitState::Closed,
        }
    }
}

/// Circuit breaker configuration
#[derive(Debug, Clone)]
pub struct CircuitBreakerConfig {
    /// Number of consecutive failures to trigger the breaker
    pub failure_threshold: u32,
    /// Seconds to wait before transitioning from Open to HalfOpen
    pub recovery_timeout_secs: u64,
    /// Number of probe requests allowed in HalfOpen state
    pub half_open_max_requests: u32,
}

impl Default for CircuitBreakerConfig {
    fn default() -> Self {
        Self {
            failure_threshold: 5,
            recovery_timeout_secs: 30,
            half_open_max_requests: 3,
        }
    }
}

/// Per-store circuit breaker instance
pub struct CircuitBreaker {
    /// Store alias (for error messages)
    alias: String,
    /// Configuration
    config: CircuitBreakerConfig,
    /// Current state (encoded as u8)
    state: AtomicU8,
    /// Consecutive failure count
    consecutive_failures: AtomicU32,
    /// Timestamp of the last failure (unix seconds)
    last_failure_time: AtomicU64,
    /// Number of requests allowed through in HalfOpen
    half_open_requests: AtomicU32,
}

impl CircuitBreaker {
    /// Create a new circuit breaker for the given store alias
    pub fn new(alias: &str, config: CircuitBreakerConfig) -> Self {
        Self {
            alias: alias.to_string(),
            config,
            state: AtomicU8::new(CircuitState::Closed as u8),
            consecutive_failures: AtomicU32::new(0),
            last_failure_time: AtomicU64::new(0),
            half_open_requests: AtomicU32::new(0),
        }
    }

    /// Create with default config
    pub fn with_defaults(alias: &str) -> Self {
        Self::new(alias, CircuitBreakerConfig::default())
    }

    /// Get current state
    pub fn state(&self) -> CircuitState {
        CircuitState::from(self.state.load(Ordering::Relaxed))
    }

    /// Check whether the next request is allowed to proceed
    pub fn allow_request(&self) -> Result<()> {
        match self.state() {
            CircuitState::Closed => Ok(()),
            CircuitState::Open => {
                if self.should_transition_to_half_open() {
                    self.set_state(CircuitState::HalfOpen);
                    self.half_open_requests.store(1, Ordering::Relaxed);
                    Ok(())
                } else {
                    Err(FederationError::CircuitBreakerOpen {
                        alias: self.alias.clone(),
                    })
                }
            }
            CircuitState::HalfOpen => {
                let current = self.half_open_requests.fetch_add(1, Ordering::Relaxed);
                if current < self.config.half_open_max_requests {
                    Ok(())
                } else {
                    Err(FederationError::CircuitBreakerOpen {
                        alias: self.alias.clone(),
                    })
                }
            }
        }
    }

    /// Record a successful operation — resets the failure counter
    pub fn record_success(&self) {
        self.consecutive_failures.store(0, Ordering::Relaxed);
        if self.state() == CircuitState::HalfOpen {
            // Probe succeeded → close the breaker
            self.set_state(CircuitState::Closed);
            self.half_open_requests.store(0, Ordering::Relaxed);
        }
    }

    /// Record a failed operation — may trip the breaker
    pub fn record_failure(&self) {
        let failures = self.consecutive_failures.fetch_add(1, Ordering::Relaxed) + 1;
        self.last_failure_time.store(Self::now_secs(), Ordering::Relaxed);

        match self.state() {
            CircuitState::Closed => {
                if failures >= self.config.failure_threshold {
                    self.set_state(CircuitState::Open);
                }
            }
            CircuitState::HalfOpen => {
                // Probe failed → back to Open
                self.set_state(CircuitState::Open);
                self.half_open_requests.store(0, Ordering::Relaxed);
            }
            CircuitState::Open => {
                // Already open, nothing to do
            }
        }
    }

    /// Get consecutive failure count
    pub fn failure_count(&self) -> u32 {
        self.consecutive_failures.load(Ordering::Relaxed)
    }

    /// Manually reset the breaker to Closed state
    pub fn reset(&self) {
        self.set_state(CircuitState::Closed);
        self.consecutive_failures.store(0, Ordering::Relaxed);
        self.half_open_requests.store(0, Ordering::Relaxed);
    }

    // -- private helpers --

    fn set_state(&self, new_state: CircuitState) {
        self.state.store(new_state as u8, Ordering::Relaxed);
    }

    fn should_transition_to_half_open(&self) -> bool {
        let last = self.last_failure_time.load(Ordering::Relaxed);
        let now = Self::now_secs();
        now.saturating_sub(last) >= self.config.recovery_timeout_secs
    }

    fn now_secs() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs()
    }
}

// =============================================================================
// Retry Executor
// =============================================================================

/// Which errors are safe to retry
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RetryableError {
    /// Store temporarily unavailable
    StoreUnavailable,
    /// Timeout
    Timeout,
    /// IO error (transient disk issues, etc.)
    IoError,
}

/// Retry configuration
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of retries (0 = no retry)
    pub max_retries: u32,
    /// Initial backoff duration in milliseconds
    pub initial_backoff_ms: u64,
    /// Maximum backoff duration in milliseconds
    pub max_backoff_ms: u64,
    /// Backoff multiplier (e.g. 2.0 for exponential)
    pub backoff_multiplier: f64,
    /// Which error kinds are retryable
    pub retryable_errors: Vec<RetryableError>,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_retries: 3,
            initial_backoff_ms: 100,
            max_backoff_ms: 5000,
            backoff_multiplier: 2.0,
            retryable_errors: vec![
                RetryableError::StoreUnavailable,
                RetryableError::Timeout,
                RetryableError::IoError,
            ],
        }
    }
}

/// Retry executor — wraps an operation with exponential-backoff retry logic
pub struct RetryExecutor {
    config: RetryConfig,
}

impl RetryExecutor {
    pub fn new(config: RetryConfig) -> Self {
        Self { config }
    }

    pub fn with_defaults() -> Self {
        Self::new(RetryConfig::default())
    }

    /// Execute `operation` with retries according to configuration.
    pub fn execute<F, T>(&self, operation: F) -> Result<T>
    where
        F: Fn() -> Result<T>,
    {
        let mut last_error: Option<FederationError> = None;
        let mut backoff_ms = self.config.initial_backoff_ms;

        for attempt in 0..=self.config.max_retries {
            match operation() {
                Ok(result) => return Ok(result),
                Err(e) => {
                    if attempt == self.config.max_retries || !self.is_retryable(&e) {
                        return Err(e);
                    }
                    last_error = Some(e);
                    std::thread::sleep(Duration::from_millis(backoff_ms));
                    backoff_ms =
                        ((backoff_ms as f64) * self.config.backoff_multiplier) as u64;
                    backoff_ms = backoff_ms.min(self.config.max_backoff_ms);
                }
            }
        }

        Err(last_error.unwrap_or_else(|| {
            FederationError::Internal("Retry exhausted with no error captured".to_string())
        }))
    }

    /// Check if an error is retryable
    pub fn is_retryable(&self, error: &FederationError) -> bool {
        match error {
            FederationError::StoreUnavailable(_) => {
                self.config.retryable_errors.contains(&RetryableError::StoreUnavailable)
            }
            FederationError::Timeout(_) => {
                self.config.retryable_errors.contains(&RetryableError::Timeout)
            }
            FederationError::IoError(_) => {
                self.config.retryable_errors.contains(&RetryableError::IoError)
            }
            _ => false,
        }
    }
}

// =============================================================================
// Timeout Controller
// =============================================================================

/// Timeout configuration for different operation types
#[derive(Debug, Clone)]
pub struct TimeoutConfig {
    /// Single-store query timeout (ms)
    pub query_single_timeout_ms: u64,
    /// Cross-store query timeout (ms)
    pub query_cross_timeout_ms: u64,
    /// Single write timeout (ms)
    pub write_single_timeout_ms: u64,
    /// Batch write timeout (ms)
    pub write_batch_timeout_ms: u64,
    /// Transaction idle timeout (ms)
    pub transaction_idle_timeout_ms: u64,
    /// Transaction commit timeout (ms)
    pub transaction_commit_timeout_ms: u64,
    /// Index rebuild timeout (ms)
    pub index_rebuild_timeout_ms: u64,
}

impl Default for TimeoutConfig {
    fn default() -> Self {
        Self {
            query_single_timeout_ms: 10_000,
            query_cross_timeout_ms: 30_000,
            write_single_timeout_ms: 5_000,
            write_batch_timeout_ms: 30_000,
            transaction_idle_timeout_ms: 60_000,
            transaction_commit_timeout_ms: 30_000,
            index_rebuild_timeout_ms: 300_000,
        }
    }
}

/// Operation kind — used to select the right timeout
#[derive(Debug, Clone, Copy)]
pub enum OperationKind {
    QuerySingle,
    QueryCross,
    WriteSingle,
    WriteBatch,
    TransactionCommit,
    IndexRebuild,
}

/// Timeout controller — enforces deadlines on operations
pub struct TimeoutController {
    config: TimeoutConfig,
}

impl TimeoutController {
    pub fn new(config: TimeoutConfig) -> Self {
        Self { config }
    }

    pub fn with_defaults() -> Self {
        Self::new(TimeoutConfig::default())
    }

    /// Get the timeout duration for a given operation kind
    pub fn timeout_for(&self, kind: OperationKind) -> Duration {
        let ms = match kind {
            OperationKind::QuerySingle => self.config.query_single_timeout_ms,
            OperationKind::QueryCross => self.config.query_cross_timeout_ms,
            OperationKind::WriteSingle => self.config.write_single_timeout_ms,
            OperationKind::WriteBatch => self.config.write_batch_timeout_ms,
            OperationKind::TransactionCommit => self.config.transaction_commit_timeout_ms,
            OperationKind::IndexRebuild => self.config.index_rebuild_timeout_ms,
        };
        Duration::from_millis(ms)
    }

    /// Execute an operation with timeout enforcement.
    ///
    /// This is a synchronous check-based approach: the caller passes a closure
    /// and we check elapsed time after it returns. For truly preemptive timeouts,
    /// use async runtimes or thread-based cancellation.
    pub fn execute_with_timeout<F, T>(&self, kind: OperationKind, operation: F) -> Result<T>
    where
        F: FnOnce() -> Result<T>,
    {
        let deadline = self.timeout_for(kind);
        let start = Instant::now();

        let result = operation();

        if start.elapsed() > deadline {
            return Err(FederationError::Timeout(format!(
                "{:?} exceeded timeout of {}ms (took {}ms)",
                kind,
                deadline.as_millis(),
                start.elapsed().as_millis(),
            )));
        }

        result
    }

    /// Check if a duration has exceeded the timeout for a given operation
    pub fn is_timed_out(&self, kind: OperationKind, elapsed: Duration) -> bool {
        elapsed > self.timeout_for(kind)
    }
}

// =============================================================================
// Resilience Config (aggregate)
// =============================================================================

/// Aggregate resilience configuration
#[derive(Debug, Clone)]
pub struct ResilienceConfig {
    pub circuit_breaker: CircuitBreakerConfig,
    pub retry: RetryConfig,
    pub timeout: TimeoutConfig,
}

impl Default for ResilienceConfig {
    fn default() -> Self {
        Self {
            circuit_breaker: CircuitBreakerConfig::default(),
            retry: RetryConfig::default(),
            timeout: TimeoutConfig::default(),
        }
    }
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn test_circuit_breaker_closed_by_default() {
        let cb = CircuitBreaker::with_defaults("test_store");
        assert_eq!(cb.state(), CircuitState::Closed);
        assert!(cb.allow_request().is_ok());
    }

    #[test]
    fn test_circuit_breaker_trips_after_threshold() {
        let cb = CircuitBreaker::new(
            "test_store",
            CircuitBreakerConfig {
                failure_threshold: 3,
                recovery_timeout_secs: 60,
                half_open_max_requests: 1,
            },
        );

        cb.record_failure();
        cb.record_failure();
        assert_eq!(cb.state(), CircuitState::Closed);

        cb.record_failure(); // 3rd failure → trips
        assert_eq!(cb.state(), CircuitState::Open);
        assert!(cb.allow_request().is_err());
    }

    #[test]
    fn test_circuit_breaker_success_resets() {
        let cb = CircuitBreaker::new(
            "test_store",
            CircuitBreakerConfig {
                failure_threshold: 2,
                recovery_timeout_secs: 60,
                half_open_max_requests: 1,
            },
        );

        cb.record_failure();
        cb.record_success(); // resets counter
        cb.record_failure();
        assert_eq!(cb.state(), CircuitState::Closed); // only 1 consecutive failure
    }

    #[test]
    fn test_circuit_breaker_reset() {
        let cb = CircuitBreaker::new(
            "test_store",
            CircuitBreakerConfig {
                failure_threshold: 1,
                recovery_timeout_secs: 9999,
                half_open_max_requests: 1,
            },
        );

        cb.record_failure();
        assert_eq!(cb.state(), CircuitState::Open);

        cb.reset();
        assert_eq!(cb.state(), CircuitState::Closed);
        assert!(cb.allow_request().is_ok());
    }

    #[test]
    fn test_retry_executor_immediate_success() {
        let executor = RetryExecutor::with_defaults();

        let result = executor.execute(|| {
            Ok(42)
        });

        assert_eq!(result.unwrap(), 42);
    }

    #[test]
    fn test_retry_executor_non_retryable_error() {
        let executor = RetryExecutor::with_defaults();

        let result: Result<i32> = executor.execute(|| {
            Err(FederationError::QueryError("bad query".to_string()))
        });

        assert!(result.is_err());
    }

    #[test]
    fn test_timeout_config_defaults() {
        let tc = TimeoutController::with_defaults();
        assert_eq!(
            tc.timeout_for(OperationKind::QuerySingle),
            Duration::from_millis(10_000)
        );
        assert_eq!(
            tc.timeout_for(OperationKind::WriteBatch),
            Duration::from_millis(30_000)
        );
    }

    #[test]
    fn test_timeout_not_exceeded() {
        let tc = TimeoutController::with_defaults();
        let result = tc.execute_with_timeout(OperationKind::QuerySingle, || Ok(123));
        assert_eq!(result.unwrap(), 123);
    }
}
