// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use anyhow::{anyhow, Context, Result};
use async_trait::async_trait;
use backoff::{self, future::retry_notify, ExponentialBackoff};
use onefuzz_telemetry::debug;
use reqwest::{Response, StatusCode};
use std::{
    sync::atomic::{AtomicUsize, Ordering},
    time::Duration,
};

pub const DEFAULT_RETRY_PERIOD: Duration = Duration::from_secs(5);
pub const MAX_RETRY_ATTEMPTS: usize = 5;

pub enum RetryCheck {
    Retry,
    Fail,
    Succeed,
}

#[derive(Debug, thiserror::Error)]
pub enum ReqwestRetryError {
    #[error("request failed with status code {status_code} and url {url}")]
    Response {
        status_code: StatusCode,
        url: reqwest::Url,
        source: anyhow::Error,
    },
    #[error("request failed to be sent")]
    SendError { source: anyhow::Error },
}

impl ReqwestRetryError {
    fn response_error(status_code: StatusCode, url: reqwest::Url, source: anyhow::Error) -> Self {
        Self::Response {
            status_code,
            url,
            source,
        }
    }

    fn send_error(source: anyhow::Error) -> Self {
        Self::SendError { source }
    }
}

fn always_retry(_: StatusCode) -> RetryCheck {
    RetryCheck::Retry
}

pub async fn send_retry_reqwest_default<
    F: Fn() -> Result<reqwest::RequestBuilder> + Send + Sync,
>(
    build_request: F,
) -> Result<Response> {
    send_retry_reqwest(
        build_request,
        |_| RetryCheck::Retry,
        DEFAULT_RETRY_PERIOD,
        MAX_RETRY_ATTEMPTS,
    )
    .await
}

pub async fn send_retry_reqwest<F, R>(
    build_request: F,
    check_status: R,
    retry_period: Duration,
    max_retry: usize,
) -> Result<Response>
where
    F: Fn() -> Result<reqwest::RequestBuilder> + Send + Sync,
    R: Fn(StatusCode) -> RetryCheck + Send + Sync,
{
    let counter = AtomicUsize::new(0);
    let op = || async {
        let attempt_count = counter.fetch_add(1, Ordering::SeqCst);
        let request = build_request()
            .map_err(|err| backoff::Error::Permanent(ReqwestRetryError::send_error(err)))?;
        let result = request
            .send()
            .await
            .with_context(|| format!("request attempt {} failed", attempt_count + 1));
        match result {
            Err(x) => {
                if attempt_count >= max_retry {
                    Err(backoff::Error::Permanent(ReqwestRetryError::send_error(x)))
                } else {
                    Err(backoff::Error::transient(ReqwestRetryError::send_error(x)))
                }
            }
            Ok(x) => {
                if x.status().is_success() {
                    Ok(x)
                } else {
                    let status = x.status();
                    let result = check_status(status);
                    let url = x.url().clone();

                    match result {
                        RetryCheck::Succeed => Ok(x),
                        RetryCheck::Fail => {
                            let content = x.text().await.unwrap_or_else(|_| "".to_string());
                            let e = anyhow!(
                                "request attempt {} failed with status code {} and content {}",
                                attempt_count + 1,
                                status,
                                content
                            );

                            Err(backoff::Error::Permanent(
                                ReqwestRetryError::response_error(status, url, e),
                            ))
                        }
                        RetryCheck::Retry => {
                            let content = x.text().await.unwrap_or_else(|_| "".to_string());
                            let e = anyhow!(
                                "request attempt {} failed with status code {} and content {}",
                                attempt_count + 1,
                                status,
                                content
                            );

                            if attempt_count >= max_retry {
                                Err(backoff::Error::Permanent(
                                    ReqwestRetryError::response_error(status, url, e),
                                ))
                            } else {
                                Err(backoff::Error::transient(
                                    ReqwestRetryError::response_error(status, url, e),
                                ))
                            }
                        }
                    }
                }
            }
        }
    };
    let result = retry_notify(
        ExponentialBackoff {
            current_interval: retry_period,
            initial_interval: retry_period,
            ..ExponentialBackoff::default()
        },
        op,
        |err: ReqwestRetryError, dur| debug!("request attempt failed after {:?}: {:?}", dur, err),
    )
    .await;

    match result {
        Ok(response) => Ok(response),
        Err(error) => Err(error.into()),
    }
}

#[async_trait]
pub trait SendRetry {
    async fn send_retry<R>(
        self,
        check_status: R,
        retry_period: Duration,
        max_retry: usize,
    ) -> Result<Response>
    where
        R: Fn(StatusCode) -> RetryCheck + Send + Sync;
    async fn send_retry_default(self) -> Result<Response>;
}

#[async_trait]
impl SendRetry for reqwest::RequestBuilder {
    async fn send_retry_default(self) -> Result<Response> {
        self.send_retry(always_retry, DEFAULT_RETRY_PERIOD, MAX_RETRY_ATTEMPTS)
            .await
    }

    async fn send_retry<R>(
        self,
        check_status: R,
        retry_period: Duration,
        max_retry: usize,
    ) -> Result<Response>
    where
        R: Fn(StatusCode) -> RetryCheck + Send + Sync,
    {
        let result = send_retry_reqwest(
            || {
                self.try_clone().ok_or_else(|| {
                    anyhow::Error::msg("This request cannot be retried because it cannot be cloned")
                })
            },
            check_status,
            retry_period,
            max_retry,
        )
        .await?;

        Ok(result)
    }
}

pub fn is_auth_failure(response: &Result<Response>) -> bool {
    // Check both cases to support `error_for_status()`.
    match response {
        Ok(response) => response.status() == StatusCode::UNAUTHORIZED,
        Err(error) => match error.downcast_ref::<ReqwestRetryError>() {
            Some(ReqwestRetryError::Response {
                status_code,
                url: _,
                source: _,
            }) => status_code == &StatusCode::UNAUTHORIZED,
            _ => false,
        },
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use wiremock::matchers::path;
    use wiremock::{Mock, MockServer, ResponseTemplate};

    async fn build_server() -> Result<MockServer> {
        let server = MockServer::start().await;
        server
            .register(Mock::given(path("/200")).respond_with(ResponseTemplate::new(200)))
            .await;
        server
            .register(Mock::given(path("/400")).respond_with(ResponseTemplate::new(400)))
            .await;
        server
            .register(Mock::given(path("/401")).respond_with(ResponseTemplate::new(401)))
            .await;
        server
            .register(Mock::given(path("/404")).respond_with(ResponseTemplate::new(404)))
            .await;
        Ok(server)
    }

    fn always_fail(_: StatusCode) -> RetryCheck {
        RetryCheck::Fail
    }

    fn succeed_400(code: StatusCode) -> RetryCheck {
        match code {
            StatusCode::BAD_REQUEST => RetryCheck::Succeed,
            _ => RetryCheck::Retry,
        }
    }

    #[tokio::test]
    async fn retry_success() -> Result<()> {
        let server = build_server().await?;
        reqwest::Client::new()
            .get(format!("{}/200", &server.uri()))
            .send_retry_default()
            .await?
            .error_for_status()?;
        Ok(())
    }

    #[tokio::test]
    async fn retry_socket_failure() -> Result<()> {
        let server = build_server().await?;
        let resp = reqwest::Client::new()
            .get(format!("{}/404", &server.uri()))
            .send_retry(always_retry, Duration::from_millis(1), 3)
            .await;

        match resp {
            Ok(result) => {
                anyhow::bail!("response should have failed: {:?}", result);
            }
            Err(err) => {
                let as_text = format!("{err:?}");
                assert!(as_text.contains("request attempt 4 failed"), "{}", as_text);
            }
        }

        Ok(())
    }

    #[tokio::test]
    async fn retry_fail_normal() -> Result<()> {
        let server = build_server().await?;
        let resp = reqwest::Client::new()
            .get(format!("{}/400", &server.uri()))
            .send_retry(always_retry, Duration::from_millis(1), 3)
            .await;

        match resp {
            Ok(result) => {
                anyhow::bail!("response should have failed: {:?}", result);
            }
            Err(err) => {
                let as_text = format!("{err:?}");
                assert!(as_text.contains("request attempt 4 failed"), "{}", as_text);
            }
        }

        Ok(())
    }

    #[tokio::test]
    async fn retry_fail_fast() -> Result<()> {
        let server = build_server().await?;
        let resp = reqwest::Client::new()
            .get(format!("{}/400", &server.uri()))
            .send_retry(always_fail, Duration::from_millis(1), 3)
            .await;

        assert!(resp.is_err(), "{resp:?}");
        let as_text = format!("{resp:?}");
        assert!(as_text.contains("request attempt 1 failed"), "{}", as_text);
        Ok(())
    }

    #[tokio::test]
    async fn retry_400_success() -> Result<()> {
        let server = build_server().await?;
        let resp = reqwest::Client::new()
            .get(format!("{}/400", &server.uri()))
            .send_retry(succeed_400, Duration::from_millis(1), 3)
            .await?;

        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
        Ok(())
    }

    #[tokio::test]
    async fn retry_400_with_retry() -> Result<()> {
        let server = build_server().await?;
        let resp = reqwest::Client::new()
            .get(format!("{}/401", &server.uri()))
            .send_retry(succeed_400, Duration::from_millis(1), 3)
            .await;

        assert!(resp.is_err(), "{resp:?}");
        let as_text = format!("{resp:?}");
        assert!(as_text.contains("request attempt 4 failed"), "{}", as_text);
        Ok(())
    }

    #[tokio::test]
    async fn test_auth_failure() -> Result<()> {
        let server = build_server().await?;

        let auth_pass = reqwest::Client::new()
            .get(format!("{}/200", &server.uri()))
            .send_retry(
                |code| match code {
                    StatusCode::UNAUTHORIZED => RetryCheck::Fail,
                    _ => RetryCheck::Retry,
                },
                DEFAULT_RETRY_PERIOD,
                MAX_RETRY_ATTEMPTS,
            )
            .await;

        let auth_fail = reqwest::Client::new()
            .get(format!("{}/401", &server.uri()))
            .send_retry(
                |code| match code {
                    StatusCode::UNAUTHORIZED => RetryCheck::Fail,
                    _ => RetryCheck::Retry,
                },
                DEFAULT_RETRY_PERIOD,
                MAX_RETRY_ATTEMPTS,
            )
            .await;

        assert!(!is_auth_failure(&auth_pass));
        assert!(is_auth_failure(&auth_fail));

        Ok(())
    }
}
