use async_mutex::Mutex as AsyncMutex;
use async_trait::async_trait;
use std::rc::Rc;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum CredentialsError {
    #[error("Failed to acquire credentials: {0}")]
    AcquireFailed(String),
    #[error("Response returned no credentials")]
    MissingCredentials,
    #[error("Response was missing credential: {0}")]
    MissingCredential(&'static str),
    #[error("{0}")]
    Other(String),
}

#[async_trait]
pub trait CredentialsProvider {
    type Credentials;

    /// Acquire new credentials from the credentials provider.
    ///
    /// Calling this method will often result in a network request to an identity provider
    async fn acquire_credentials(&self) -> Result<Self::Credentials, CredentialsError>;

    /// Check whether the credentials since the last call to `acquire_credentials` have expired
    ///
    /// This method should not make any network requests to identity providers so may still return
    /// true even if a token has been revoked.
    fn is_expired(&self) -> bool;
}

/// A proxy and cache for a [`CredentialsProvider`] that handles refetching when credentials have
/// expired
pub struct CredentialsProxy<T, C> {
    provider: Box<dyn CredentialsProvider<Credentials = C>>,
    generate: Box<dyn Fn(C) -> T>,
    cache: AsyncMutex<Option<Rc<T>>>,
}

impl<T, C> CredentialsProxy<T, C> {
    /// Create a new CredentialsProxy from a provider and a generate function.
    pub fn new<P: CredentialsProvider<Credentials = C> + 'static, F: (Fn(C) -> T) + 'static>(
        provider: P,
        generate: F,
    ) -> CredentialsProxy<T, C> {
        Self {
            provider: Box::new(provider),
            generate: Box::new(generate),
            cache: AsyncMutex::new(None),
        }
    }

    /// Return AWS credentials based on the internal provider
    ///
    /// If there are no cached credentials or the credentials are expired, the proxy will use the
    /// internal CredentialsProvider to get new ones.
    pub async fn get(&self) -> Result<Rc<T>, CredentialsError> {
        // Lock the cache to prevent race conditions where two tasks update the credentials at the
        // same time.
        let mut cache_guard = self.cache.lock().await;

        if let Some(cached) = &*cache_guard {
            if !self.provider.is_expired() {
                return Ok(cached.clone());
            }
        }

        let creds = self.provider.acquire_credentials().await?;

        let value = Rc::new((self.generate)(creds));

        cache_guard.replace(value.clone());

        Ok(value)
    }
}

#[cfg(test)]
mod test {
    use super::{AsyncMutex, CredentialsError, CredentialsProvider, CredentialsProxy};
    use async_trait::async_trait;
    use std::{
        borrow::{Borrow, BorrowMut},
        ops::DerefMut,
        sync::{
            atomic::{AtomicBool, Ordering},
            Arc,
        },
    };

    struct TestCredentialsProvider {
        credentials: Arc<AsyncMutex<String>>,
        expired: Arc<AtomicBool>,
    }

    #[async_trait]
    impl CredentialsProvider for TestCredentialsProvider {
        type Credentials = String;

        async fn acquire_credentials(&self) -> Result<Self::Credentials, CredentialsError> {
            let mutex: &AsyncMutex<_> = self.credentials.borrow();
            let value = &*mutex.lock().await;
            Ok(value.clone())
        }

        fn is_expired(&self) -> bool {
            self.expired.load(Ordering::Relaxed)
        }
    }

    fn create_provider() -> (
        TestCredentialsProvider,
        Arc<AsyncMutex<String>>,
        Arc<AtomicBool>,
    ) {
        let expired = Arc::new(AtomicBool::new(false));
        let credentials = Arc::new(AsyncMutex::new("secret-credentials".into()));

        let provider = TestCredentialsProvider {
            expired: expired.clone(),
            credentials: credentials.clone(),
        };

        (provider, credentials, expired)
    }

    #[tokio::test]
    async fn test_generate_credentials() {
        let (provider, _, _) = create_provider();

        let access_key_proxy = CredentialsProxy::new(provider, |creds| format!("{}-cached", creds));

        let result = access_key_proxy.get().await.unwrap();

        assert_eq!(&*result, "secret-credentials-cached");
    }

    #[tokio::test]
    async fn test_use_cached_value() {
        let (provider, credentials, _) = create_provider();

        let access_key_proxy = CredentialsProxy::new(provider, |creds| format!("{}-cached", creds));

        let result = access_key_proxy.get().await.unwrap();

        assert_eq!(&*result, "secret-credentials-cached");

        {
            let mut guard = credentials.lock_arc().await;
            let target = guard.borrow_mut().deref_mut();
            *target = String::from("new-secret-credentials");
        }

        let result = access_key_proxy.get().await.unwrap();

        // Even though the underlying credentials have changed, it is still returning the cached /
        // memoed value
        assert_eq!(&*result, "secret-credentials-cached");
    }

    #[tokio::test]
    async fn test_regen_after_expiry() {
        let (provider, credentials, expired) = create_provider();

        let access_key_proxy = CredentialsProxy::new(provider, |creds| format!("{}-cached", creds));

        let result = access_key_proxy.get().await.unwrap();

        assert_eq!(&*result, "secret-credentials-cached");

        {
            let mut guard = credentials.lock_arc().await;
            let target = guard.borrow_mut().deref_mut();
            *target = String::from("new-secret-credentials")
        }

        expired.store(true, Ordering::Relaxed);

        let result = access_key_proxy.get().await.unwrap();

        // Since expired returned true - the credentials are refreshed
        assert_eq!(&*result, "new-secret-credentials-cached");
    }
}