use std::time::Duration;

use rand::Rng;
use tokio::time::sleep;
use tracing::warn;

use crate::error::{A2aClientError, Result};

#[derive(Clone, Debug)]
pub struct RetryPolicy {
    pub max_retries: u32,
    pub initial_delay_ms: u64,
    pub max_delay_ms: u64,
    pub backoff_multiplier: f64,
    pub jitter: bool,
}

impl Default for RetryPolicy {
    fn default() -> Self {
        Self {
            max_retries: 3,
            initial_delay_ms: 100,
            max_delay_ms: 5000,
            backoff_multiplier: 2.0,
            jitter: true,
        }
    }
}

impl RetryPolicy {
    /// Execute an operation with exponential backoff retry
    pub async fn execute<F, Fut, T>(&self, mut operation: F) -> Result<T>
    where
        F: FnMut() -> Fut,
        Fut: std::future::Future<Output = Result<T>>,
    {
        let mut attempts = 0;
        let mut delay_ms = self.initial_delay_ms;

        loop {
            match operation().await {
                Ok(result) => return Ok(result),
                Err(e) if !Self::is_retryable(&e) => {
                    // Non-retryable error (4xx client errors)
                    return Err(e);
                }
                Err(e) if attempts >= self.max_retries => {
                    // Exhausted retries
                    return Err(e);
                }
                Err(e) => {
                    attempts += 1;

                    // Apply jitter if enabled (randomize ±20%)
                    let actual_delay = if self.jitter {
                        let jitter_factor = rand::rng().random_range(0.8..1.2);
                        (delay_ms as f64 * jitter_factor) as u64
                    } else {
                        delay_ms
                    };

                    warn!(
                        error = %e,
                        attempt = attempts,
                        max_retries = self.max_retries,
                        delay_ms = actual_delay,
                        "Retrying failed operation"
                    );

                    sleep(Duration::from_millis(actual_delay)).await;

                    // Calculate next delay with exponential backoff
                    delay_ms =
                        ((delay_ms as f64 * self.backoff_multiplier) as u64).min(self.max_delay_ms);
                }
            }
        }
    }

    /// Determine if an error is retryable
    fn is_retryable(error: &A2aClientError) -> bool {
        match error {
            // Retry network errors
            A2aClientError::Timeout(_) => true,
            A2aClientError::ConnectionRefused(_) => true,
            A2aClientError::HttpError(e) => {
                // Retry on connection/timeout errors
                e.is_timeout() || e.is_connect()
            }
            // Retry server errors (5xx)
            A2aClientError::ServerError { code, .. } => *code >= 500 && *code < 600,
            // Do NOT retry client errors (4xx) or deserialization errors
            A2aClientError::InvalidResponse => false,
            A2aClientError::TaskNotFound(_) => false,
            A2aClientError::SerdeError(_) => false,
            A2aClientError::InternalError(_) => false,
        }
    }
}

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

    #[tokio::test]
    async fn test_retry_succeeds_eventually() {
        use std::sync::atomic::{AtomicU32, Ordering};
        use std::sync::Arc;

        let policy = RetryPolicy {
            max_retries: 3,
            initial_delay_ms: 10,
            max_delay_ms: 100,
            backoff_multiplier: 2.0,
            jitter: false,
        };

        let attempt = Arc::new(AtomicU32::new(0));
        let attempt_clone = attempt.clone();

        #[allow(clippy::excessive_nesting)]
        let result = policy
            .execute(|| {
                let attempt = attempt_clone.clone();
                async move {
                    let current = attempt.fetch_add(1, Ordering::SeqCst);
                    if current < 2 {
                        Err(A2aClientError::Timeout("test".to_string()))
                    } else {
                        Ok("success")
                    }
                }
            })
            .await;

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "success");
        assert_eq!(attempt.load(Ordering::SeqCst), 3);
    }

    #[tokio::test]
    async fn test_retry_exhausted() {
        use std::sync::atomic::{AtomicU32, Ordering};
        use std::sync::Arc;

        let policy = RetryPolicy {
            max_retries: 2,
            initial_delay_ms: 10,
            max_delay_ms: 100,
            backoff_multiplier: 2.0,
            jitter: false,
        };

        let attempt = Arc::new(AtomicU32::new(0));
        let attempt_clone = attempt.clone();

        let result = policy
            .execute(|| {
                let attempt = attempt_clone.clone();
                async move {
                    attempt.fetch_add(1, Ordering::SeqCst);
                    Err::<(), _>(A2aClientError::Timeout("test".to_string()))
                }
            })
            .await;

        assert!(result.is_err());
        assert_eq!(attempt.load(Ordering::SeqCst), 3); // initial + 2 retries
    }

    #[tokio::test]
    async fn test_non_retryable_error() {
        use std::sync::atomic::{AtomicU32, Ordering};
        use std::sync::Arc;

        let policy = RetryPolicy::default();

        let attempt = Arc::new(AtomicU32::new(0));
        let attempt_clone = attempt.clone();

        let result = policy
            .execute(|| {
                let attempt = attempt_clone.clone();
                async move {
                    attempt.fetch_add(1, Ordering::SeqCst);
                    Err::<(), _>(A2aClientError::TaskNotFound("test".to_string()))
                }
            })
            .await;

        assert!(result.is_err());
        assert_eq!(attempt.load(Ordering::SeqCst), 1); // No retries for
                                                       // non-retryable error
    }
}
