use std::time::Duration;

/// Retries `operation` with exponentially increasing delays.
///
/// `max_attempts` includes the initial attempt and must be greater than zero.
pub fn retry_with_backoff<T, E, F>(
    operation: F,
    max_attempts: usize,
    initial_delay: Duration,
    max_delay: Duration,
) -> Result<T, E>
where
    F: FnMut() -> Result<T, E>,
{
    retry_with_backoff_using(
        operation,
        max_attempts,
        initial_delay,
        max_delay,
        std::thread::sleep,
    )
}

fn retry_with_backoff_using<T, E, F, S>(
    mut operation: F,
    max_attempts: usize,
    initial_delay: Duration,
    max_delay: Duration,
    mut sleep: S,
) -> Result<T, E>
where
    F: FnMut() -> Result<T, E>,
    S: FnMut(Duration),
{
    assert!(max_attempts > 0, "max_attempts must be greater than zero");

    let mut delay = initial_delay.min(max_delay);

    for attempt in 1..=max_attempts {
        match operation() {
            Ok(value) => return Ok(value),
            Err(error) if attempt == max_attempts => return Err(error),
            Err(_) => {
                sleep(delay);
                delay = delay.saturating_mul(2).min(max_delay);
            }
        }
    }

    unreachable!()
}

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

    #[test]
    fn succeeds_after_transient_failures() {
        let mut attempts = 0;
        let mut delays = Vec::new();

        let result = retry_with_backoff_using(
            || {
                attempts += 1;
                if attempts < 4 {
                    Err("temporary failure")
                } else {
                    Ok(42)
                }
            },
            5,
            Duration::from_millis(10),
            Duration::from_millis(25),
            |delay| delays.push(delay),
        );

        assert_eq!(result, Ok(42));
        assert_eq!(attempts, 4);
        assert_eq!(
            delays,
            [
                Duration::from_millis(10),
                Duration::from_millis(20),
                Duration::from_millis(25),
            ]
        );
    }

    #[test]
    fn returns_last_error_without_sleeping_after_final_attempt() {
        let mut attempts = 0;
        let mut sleeps = 0;

        let result: Result<(), usize> = retry_with_backoff_using(
            || {
                attempts += 1;
                Err(attempts)
            },
            3,
            Duration::from_millis(1),
            Duration::from_secs(1),
            |_| sleeps += 1,
        );

        assert_eq!(result, Err(3));
        assert_eq!(attempts, 3);
        assert_eq!(sleeps, 2);
    }
}