//! Generic implementation of CircuitBreakerStore backed by any KVStore.

use super::{
    CircuitBreakerConfig, CircuitBreakerRecord, CircuitBreakerStore, CircuitState, FailureType,
};
use crate::key_codec;
use crate::{KVStore, Result, StorageError};
use async_trait::async_trait;
use std::sync::Arc;
use tracing::{debug, instrument, warn};

/// Generic implementation of `CircuitBreakerStore` backed by any `KVStore`.
pub struct GenericCircuitBreakerStore<S> {
    store: Arc<S>,
    config: CircuitBreakerConfig,
}

// Manual Clone implementation because Arc<S> is Clone even if S is not
impl<S> Clone for GenericCircuitBreakerStore<S> {
    fn clone(&self) -> Self {
        Self { store: Arc::clone(&self.store), config: self.config }
    }
}

impl<S> GenericCircuitBreakerStore<S> {
    /// Create a new circuit breaker store with the given config.
    ///
    /// Takes an `Arc<S>` to enable sharing the store across components.
    pub fn new(store: Arc<S>, config: CircuitBreakerConfig) -> Self {
        Self { store, config }
    }

    /// Create a new circuit breaker store with default config.
    pub fn with_defaults(store: Arc<S>) -> Self {
        Self::new(store, CircuitBreakerConfig::default())
    }

    /// Get the configuration.
    pub fn config(&self) -> &CircuitBreakerConfig {
        &self.config
    }
}

impl<S: KVStore> GenericCircuitBreakerStore<S> {
    /// Load a circuit breaker record from storage.
    async fn load_record(&self, agent_id: &[u8; 32]) -> Result<Option<CircuitBreakerRecord>> {
        let agent_hex = hex::encode(agent_id);
        let key = key_codec::circuit_breaker_key(&agent_hex);

        match self.store.get(&key).await? {
            Some(data) => {
                let record: CircuitBreakerRecord = stemedb_core::serde::deserialize(&data)
                    .map_err(|e| StorageError::Serialization(e.to_string()))?;
                Ok(Some(record))
            }
            None => Ok(None),
        }
    }

    /// Save a circuit breaker record to storage.
    async fn save_record(&self, record: &CircuitBreakerRecord) -> Result<()> {
        let agent_hex = hex::encode(record.agent_id);
        let key = key_codec::circuit_breaker_key(&agent_hex);

        let data = stemedb_core::serde::serialize(record)
            .map_err(|e| StorageError::Serialization(e.to_string()))?;

        self.store.put(&key, &data).await
    }

    /// Delete a circuit breaker record from storage.
    async fn delete_record(&self, agent_id: &[u8; 32]) -> Result<()> {
        let agent_hex = hex::encode(agent_id);
        let key = key_codec::circuit_breaker_key(&agent_hex);
        self.store.delete(&key).await
    }
}

#[async_trait]
impl<S: KVStore + 'static> CircuitBreakerStore for GenericCircuitBreakerStore<S> {
    #[instrument(skip(self), fields(agent = %hex::encode(agent_id)))]
    async fn get_circuit(&self, agent_id: &[u8; 32]) -> Result<Option<CircuitBreakerRecord>> {
        self.load_record(agent_id).await
    }

    #[instrument(skip(self), fields(agent = %hex::encode(agent_id), failure_type = %failure_type.name()))]
    async fn record_failure(
        &self,
        agent_id: &[u8; 32],
        failure_type: FailureType,
        timestamp: u64,
    ) -> Result<CircuitBreakerRecord> {
        // Load or create record
        let mut record = self.load_record(agent_id).await?.unwrap_or_else(|| {
            debug!(agent = %hex::encode(agent_id), "Creating new circuit breaker record");
            CircuitBreakerRecord::new(*agent_id)
        });

        // Handle based on current state
        match record.state {
            CircuitState::Closed => {
                // Add failure and check threshold
                let failure_count =
                    record.add_failure(failure_type, timestamp, self.config.failure_window_secs);

                debug!(
                    agent = %hex::encode(agent_id),
                    failure_count = failure_count,
                    threshold = self.config.failure_threshold,
                    "Recorded failure in Closed state"
                );

                if failure_count >= self.config.failure_threshold as usize {
                    record.trip(timestamp);
                    warn!(
                        agent = %hex::encode(agent_id),
                        trip_count = record.trip_count,
                        "Circuit breaker tripped"
                    );
                }
            }
            CircuitState::HalfOpen => {
                // Failure in HalfOpen → trip back to Open
                record.add_failure(failure_type, timestamp, self.config.failure_window_secs);
                record.trip(timestamp);
                warn!(
                    agent = %hex::encode(agent_id),
                    trip_count = record.trip_count,
                    "Circuit breaker re-tripped from HalfOpen"
                );
            }
            CircuitState::Open => {
                // Already open, just record the failure
                record.add_failure(failure_type, timestamp, self.config.failure_window_secs);
                debug!(
                    agent = %hex::encode(agent_id),
                    "Recorded failure while circuit is Open"
                );
            }
        }

        self.save_record(&record).await?;
        Ok(record)
    }

    #[instrument(skip(self), fields(agent = %hex::encode(agent_id)))]
    async fn record_success(&self, agent_id: &[u8; 32], timestamp: u64) -> Result<()> {
        let record = match self.load_record(agent_id).await? {
            Some(r) => r,
            None => {
                // No record means agent is in good standing, nothing to do
                debug!(agent = %hex::encode(agent_id), "No circuit breaker record, ignoring success");
                return Ok(());
            }
        };

        match record.state {
            CircuitState::HalfOpen => {
                let mut record = record;
                let should_close =
                    record.record_half_open_success(self.config.half_open_success_threshold);

                if should_close {
                    record.close();
                    debug!(
                        agent = %hex::encode(agent_id),
                        "Circuit closed after successful HalfOpen test"
                    );
                    // Delete the record to clean up storage
                    self.delete_record(agent_id).await?;
                } else {
                    debug!(
                        agent = %hex::encode(agent_id),
                        successes = record.half_open_successes,
                        threshold = self.config.half_open_success_threshold,
                        "HalfOpen success recorded"
                    );
                    self.save_record(&record).await?;
                }
            }
            CircuitState::Closed => {
                // Success in Closed state is normal, no action needed
                debug!(agent = %hex::encode(agent_id), "Success in Closed state, ignoring");

                // Prune old failures if enough time has passed
                let mut record = record;
                let cutoff = timestamp.saturating_sub(self.config.failure_window_secs);
                record.failures.retain(|f| f.timestamp >= cutoff);

                if record.failures.is_empty() && record.trip_count == 0 {
                    // Clean up record if no failures and never tripped
                    self.delete_record(agent_id).await?;
                } else {
                    self.save_record(&record).await?;
                }
            }
            CircuitState::Open => {
                // Success shouldn't happen in Open state (requests should be blocked)
                // But if it does, treat it as if we're in HalfOpen
                warn!(
                    agent = %hex::encode(agent_id),
                    "Unexpected success in Open state"
                );
            }
        }

        Ok(())
    }

    #[instrument(skip(self), fields(agent = %hex::encode(agent_id)))]
    async fn reset_circuit(&self, agent_id: &[u8; 32]) -> Result<()> {
        debug!(agent = %hex::encode(agent_id), "Resetting circuit breaker");
        self.delete_record(agent_id).await
    }

    #[instrument(skip(self))]
    async fn list_tripped(&self, limit: usize) -> Result<Vec<CircuitBreakerRecord>> {
        let entries = self.store.scan_prefix(&key_codec::circuit_breaker_scan_prefix()).await?;

        let mut tripped = Vec::new();
        for (_key, data) in entries {
            if tripped.len() >= limit {
                break;
            }

            match stemedb_core::serde::deserialize::<CircuitBreakerRecord>(&data) {
                Ok(record) if record.state != CircuitState::Closed => {
                    tripped.push(record);
                }
                Ok(_) => {} // Skip Closed circuits
                Err(e) => {
                    debug!(error = %e, "Skipping malformed circuit breaker record");
                }
            }
        }

        // Sort by last failure time (most recent first)
        tripped.sort_by(|a, b| b.last_failure_time.cmp(&a.last_failure_time));

        debug!(count = tripped.len(), limit = limit, "Listed tripped circuit breakers");

        Ok(tripped)
    }

    #[instrument(skip(self), fields(agent = %hex::encode(agent_id)))]
    async fn check_allowed(&self, agent_id: &[u8; 32], current_time: u64) -> Result<bool> {
        let record = match self.load_record(agent_id).await? {
            Some(r) => r,
            None => {
                // No record means agent is allowed
                return Ok(true);
            }
        };

        match record.state {
            CircuitState::Closed => Ok(true),
            CircuitState::HalfOpen => {
                // Allow one request for testing
                debug!(agent = %hex::encode(agent_id), "Allowing HalfOpen test request");
                Ok(true)
            }
            CircuitState::Open => {
                // Check if timeout has elapsed
                if record.open_timeout_elapsed(current_time, self.config.open_duration_secs) {
                    // Transition to HalfOpen
                    let mut record = record;
                    record.half_open();
                    self.save_record(&record).await?;
                    debug!(agent = %hex::encode(agent_id), "Circuit transitioned to HalfOpen");
                    Ok(true)
                } else {
                    let retry_after =
                        record.seconds_until_retry(current_time, self.config.open_duration_secs);
                    debug!(
                        agent = %hex::encode(agent_id),
                        retry_after = retry_after,
                        "Circuit is Open, request blocked"
                    );
                    Ok(false)
                }
            }
        }
    }

    #[instrument(skip(self), fields(agent = %hex::encode(agent_id)))]
    async fn retry_after(&self, agent_id: &[u8; 32], current_time: u64) -> Result<Option<u64>> {
        let record = match self.load_record(agent_id).await? {
            Some(r) => r,
            None => return Ok(None),
        };

        match record.state {
            CircuitState::Open => {
                let secs = record.seconds_until_retry(current_time, self.config.open_duration_secs);
                Ok(Some(secs))
            }
            _ => Ok(None),
        }
    }
}

#[cfg(test)]
#[path = "tests.rs"]
mod tests;
