// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! PDK Spike Control Library
//!
//! Provides spike control functionality for Flex Gateway policies: rate limiting with
//! optional retry-and-delay logic to absorb temporary traffic spikes.
//!
//! When a request exceeds its quota, instead of rejecting it immediately, the library
//! can suspend the request for a configurable delay and re-check the quota, retrying
//! up to a maximum number of attempts before finally rejecting with a 429 response.
//!
//! **Without retry** (queuing disabled): requests exceeding the quota are rejected
//! immediately, with no async delay.
//!
//! **With retry** (queuing enabled): requests are delayed and retried on each
//! quota check using a non-blocking async timer, allowing spikes in traffic to
//! succeed if quota becomes available within the configured attempts window.
//!
//! ## Highlights
//!
//! - Per-call retry control via the `with_retry` flag on [`SpikeControlHandler::is_allowed`]
//! - Non-blocking async delays
//! - Optional clustered mode for distributed rate limiting across gateway nodes
//! - Structured error type with quota statistics attached to rejections
//!
//! ## Primary types
//!
//! - [`SpikeControlBuilder`]: PDK-injected builder; entry point for constructing a handler
//! - [`SpikeControlHandler`]: owns the rate-limit instance and executes the retry loop
//! - [`SpikeControlResult`]: type alias for the result of a [`SpikeControlHandler::is_allowed`] call
//! - [`SpikeControlError`]: error types for spike control operations
//! - [`RetryConfig`]: value type describing delay duration, max attempts and enabled state

use pdk_core::classy::extract::context::ConfigureContext;
use pdk_core::classy::extract::{Extract, FromContext};
use pdk_core::classy::timer::Timer;
use pdk_core::logger;
use pdk_core::policy_context::api::Tier;
use rate_limit_lib::{
    RateLimit, RateLimitBuilder, RateLimitBuilderInstance, RateLimitError, RateLimitInstance,
    RateLimitResult, RateLimitStatistics,
};
use std::convert::Infallible;
use std::rc::Rc;
use std::time::Duration;
use thiserror::Error;

/// Configuration for retry behavior when requests exceed rate limits.
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Time to wait between retry attempts in milliseconds.
    pub delay_time_in_millis: u64,

    /// Maximum number of retry attempts before rejecting the request.
    /// If 1, the request is rejected immediately without retrying.
    pub max_attempts: u32,

    /// Whether retry logic is enabled.
    pub enabled: bool,
}

impl RetryConfig {
    /// Create a new retry configuration.
    pub fn new(delay_time_in_millis: u64, max_attempts: u32, enabled: bool) -> Self {
        Self {
            delay_time_in_millis,
            max_attempts,
            enabled,
        }
    }

    /// Create a configuration with retries disabled (immediate rejection).
    pub fn disabled() -> Self {
        Self {
            delay_time_in_millis: 0,
            max_attempts: 1,
            enabled: false,
        }
    }
}

/// Result type for spike control operations.
pub type SpikeControlResult = Result<RateLimitStatistics, SpikeControlError>;

/// Errors that can occur during spike control checks.
#[non_exhaustive]
#[derive(Error)]
pub enum SpikeControlError {
    /// Request exceeded rate limit and all retry attempts.
    #[error("Too many requests")]
    TooManyRequests(RateLimitStatistics),

    /// Rate limit check failed due to an internal error.
    #[error("Rate limit error: {0}")]
    RateLimitError(#[from] RateLimitError),

    /// Failed to build the spike control handler.
    #[error("Build error: {0}")]
    BuildError(String),
}

impl std::fmt::Debug for SpikeControlError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SpikeControlError::TooManyRequests(stats) => f
                .debug_struct("TooManyRequests")
                .field("remaining", &stats.remaining)
                .field("limit", &stats.limit)
                .field("reset", &stats.reset)
                .finish(),
            SpikeControlError::RateLimitError(e) => {
                f.debug_tuple("RateLimitError").field(e).finish()
            }
            SpikeControlError::BuildError(msg) => f.debug_tuple("BuildError").field(msg).finish(),
        }
    }
}

/// Builder for creating spike control instances.
pub struct SpikeControlBuilder {
    rate_limit_builder: RateLimitBuilder,
}

impl FromContext<ConfigureContext> for SpikeControlBuilder {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        let rate_limit_builder: RateLimitBuilder = context.extract()?;
        Ok(Self { rate_limit_builder })
    }
}

impl SpikeControlBuilder {
    /// Creates a new [`SpikeControlBuilderInstance`] that can be configured
    /// with buckets, retry settings and clustering options before building.
    #[allow(clippy::new_ret_no_self)]
    pub fn new(&self, name: String) -> SpikeControlBuilderInstance {
        SpikeControlBuilderInstance {
            rl_builder: self.rate_limit_builder.new(name),
            retry_config: RetryConfig::disabled(),
            retry_timer: None,
            clustered: false,
        }
    }
}

/// Instance builder for configuring a specific spike control handler.
///
/// This struct is returned by [`SpikeControlBuilder::new`] and provides
/// methods to configure rate limiting buckets, retry behavior and clustering.
pub struct SpikeControlBuilderInstance {
    rl_builder: RateLimitBuilderInstance,
    retry_config: RetryConfig,
    retry_timer: Option<Rc<Timer>>,
    clustered: bool,
}

impl SpikeControlBuilderInstance {
    /// Adds a rate limiting bucket with configured tiers.
    pub fn with_bucket(mut self, name: String, tiers: Vec<Tier>) -> Self {
        self.rl_builder = self.rl_builder.buckets(vec![(name, tiers)]);
        self
    }

    /// Configures retry behavior for rate-limited requests.
    pub fn with_retry(mut self, delay_time_in_millis: u64, max_attempts: u32) -> Self {
        self.retry_config = RetryConfig::new(delay_time_in_millis, max_attempts, true);
        self
    }

    /// Sets the ticker used for retry delays.
    pub fn with_ticker(mut self, ticker: Rc<Timer>) -> Self {
        self.retry_timer = Some(ticker);
        self
    }

    /// Enables clustered mode for distributed rate limiting across multiple nodes.
    pub fn clustered(mut self) -> Self {
        self.clustered = true;
        self
    }

    /// Builds the spike control handler with the configured options.
    pub fn build(self) -> Result<SpikeControlHandler, SpikeControlError> {
        if (self.retry_config.enabled || self.clustered) && self.retry_timer.is_none() {
            return Err(SpikeControlError::BuildError(
                "Ticker is required when retry or clustered mode is enabled".to_string(),
            ));
        }

        let mut rl_builder = self.rl_builder;
        if self.clustered {
            if let Some(ticker) = self.retry_timer.as_ref() {
                rl_builder = rl_builder.clustered(Rc::clone(ticker));
            }
        }

        let rate_limit = rl_builder
            .build()
            .map_err(|e| SpikeControlError::BuildError(e.to_string()))?;

        Ok(SpikeControlHandler {
            rate_limit,
            retry_config: self.retry_config,
            retry_timer: self.retry_timer,
        })
    }
}

/// Handler for spike control with retry logic.
pub struct SpikeControlHandler {
    rate_limit: RateLimitInstance,
    retry_config: RetryConfig,
    retry_timer: Option<Rc<Timer>>,
}

impl SpikeControlHandler {
    /// Check if a request is allowed, with optional retry logic.
    ///
    /// If the request exceeds the rate limit and `with_retry` is true, this will
    /// delay and retry up to `max_attempts` times before rejecting.
    /// When the request is allowed quota statistics are returned.
    /// If the request is rejected the appropriate error is returned.
    pub async fn is_allowed(
        &self,
        bucket: &str,
        key: &str,
        weight: usize,
        with_retry: bool,
    ) -> SpikeControlResult {
        let max_attempts = if with_retry && self.retry_config.enabled {
            self.retry_config.max_attempts
        } else {
            1
        };
        let retry_delay = Duration::from_millis(self.retry_config.delay_time_in_millis);

        for attempt in 0..max_attempts {
            match self.rate_limit.is_allowed(bucket, key, weight).await {
                Ok(RateLimitResult::Allowed(stats)) => {
                    if attempt > 0 {
                        logger::debug!(
                            "Request allowed on attempt {}/{}, quota remaining: {}",
                            attempt + 1,
                            max_attempts,
                            stats.remaining
                        );
                    }
                    return Ok(stats);
                }
                Ok(RateLimitResult::TooManyRequests(stats)) => {
                    // Check if we should retry or reject
                    if attempt + 1 >= max_attempts {
                        // Max attempts reached or no retries, reject
                        logger::debug!(
                            "Request rejected after {} attempt(s), quota: {}/{}",
                            attempt + 1,
                            stats.remaining,
                            stats.limit
                        );
                        return Err(SpikeControlError::TooManyRequests(stats));
                    }

                    // Delay and retry
                    logger::debug!(
                        "Request delayed, attempt {}/{}, waiting {}ms",
                        attempt + 1,
                        max_attempts,
                        self.retry_config.delay_time_in_millis
                    );
                    let Some(timer) = self.retry_timer.as_ref() else {
                        logger::error!(
                            "Retry is enabled but no retry ticker is configured at runtime"
                        );
                        return Err(SpikeControlError::BuildError(
                            "Retry is enabled but no retry ticker is configured".to_string(),
                        ));
                    };

                    if !timer.sleep(retry_delay).await {
                        logger::warn!("Retry sleep interrupted by host shutdown");
                        return Err(SpikeControlError::BuildError(
                            "Retry sleep interrupted by host shutdown".to_string(),
                        ));
                    }
                }
                Err(e) => {
                    logger::debug!("Rate limit error: {e}");
                    return Err(SpikeControlError::RateLimitError(e));
                }
            }
        }

        // Should never reach here, but handle defensively
        logger::error!("Unexpected: exceeded retry loop without decision");
        Err(SpikeControlError::TooManyRequests(RateLimitStatistics {
            remaining: 0,
            limit: 0,
            reset: 0,
        }))
    }

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

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

    #[test]
    fn retry_config_disabled_creates_immediate_rejection() {
        let config = RetryConfig::disabled();
        assert_eq!(config.max_attempts, 1);
        assert!(!config.enabled);
        assert_eq!(config.delay_time_in_millis, 0);
    }

    #[test]
    fn retry_config_new_accepts_custom_values() {
        let config = RetryConfig::new(500, 3, true);
        assert_eq!(config.delay_time_in_millis, 500);
        assert_eq!(config.max_attempts, 3);
        assert!(config.enabled);
    }

    #[test]
    fn retry_config_with_zero_delay() {
        let config = RetryConfig::new(0, 5, true);
        assert_eq!(config.delay_time_in_millis, 0);
        assert_eq!(config.max_attempts, 5);
        assert!(config.enabled);
    }

    #[test]
    fn retry_config_enabled_false_overrides_attempts() {
        let config = RetryConfig::new(1000, 10, false);
        assert!(!config.enabled);
        // Even with 10 attempts configured, disabled means it won't retry
    }

    #[test]
    fn spike_control_handler_stores_config() {
        let config = RetryConfig::new(1000, 5, true);
        let retrieved_config = config.clone();
        assert_eq!(retrieved_config.delay_time_in_millis, 1000);
        assert_eq!(retrieved_config.max_attempts, 5);
    }

    #[test]
    fn spike_control_error_debug_format_too_many_requests() {
        let stats = RateLimitStatistics {
            remaining: 0,
            limit: 100,
            reset: 60000,
        };
        let error = SpikeControlError::TooManyRequests(stats);
        let debug_str = format!("{error:?}");
        assert!(debug_str.contains("TooManyRequests"));
        assert!(debug_str.contains("remaining"));
        assert!(debug_str.contains("limit"));
        assert!(debug_str.contains("reset"));
    }

    #[test]
    fn spike_control_error_debug_format_build_error() {
        let error = SpikeControlError::BuildError("test error".to_string());
        let debug_str = format!("{error:?}");
        assert!(debug_str.contains("BuildError"));
        assert!(debug_str.contains("test error"));
    }

    #[test]
    fn spike_control_error_display_format() {
        let stats = RateLimitStatistics {
            remaining: 0,
            limit: 100,
            reset: 60000,
        };
        let error = SpikeControlError::TooManyRequests(stats);
        assert_eq!(error.to_string(), "Too many requests");

        let build_error = SpikeControlError::BuildError("build failed".to_string());
        assert_eq!(build_error.to_string(), "Build error: build failed");
    }

    #[test]
    fn retry_config_clone_works() {
        let config1 = RetryConfig::new(2000, 7, true);
        let config2 = config1.clone();
        assert_eq!(config1.delay_time_in_millis, config2.delay_time_in_millis);
        assert_eq!(config1.max_attempts, config2.max_attempts);
        assert_eq!(config1.enabled, config2.enabled);
    }

    #[test]
    fn retry_config_disabled_does_not_require_timer() {
        let config = RetryConfig::disabled();
        assert!(!config.enabled);
    }

    #[test]
    fn retry_config_enabled_requires_timer() {
        let config = RetryConfig::new(250, 3, true);
        assert!(config.enabled);
    }

    #[test]
    fn retry_config_enabled_with_values_is_valid() {
        let config = RetryConfig::new(250, 3, true);
        assert_eq!(config.delay_time_in_millis, 250);
        assert_eq!(config.max_attempts, 3);
    }
}
