use std::future::Future;
use std::time::Duration;

use derive_new::new;
use derive_setters::Setters;
use tokio::time::sleep;
use tracing::warn;

#[derive(Debug, Clone, new, Setters)]
#[setters(prefix = "with_")]
pub struct Retry {
    #[setters(skip)]
    pub max_attempts: u32,

    #[new(value = "Duration::from_millis(100)")]
    pub initial_backoff: Duration,

    #[new(value = "Duration::from_secs(2)")]
    pub max_backoff: Duration,

    #[new(value = "true")]
    pub retry_on_no_reply: bool,

    #[new(value = "true")]
    pub retry_on_transport_error: bool,
}

impl Retry {
    fn validate(&self) -> crate::bus::Result<()> {
        if self.max_attempts == 0 {
            return Err(crate::bus::Error::InvalidRetry(
                "max_attempts must be greater than zero".to_string(),
            ));
        }
        if self.initial_backoff.is_zero() {
            return Err(crate::bus::Error::InvalidRetry(
                "initial_backoff must be greater than zero".to_string(),
            ));
        }
        if self.max_backoff.is_zero() {
            return Err(crate::bus::Error::InvalidRetry(
                "max_backoff must be greater than zero".to_string(),
            ));
        }
        if self.initial_backoff > self.max_backoff {
            return Err(crate::bus::Error::InvalidRetry(
                "initial_backoff must be less than or equal to max_backoff".to_string(),
            ));
        }
        Ok(())
    }
}

pub async fn retry_query<T, F, Fut>(
    name: &str,
    retry: &Retry,
    mut attempt: F,
) -> crate::bus::Result<Option<T>>
where
    F: FnMut() -> Fut,
    Fut: Future<Output = crate::bus::Result<Option<T>>>,
{
    retry.validate()?;

    let mut backoff = retry.initial_backoff;
    for attempt_index in 1..=retry.max_attempts {
        match attempt().await {
            Ok(Some(response)) => return Ok(Some(response)),
            Ok(None) => {
                if !retry.retry_on_no_reply || attempt_index == retry.max_attempts {
                    return Ok(None);
                }
                warn!(
                    query = name,
                    attempt = attempt_index,
                    max_attempts = retry.max_attempts,
                    backoff_ms = backoff.as_millis(),
                    "query returned no reply, retrying"
                );
            }
            Err(crate::bus::Error::Transport(error)) => {
                if !retry.retry_on_transport_error || attempt_index == retry.max_attempts {
                    return Err(crate::bus::Error::Transport(error));
                }
                warn!(
                    query = name,
                    attempt = attempt_index,
                    max_attempts = retry.max_attempts,
                    backoff_ms = backoff.as_millis(),
                    error = %error,
                    "query failed with transport error, retrying"
                );
            }
            Err(error) => return Err(error),
        }

        sleep(backoff).await;
        backoff = std::cmp::min(backoff.saturating_mul(2), retry.max_backoff);
    }

    Ok(None)
}

#[cfg(test)]
mod tests {
    use super::{Retry, retry_query};
    use crate::bus::{Error, Result};
    use std::sync::{
        Arc,
        atomic::{AtomicU32, Ordering},
    };
    use std::time::Duration;

    #[tokio::test]
    async fn retry_query_retries_no_reply_until_success() -> Result<()> {
        let attempts = Arc::new(AtomicU32::new(0));
        let retry = Retry::new(3).with_initial_backoff(Duration::from_millis(1));
        let result = retry_query("robot/r1/asset/bundle", &retry, {
            let attempts = attempts.clone();
            move || {
                let attempts = attempts.clone();
                async move {
                    let current = attempts.fetch_add(1, Ordering::SeqCst);
                    if current < 2 {
                        Ok(None)
                    } else {
                        Ok(Some("model"))
                    }
                }
            }
        })
        .await?;

        assert_eq!(result, Some("model"));
        assert_eq!(attempts.load(Ordering::SeqCst), 3);
        Ok(())
    }

    #[tokio::test]
    async fn retry_query_retries_transport_errors_until_success() -> Result<()> {
        let attempts = Arc::new(AtomicU32::new(0));
        let retry = Retry::new(2).with_initial_backoff(Duration::from_millis(1));
        let result = retry_query("robot/r1/asset/bundle", &retry, {
            let attempts = attempts.clone();
            move || {
                let attempts = attempts.clone();
                async move {
                    let current = attempts.fetch_add(1, Ordering::SeqCst);
                    if current == 0 {
                        Err(Error::Transport(std::io::Error::other("transient").into()))
                    } else {
                        Ok(Some("model"))
                    }
                }
            }
        })
        .await?;

        assert_eq!(result, Some("model"));
        assert_eq!(attempts.load(Ordering::SeqCst), 2);
        Ok(())
    }

    #[tokio::test]
    async fn retry_query_rejects_invalid_retry() {
        let error = retry_query::<(), _, _>("robot/r1/asset/bundle", &Retry::new(0), || async {
            Ok(Some(()))
        })
        .await
        .expect_err("invalid retry should fail");
        assert!(matches!(error, Error::InvalidRetry(_)));
    }
}
