use embassy_time::{Duration, Timer};
use log::warn;

#[derive(Clone, Copy, Debug)]
pub struct Backoff {
    delay: Duration,
    attempts: Attempts,
}

#[derive(Clone, Copy, Debug)]
pub enum Attempts {
    Finite(u64),
    Forever
}

impl Iterator for Attempts {
    type Item = Self;

    fn next(&mut self) -> Option<Self::Item> {
        match *self {
            Attempts::Forever => Some(Attempts::Forever),
            Attempts::Finite(0) => None,
            Attempts::Finite(n) => {
                *self = Attempts::Finite(n - 1);
                Some(Attempts::Finite(n))
            }
        }
    }
}

impl Backoff {
    pub fn from_millis(millis: u64) -> Self {
        Self {
            delay: Duration::from_millis(millis),
            attempts: Attempts::Finite(3),
        }
    }

    pub fn from_secs(secs: u64) -> Self {
        Self {
            delay: Duration::from_secs(secs),
            attempts: Attempts::Finite(3),
        }
    }

    pub fn forever(self) -> Self {
        Self {
            attempts: Attempts::Forever,
            ..self
        }
    }

    pub async fn attempt<F: AsyncFnMut() -> Result<T, E>, T, E: core::fmt::Debug>(self, mut f: F) -> Result<T, E> {
        let mut latest= f().await;
        let mut delay = self.delay;
        'outer: loop {
            for attempt in self.attempts {
                match &latest {
                    Err(err) => {
                        match attempt {
                            Attempts::Finite(1) => {
                                warn!("Operation failed on final attempt: {err:?}");
                                break 'outer
                            }
                            attempt => {
                                warn!("Operation failed: {err:?} Retrying attempt {attempt:?} after {delay}");
                                Timer::after(delay).await;
                                delay *= 2;
                                latest = f().await
                            }
                        }
                    },
                    _ => break 'outer
                }
            }
        }

        latest
    }
}