use std::fmt;
use std::time::Duration;

pub fn retry<T, TOk, TError>(f: T) -> Result<TOk, RetryError<TError>>
where
    T: Fn() -> Result<TOk, Retryable<TError>>,
{
    builder().build().run(f)
}

#[cfg(feature = "tokio")]
pub async fn retry_async<T, TFut, TOk, TError>(f: T) -> Result<TOk, RetryError<TError>>
where
    T: Fn() -> TFut + Send + 'static,
    TFut: Future<Output = Result<TOk, Retryable<TError>>>,
{
    builder().build().run_async(f).await
}

pub fn builder() -> RetryBuilder {
    RetryBuilder {
        max_attempts: 3,
        retry_strategy: RetryStrategy::ExponentialBackoff,
        initial_delay: Duration::from_millis(100),
        max_delay: Duration::from_secs(30),
        jitter: true,
    }
}

pub struct RetryBuilder {
    max_attempts: usize,
    retry_strategy: RetryStrategy,
    initial_delay: Duration,
    max_delay: Duration,
    jitter: bool,
}

pub enum RetryStrategy {
    LinearBackoff,
    ExponentialBackoff,
}

impl RetryBuilder {
    pub fn with_max_attempts(mut self, attempts: usize) -> Self {
        self.max_attempts = attempts;
        self
    }

    pub fn with_exponential_backoff(mut self) -> Self {
        self.retry_strategy = RetryStrategy::ExponentialBackoff;
        self
    }

    pub fn with_linear_backoff(mut self) -> Self {
        self.retry_strategy = RetryStrategy::LinearBackoff;
        self
    }

    pub fn with_initial_delay(mut self, delay: Duration) -> Self {
        self.initial_delay = delay;
        self
    }

    pub fn with_max_delay(mut self, delay: Duration) -> Self {
        self.max_delay = delay;
        self
    }

    pub fn with_jitter(mut self, jitter: bool) -> Self {
        self.jitter = jitter;
        self
    }

    pub fn build(self) -> RetryOptions {
        RetryOptions {
            max_attempts: self.max_attempts,
            retry_strategy: self.retry_strategy,
            initial_delay: self.initial_delay,
            max_delay: self.max_delay,
            jitter: self.jitter,
        }
    }
}

pub struct RetryOptions {
    max_attempts: usize,
    retry_strategy: RetryStrategy,
    initial_delay: Duration,
    max_delay: Duration,
    jitter: bool,
}

impl RetryOptions {
    pub fn run<T, TOk, TError>(self, f: T) -> Result<TOk, RetryError<TError>>
    where
        T: Fn() -> Result<TOk, Retryable<TError>>,
    {
        let mut last_error = None;

        for attempt in 1..=self.max_attempts {
            match f() {
                Ok(success) => return Ok(success),
                Err(Retryable::Permanent { error }) => {
                    return Err(RetryError::Permanent { error });
                }
                Err(Retryable::Transient { error }) => {
                    last_error = Some(error);

                    if attempt < self.max_attempts {
                        let delay = self.get_delay(attempt);
                        std::thread::sleep(delay);
                    }
                }
            }
        }

        Err(RetryError::Exhausted {
            last_error: last_error.unwrap(),
            attempts: self.max_attempts,
        })
    }

    #[cfg(feature = "tokio")]
    pub async fn run_async<T, TFut, TOk, TError>(self, f: T) -> Result<TOk, RetryError<TError>>
    where
        T: Fn() -> TFut + Send + 'static,
        TFut: Future<Output = Result<TOk, Retryable<TError>>>,
    {
        let mut last_error = None;

        for attempt in 1..=self.max_attempts {
            match f().await {
                Ok(success) => return Ok(success),
                Err(Retryable::Permanent { error }) => {
                    return Err(RetryError::Permanent { error });
                }
                Err(Retryable::Transient { error }) => {
                    last_error = Some(error);

                    if attempt < self.max_attempts {
                        let delay = self.get_delay(attempt);
                        tokio::time::sleep(delay).await;
                    }
                }
            }
        }

        Err(RetryError::Exhausted {
            last_error: last_error.unwrap(),
            attempts: self.max_attempts,
        })
    }

    fn get_delay(&self, attempt: usize) -> Duration {
        let base = match self.retry_strategy {
            RetryStrategy::LinearBackoff => self.initial_delay * attempt as u32,
            RetryStrategy::ExponentialBackoff => self.initial_delay * 2u32.pow(attempt as u32 - 1),
        };

        let capped = base.min(self.max_delay);

        if self.jitter {
            // Random factor in [0.5, 1.0) to avoid thundering herd
            let jitter_factor = 0.5 + fastrand::f64() * 0.5;
            capped.mul_f64(jitter_factor)
        } else {
            capped
        }
    }
}

pub enum Retryable<E> {
    Transient { error: E },
    Permanent { error: E },
}

impl<E> Retryable<E> {
    pub fn permanent(error: E) -> Self {
        Self::Permanent { error }
    }

    pub fn transient(error: E) -> Self {
        Self::Transient { error }
    }
}

impl<E> From<E> for Retryable<E> {
    fn from(value: E) -> Self {
        Self::Transient { error: value }
    }
}

#[derive(Debug)]
pub enum RetryError<E> {
    /// All retry attempts were exhausted.
    Exhausted {
        last_error: E,
        attempts: usize,
    },
    Permanent {
        error: E,
    },
}

impl<E> RetryError<E> {
    /// Extract the underlying error, discarding retry context.
    pub fn into_inner(self) -> E {
        match self {
            RetryError::Exhausted { last_error, .. } => last_error,
            RetryError::Permanent { error } => error,
        }
    }
}

impl<E: fmt::Display> fmt::Display for RetryError<E> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            RetryError::Exhausted {
                last_error,
                attempts,
            } => write!(f, "retry exhausted after {attempts} attempts: {last_error}"),
            RetryError::Permanent { error } => {
                write!(f, "{error}")
            }
        }
    }
}

impl<E: std::error::Error + 'static> std::error::Error for RetryError<E> {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            RetryError::Exhausted { last_error, .. } => Some(last_error),
            RetryError::Permanent { error } => Some(error),
        }
    }
}
