use std::time::Duration;

#[derive(Clone, Debug)]
pub(crate) struct RetryConfig {
    pub(crate) max_retries: u32,
    pub(crate) retry_on_429: bool,
    pub(crate) respect_retry_after: bool,
    pub(crate) base_backoff: Duration,
    pub(crate) max_backoff: Duration,
    pub(crate) jitter: Option<Duration>,
    pub(crate) total_retry_budget: Option<Duration>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum RetryDecision {
    DoNotRetry,
    RetryAfter(Duration),
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_retries: 3,
            retry_on_429: false,
            respect_retry_after: false,
            base_backoff: Duration::from_millis(50),
            max_backoff: Duration::from_millis(250),
            jitter: None,
            total_retry_budget: None,
        }
    }
}

impl RetryConfig {
    pub(crate) fn classify_status(
        &self,
        status: reqwest::StatusCode,
        attempt: u32,
        retry_after: Option<Duration>,
        elapsed: Duration,
    ) -> RetryDecision {
        if attempt >= self.max_retries {
            return RetryDecision::DoNotRetry;
        }

        let wait = if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
            if !self.retry_on_429 {
                return RetryDecision::DoNotRetry;
            }

            if self.respect_retry_after {
                retry_after.unwrap_or_else(|| self.backoff(attempt + 1))
            } else {
                self.backoff(attempt + 1)
            }
        } else if status.is_server_error() {
            self.backoff(attempt + 1)
        } else {
            return RetryDecision::DoNotRetry;
        };

        match self.finalize_wait(wait, elapsed, current_jitter_seed()) {
            Some(wait) => RetryDecision::RetryAfter(wait),
            None => RetryDecision::DoNotRetry,
        }
    }

    fn backoff(&self, attempt: u32) -> Duration {
        let factor = 1u32
            .checked_shl(attempt.saturating_sub(1))
            .unwrap_or(u32::MAX);
        let millis = self.base_backoff.as_millis();
        let scaled = millis.saturating_mul(u128::from(factor));
        let bounded = scaled.min(self.max_backoff.as_millis());
        let bounded_u64 = u64::try_from(bounded).unwrap_or(u64::MAX);
        Duration::from_millis(bounded_u64)
    }

    fn finalize_wait(
        &self,
        wait: Duration,
        elapsed: Duration,
        jitter_seed: u128,
    ) -> Option<Duration> {
        let wait = wait.min(self.max_backoff);

        let remaining_budget = match self.total_retry_budget {
            Some(total_retry_budget) => {
                let remaining_budget = total_retry_budget.checked_sub(elapsed)?;
                if remaining_budget.is_zero() {
                    return None;
                }

                Some(remaining_budget)
            }
            None => None,
        };

        let wait = if let Some(remaining_budget) = remaining_budget {
            if remaining_budget.is_zero() {
                return None;
            }

            self.apply_jitter_with_seed(wait.min(remaining_budget), jitter_seed)
                .min(remaining_budget)
        } else {
            self.apply_jitter_with_seed(wait, jitter_seed)
        };

        Some(wait)
    }

    fn apply_jitter_with_seed(&self, wait: Duration, jitter_seed: u128) -> Duration {
        let Some(max_jitter) = self.jitter else {
            return wait;
        };

        let jitter_nanos = max_jitter.as_nanos();
        if jitter_nanos == 0 {
            return wait;
        }

        let extra = (jitter_seed % jitter_nanos.saturating_add(1)).min(u128::from(u64::MAX));
        let extra = Duration::from_nanos(u64::try_from(extra).unwrap_or(u64::MAX));

        wait.saturating_add(extra).min(self.max_backoff)
    }
}

fn current_jitter_seed() -> u128 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos()
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::{RetryConfig, RetryDecision};

    #[test]
    fn default_config_keeps_429_retries_disabled() {
        let config = RetryConfig::default();

        assert_eq!(
            config.classify_status(
                reqwest::StatusCode::TOO_MANY_REQUESTS,
                0,
                Some(std::time::Duration::from_secs(1)),
                std::time::Duration::ZERO,
            ),
            RetryDecision::DoNotRetry
        );
    }

    #[test]
    fn retry_after_is_honored_when_enabled() {
        let config = RetryConfig {
            retry_on_429: true,
            respect_retry_after: true,
            max_backoff: Duration::from_secs(10),
            ..RetryConfig::default()
        };

        assert_eq!(
            config.classify_status(
                reqwest::StatusCode::TOO_MANY_REQUESTS,
                0,
                Some(Duration::from_secs(3)),
                Duration::ZERO,
            ),
            RetryDecision::RetryAfter(Duration::from_secs(3))
        );
    }

    #[test]
    fn retry_after_larger_than_remaining_budget_is_clamped() {
        let config = RetryConfig {
            retry_on_429: true,
            respect_retry_after: true,
            max_backoff: Duration::from_secs(10),
            total_retry_budget: Some(Duration::from_millis(50)),
            ..RetryConfig::default()
        };

        assert_eq!(
            config.classify_status(
                reqwest::StatusCode::TOO_MANY_REQUESTS,
                0,
                Some(Duration::from_millis(45)),
                Duration::from_millis(40),
            ),
            RetryDecision::RetryAfter(Duration::from_millis(10))
        );
    }

    #[test]
    fn retry_after_with_jitter_still_does_not_exceed_remaining_budget() {
        let config = RetryConfig {
            jitter: Some(Duration::from_millis(50)),
            total_retry_budget: Some(Duration::from_millis(50)),
            max_backoff: Duration::from_secs(1),
            ..RetryConfig::default()
        };

        assert_eq!(
            config.finalize_wait(
                Duration::from_millis(45),
                Duration::from_millis(40),
                Duration::from_millis(50).as_nanos(),
            ),
            Some(Duration::from_millis(10))
        );
    }

    #[test]
    fn budget_without_jitter_respects_remaining_budget() {
        let config = RetryConfig {
            total_retry_budget: Some(Duration::from_millis(100)),
            ..RetryConfig::default()
        };

        assert_eq!(
            config.finalize_wait(Duration::from_millis(50), Duration::from_millis(90), 0),
            Some(Duration::from_millis(10))
        );
    }

    #[test]
    fn budget_with_jitter_still_does_not_exceed_remaining_budget() {
        let config = RetryConfig {
            jitter: Some(Duration::from_millis(50)),
            total_retry_budget: Some(Duration::from_millis(100)),
            max_backoff: Duration::from_secs(1),
            ..RetryConfig::default()
        };

        assert_eq!(
            config.finalize_wait(
                Duration::from_millis(10),
                Duration::from_millis(90),
                Duration::from_millis(50).as_nanos(),
            ),
            Some(Duration::from_millis(10))
        );
    }

    #[test]
    fn exhausted_budget_on_429_returns_do_not_retry() {
        let config = RetryConfig {
            retry_on_429: true,
            respect_retry_after: true,
            total_retry_budget: Some(Duration::from_millis(50)),
            ..RetryConfig::default()
        };

        assert_eq!(
            config.classify_status(
                reqwest::StatusCode::TOO_MANY_REQUESTS,
                0,
                Some(Duration::from_millis(1)),
                Duration::from_millis(50),
            ),
            RetryDecision::DoNotRetry
        );
    }

    #[test]
    fn exhausted_budget_returns_do_not_retry() {
        let config = RetryConfig {
            total_retry_budget: Some(Duration::from_millis(100)),
            ..RetryConfig::default()
        };

        assert_eq!(
            config.classify_status(
                reqwest::StatusCode::INTERNAL_SERVER_ERROR,
                0,
                None,
                Duration::from_millis(100),
            ),
            RetryDecision::DoNotRetry
        );
    }
}
