use http::Extensions;
use reqwest::{Request, Response};
use reqwest_middleware::{Middleware, Next, Result};
use reqwest_retry::{
    DefaultRetryableStrategy, RetryPolicy, RetryTransientMiddleware, RetryableStrategy,
};

/// Custom extension options for the request
///
/// The `FWClient` uses a custom [`middleware`](FWRetryTransientMiddleware) to handle retry logic.
/// the `FWOptions` struct allows you to control the behavior of the middleware.
#[derive(Debug, Clone)]
pub struct FWOptions {
    /// If true, the middleware will skip retry logic.
    pub skip_retry: bool,
    /// If true, the middleware will remove the `Authorization` header from the request.
    pub no_auth: bool,
}

pub struct FWRetryTransientMiddleware<
    T: RetryPolicy + Send + Sync + 'static,
    R: RetryableStrategy + Send + Sync + 'static = DefaultRetryableStrategy,
>(Option<RetryTransientMiddleware<T, R>>);

impl<T, R> FWRetryTransientMiddleware<T, R>
where
    T: RetryPolicy + Send + Sync + 'static,
    R: RetryableStrategy + Send + Sync + 'static,
{
    /// Create a new instance of `FWRetryTransientMiddleware` with the provided
    /// [`RetryTransientMiddleware`].
    pub fn new(retry: Option<RetryTransientMiddleware<T, R>>) -> Self {
        Self(retry)
    }
}

#[async_trait::async_trait]
impl<T, R> Middleware for FWRetryTransientMiddleware<T, R>
where
    T: RetryPolicy + Send + Sync,
    R: RetryableStrategy + Send + Sync + 'static,
{
    async fn handle(
        &self,
        mut req: Request,
        extensions: &mut Extensions,
        next: Next<'_>,
    ) -> Result<Response> {
        // If the request cannot be cloned (try_clone returns None) we skip retry logic
        let mut skip_retry = req.try_clone().is_none();
        if let Some(options) = extensions.get::<FWOptions>() {
            if options.no_auth {
                req.headers_mut().remove("Authorization");
            }
            if options.skip_retry {
                skip_retry = true;
            }
        }
        if skip_retry {
            return next.run(req, extensions).await;
        }
        match &self.0 {
            Some(middleware) => middleware.handle(req, extensions, next).await,
            None => next.run(req, extensions).await,
        }
    }
}
