ironclad-llm 0.9.7

use std::collections::HashMap;
use std::time::{Duration, Instant};

use ironclad_core::config::CircuitBreakerConfig;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CircuitState {
    Closed,
    Open,
    HalfOpen,
}

#[derive(Debug)]
struct CircuitBreaker {
    state: CircuitState,
    failure_count: u32,
    last_failure_at: Option<Instant>,
    /// First failure timestamp in the current accumulation window.  The counter
    /// resets only when `now - window_start > window` (rolling window), NOT when
    /// the gap between two consecutive failures exceeds the window.
    window_start: Option<Instant>,
    cooldown: Duration,
    max_cooldown: Duration,
    threshold: u32,
    window: Duration,
    /// When true, the breaker was tripped by a credit/billing error and will
    /// NOT auto-recover to HalfOpen. It stays Open until explicitly `reset()`.
    credit_tripped: bool,
    /// When true, breaker was opened by an explicit operator kill-switch.
    operator_forced_open: bool,
    /// Soft-pressure signal from capacity monitoring; used to deprioritize a
    /// provider before hard failures occur.
    preemptive_half_open: bool,
}

impl CircuitBreaker {
    fn new(config: &CircuitBreakerConfig) -> Self {
        Self {
            state: CircuitState::Closed,
            failure_count: 0,
            last_failure_at: None,
            window_start: None,
            cooldown: Duration::from_secs(config.cooldown_seconds),
            max_cooldown: Duration::from_secs(config.max_cooldown_seconds),
            threshold: config.threshold,
            window: Duration::from_secs(config.window_seconds),
            credit_tripped: false,
            operator_forced_open: false,
            preemptive_half_open: false,
        }
    }

    fn effective_state(&self) -> CircuitState {
        match self.state {
            CircuitState::Open => {
                if self.credit_tripped || self.operator_forced_open {
                    return CircuitState::Open;
                }
                if let Some(last) = self.last_failure_at
                    && last.elapsed() >= self.cooldown
                {
                    return CircuitState::HalfOpen;
                }
                CircuitState::Open
            }
            CircuitState::Closed if self.preemptive_half_open => CircuitState::HalfOpen,
            other => other,
        }
    }
}

#[derive(Debug)]
pub struct CircuitBreakerRegistry {
    breakers: HashMap<String, CircuitBreaker>,
    config: CircuitBreakerConfig,
}

#[cfg(test)]
impl CircuitBreakerRegistry {
    fn force_half_open(&mut self, provider: &str) {
        let cb = self.get_or_create(provider);
        cb.last_failure_at = Some(Instant::now() - cb.cooldown - Duration::from_millis(1));
    }
}

impl CircuitBreakerRegistry {
    pub fn new(config: &CircuitBreakerConfig) -> Self {
        Self {
            breakers: HashMap::new(),
            config: config.clone(),
        }
    }

    fn get_or_create(&mut self, provider: &str) -> &mut CircuitBreaker {
        let config = self.config.clone();
        self.breakers
            .entry(provider.to_string())
            .or_insert_with(|| CircuitBreaker::new(&config))
    }

    pub fn is_blocked(&self, provider: &str) -> bool {
        match self.breakers.get(provider) {
            Some(cb) => cb.effective_state() == CircuitState::Open,
            None => false,
        }
    }

    pub fn record_success(&mut self, provider: &str) {
        let base_cooldown = self.config.cooldown_seconds;
        let cb = self.get_or_create(provider);
        match cb.effective_state() {
            CircuitState::HalfOpen => {
                if cb.preemptive_half_open {
                    cb.failure_count = 0;
                    cb.window_start = None;
                    cb.preemptive_half_open = false;
                    return;
                }
                cb.state = CircuitState::Closed;
                cb.failure_count = 0;
                cb.window_start = None;
                cb.cooldown = Duration::from_secs(base_cooldown);
            }
            CircuitState::Closed => {
                cb.failure_count = 0;
                cb.window_start = None;
            }
            CircuitState::Open => {}
        }
    }

    pub fn record_failure(&mut self, provider: &str) {
        let cb = self.get_or_create(provider);
        match cb.effective_state() {
            CircuitState::HalfOpen => {
                cb.state = CircuitState::Open;
                cb.last_failure_at = Some(Instant::now());
                let doubled = cb.cooldown * 2;
                cb.cooldown = doubled.min(cb.max_cooldown);
            }
            CircuitState::Closed => {
                let now = Instant::now();
                // Rolling-window accumulation: reset counter only when the
                // *entire window* has elapsed since the first failure in the
                // current accumulation period — NOT on the gap between two
                // consecutive failures.  This ensures failures spaced ~window
                // apart still accumulate toward the threshold.
                if let Some(start) = cb.window_start {
                    if now.duration_since(start) > cb.window {
                        cb.failure_count = 0;
                        cb.window_start = Some(now);
                    }
                } else {
                    cb.window_start = Some(now);
                }
                cb.failure_count += 1;
                cb.last_failure_at = Some(now);
                if cb.failure_count >= cb.threshold {
                    cb.state = CircuitState::Open;
                }
            }
            CircuitState::Open => {}
        }
    }

    pub fn record_credit_error(&mut self, provider: &str) {
        let cb = self.get_or_create(provider);
        cb.state = CircuitState::Open;
        cb.last_failure_at = Some(Instant::now());
        cb.credit_tripped = true;
        cb.operator_forced_open = false;
    }

    pub fn reset(&mut self, provider: &str) {
        let base_cooldown = self.config.cooldown_seconds;
        let cb = self.get_or_create(provider);
        cb.state = CircuitState::Closed;
        cb.failure_count = 0;
        cb.last_failure_at = None;
        cb.window_start = None;
        cb.cooldown = Duration::from_secs(base_cooldown);
        cb.credit_tripped = false;
        cb.operator_forced_open = false;
        cb.preemptive_half_open = false;
    }

    /// Returns true if the provider's breaker was tripped by a credit/billing error.
    pub fn is_credit_tripped(&self, provider: &str) -> bool {
        self.breakers
            .get(provider)
            .is_some_and(|cb| cb.credit_tripped)
    }

    /// Returns true if the provider is held open by operator kill-switch.
    pub fn is_operator_forced_open(&self, provider: &str) -> bool {
        self.breakers
            .get(provider)
            .is_some_and(|cb| cb.operator_forced_open)
    }

    pub fn get_state(&self, provider: &str) -> CircuitState {
        match self.breakers.get(provider) {
            Some(cb) => cb.effective_state(),
            None => CircuitState::Closed,
        }
    }

    pub fn list_providers(&self) -> Vec<(String, CircuitState)> {
        self.breakers
            .iter()
            .map(|(name, cb)| (name.clone(), cb.effective_state()))
            .collect()
    }

    /// Force a provider's circuit breaker open (kill-switch). Stays open
    /// until explicitly `reset()`, without being marked as a credit trip.
    pub fn force_open(&mut self, provider: &str) {
        let cb = self.get_or_create(provider);
        cb.state = CircuitState::Open;
        cb.last_failure_at = Some(Instant::now());
        cb.credit_tripped = false;
        cb.operator_forced_open = true;
    }

    /// Toggle soft capacity pressure state for a provider.
    pub fn set_capacity_pressure(&mut self, provider: &str, pressured: bool) {
        let cb = self.get_or_create(provider);
        if cb.credit_tripped || cb.operator_forced_open || cb.state == CircuitState::Open {
            return;
        }
        cb.preemptive_half_open = pressured;
    }

    /// Update the config used for newly-created breakers. Existing breakers
    /// retain their current thresholds — only breakers created after this
    /// call will use the updated values.
    pub fn sync_config(&mut self, config: &CircuitBreakerConfig) {
        self.config = config.clone();
    }
}

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

    fn test_config() -> CircuitBreakerConfig {
        CircuitBreakerConfig {
            threshold: 3,
            window_seconds: 60,
            cooldown_seconds: 2,
            max_cooldown_seconds: 30,
        }
    }

    #[test]
    fn normal_operation() {
        let mut reg = CircuitBreakerRegistry::new(&test_config());
        assert!(!reg.is_blocked("openai"));
        assert_eq!(reg.get_state("openai"), CircuitState::Closed);

        reg.record_success("openai");
        assert_eq!(reg.get_state("openai"), CircuitState::Closed);
        assert!(!reg.is_blocked("openai"));
    }

    #[test]
    fn trip_on_threshold() {
        let mut reg = CircuitBreakerRegistry::new(&test_config());

        reg.record_failure("openai");
        assert_eq!(reg.get_state("openai"), CircuitState::Closed);
        reg.record_failure("openai");
        assert_eq!(reg.get_state("openai"), CircuitState::Closed);
        reg.record_failure("openai");
        assert_eq!(reg.get_state("openai"), CircuitState::Open);
        assert!(reg.is_blocked("openai"));
    }

    #[test]
    fn recovery_after_cooldown() {
        let config = CircuitBreakerConfig {
            threshold: 1,
            cooldown_seconds: 0,
            ..test_config()
        };
        let mut reg = CircuitBreakerRegistry::new(&config);

        reg.record_failure("openai");
        // 0s cooldown means effective_state transitions to HalfOpen immediately
        std::thread::sleep(Duration::from_millis(5));
        assert_eq!(reg.get_state("openai"), CircuitState::HalfOpen);

        reg.record_success("openai");
        assert_eq!(reg.get_state("openai"), CircuitState::Closed);
        assert!(!reg.is_blocked("openai"));
    }

    #[test]
    fn credit_error_immediate_trip() {
        let mut reg = CircuitBreakerRegistry::new(&test_config());
        assert_eq!(reg.get_state("anthropic"), CircuitState::Closed);

        reg.record_credit_error("anthropic");
        assert_eq!(reg.get_state("anthropic"), CircuitState::Open);
        assert!(reg.is_blocked("anthropic"));
    }

    #[test]
    fn reset_clears_state() {
        let mut reg = CircuitBreakerRegistry::new(&test_config());
        reg.record_credit_error("openai");
        assert!(reg.is_blocked("openai"));

        reg.reset("openai");
        assert!(!reg.is_blocked("openai"));
        assert_eq!(reg.get_state("openai"), CircuitState::Closed);
    }

    #[test]
    fn half_open_failure_doubles_cooldown() {
        let config = CircuitBreakerConfig {
            threshold: 1,
            cooldown_seconds: 1,
            max_cooldown_seconds: 8,
            ..test_config()
        };
        let mut reg = CircuitBreakerRegistry::new(&config);

        reg.record_failure("openai");
        reg.force_half_open("openai");
        assert_eq!(reg.get_state("openai"), CircuitState::HalfOpen);

        reg.record_failure("openai");
        // Should be Open again with doubled cooldown (1s -> 2s)
        assert_eq!(reg.get_state("openai"), CircuitState::Open);
    }

    #[test]
    fn credit_error_never_auto_recovers() {
        let config = CircuitBreakerConfig {
            threshold: 1,
            cooldown_seconds: 0,
            ..test_config()
        };
        let mut reg = CircuitBreakerRegistry::new(&config);

        reg.record_credit_error("anthropic");
        assert_eq!(reg.get_state("anthropic"), CircuitState::Open);
        assert!(reg.is_credit_tripped("anthropic"));

        // Even with 0s cooldown, credit-tripped breakers stay Open (no HalfOpen)
        std::thread::sleep(Duration::from_millis(5));
        assert_eq!(reg.get_state("anthropic"), CircuitState::Open);
        assert!(reg.is_blocked("anthropic"));
    }

    #[test]
    fn credit_error_clears_on_manual_reset() {
        let mut reg = CircuitBreakerRegistry::new(&test_config());
        reg.record_credit_error("anthropic");
        assert!(reg.is_credit_tripped("anthropic"));
        assert!(reg.is_blocked("anthropic"));

        reg.reset("anthropic");
        assert!(!reg.is_credit_tripped("anthropic"));
        assert!(!reg.is_blocked("anthropic"));
        assert_eq!(reg.get_state("anthropic"), CircuitState::Closed);
    }

    #[test]
    fn transient_failure_still_auto_recovers() {
        let config = CircuitBreakerConfig {
            threshold: 1,
            cooldown_seconds: 0,
            ..test_config()
        };
        let mut reg = CircuitBreakerRegistry::new(&config);

        reg.record_failure("openai");
        assert!(!reg.is_credit_tripped("openai"));

        // With 0s cooldown, transient failures auto-recover to HalfOpen immediately
        std::thread::sleep(Duration::from_millis(5));
        assert_eq!(reg.get_state("openai"), CircuitState::HalfOpen);
        assert!(!reg.is_blocked("openai"));
    }

    #[test]
    fn capacity_pressure_sets_half_open_without_blocking() {
        let mut reg = CircuitBreakerRegistry::new(&test_config());
        reg.set_capacity_pressure("openai", true);
        assert_eq!(reg.get_state("openai"), CircuitState::HalfOpen);
        assert!(!reg.is_blocked("openai"));

        reg.set_capacity_pressure("openai", false);
        assert_eq!(reg.get_state("openai"), CircuitState::Closed);
        assert!(!reg.is_blocked("openai"));
    }

    #[test]
    fn force_open_blocks_and_stays_open() {
        let mut reg = CircuitBreakerRegistry::new(&test_config());
        assert!(!reg.is_blocked("openai"));

        reg.force_open("openai");
        assert!(reg.is_blocked("openai"));
        assert_eq!(reg.get_state("openai"), CircuitState::Open);
        assert!(reg.is_operator_forced_open("openai"));
        assert!(!reg.is_credit_tripped("openai"));

        // Should NOT auto-recover to HalfOpen (operator-forced open prevents it)
        std::thread::sleep(Duration::from_millis(5));
        assert!(reg.is_blocked("openai"));

        // reset() clears the force-open
        reg.reset("openai");
        assert!(!reg.is_blocked("openai"));
        assert!(!reg.is_operator_forced_open("openai"));
        assert_eq!(reg.get_state("openai"), CircuitState::Closed);
    }
}