use std::{
    sync::{Arc, Weak},
    time::{Duration, SystemTime},
};

use arc_swap::{ArcSwap, ArcSwapOption};
use async_trait::async_trait;
use tokio::runtime::Handle;
use tokio::sync::{
    Notify,
    watch::{self, Receiver, Sender},
};
use kameleoon_trace_logging::trace_method;

use crate::{
    debug_log, error_log,
    logging::LogLevel,
    network::{dto::Jwt, network_manager::automation_api::AutomationApiNetworkManager},
};

// MARK: - ATS Trait

#[async_trait]
pub(crate) trait AccessTokenSource: Send + Sync {
    fn get_token_now(&self) -> Option<String>;
    async fn discard_token(&self, token: String);
    #[cfg(test)]
    async fn get_token(&self) -> Option<String>; // can be used in case if token is required for request, but not optional
}

// MARK: - ATS Impl

pub(crate) const TOKEN_EXPIRATION_GAP: Duration = Duration::from_secs(60);
pub(crate) const RETRY_INTERVAL: Duration = Duration::from_secs(300);
const BEARER_PREFIX: &str = "Bearer ";

#[allow(dead_code)]
pub(crate) struct AccessTokenSourceImpl {
    automation_api: Weak<dyn AutomationApiNetworkManager>,
    client_id: String,
    client_secret: String,
    retry_interval: Duration,
    runtime_handle: Handle,
    stop_polling_notify: ArcSwapOption<Notify>,
    cached_token: ArcSwap<ExpiringToken>,
    rx: Receiver<bool>,
    tx: Sender<bool>,
}

#[async_trait]
impl AccessTokenSource for Arc<AccessTokenSourceImpl> {
    #[trace_method(LogLevel::Debug, skip(self))]
    fn get_token_now(&self) -> Option<String> {
        let cached_token = self.cached_token.load();
        return if cached_token.is_valid() { Some(cached_token.value.clone()) } else { None };
    }

    async fn discard_token(&self, token: String) {
        //~ [42126]
        debug_log!("CALL AccessTokenSourceImpl::discard_token(token: {token})");
        // async keyword only ensures async context - required for tokio::spawn
        let cached_token = self.cached_token.load();
        if token == cached_token.value {
            self.cached_token.compare_and_swap(&cached_token, Arc::new(ExpiringToken::default()));
            drop(cached_token);
            if !self.cached_token.load().is_valid() {
                let ats = self.clone();
                self.runtime_handle.spawn(async move {
                    ats.fetch_token().await;
                });
            }
        }
        debug_log!("RET AccessTokenSourceImpl::discard_token(token: {token})");
    }

    #[cfg(test)]
    async fn get_token(&self) -> Option<String> {
        //~ [42126]
        debug_log!("CALL AccessTokenSourceImpl::get_token()");
        let result = match self.get_token_now() {
            Some(token) => Some(token),
            None => {
                let mut rx = self.rx.clone();
                rx.changed().await.map_or_else(
                    |err| {
                        error_log!("AccessTokenSource::get_token failed to await token fetch: {err}");
                        None
                    },
                    |_| self.get_token_now(),
                )
            }
        };
        debug_log!("RET AccessTokenSourceImpl::get_token() → {:?}", result);
        result
    }
}

impl Drop for AccessTokenSourceImpl {
    #[trace_method(LogLevel::Debug, skip(self))]
    fn drop(&mut self) {
        if let Some(notify) = self.stop_polling_notify.swap(None) {
            notify.notify_one();
        }
    }
}

impl AccessTokenSourceImpl {
    #[trace_method(LogLevel::Debug, skip(retval, automation_api, runtime_handle))]
    pub fn new(
        automation_api: Weak<dyn AutomationApiNetworkManager>,
        client_id: String,
        client_secret: String,
        retry_interval: Duration,
        runtime_handle: Handle,
    ) -> Self {
        let (tx, rx) = watch::channel(false);
        return Self {
            automation_api,
            client_id,
            client_secret,
            retry_interval,
            runtime_handle,
            stop_polling_notify: ArcSwapOption::empty(),
            cached_token: ArcSwap::new(Arc::new(ExpiringToken::default())),
            rx,
            tx,
        };
    }

    pub fn start_polling(self: &Arc<Self>) {
        let notify = Arc::new(Notify::new());
        self.stop_polling_notify.store(Some(notify.clone()));
        let retry_interval = self.retry_interval;
        let ats = Arc::downgrade(self);
        self.runtime_handle.spawn(async move {
            let mut interval = tokio::time::Duration::ZERO;
            loop {
                tokio::select! {
                    _ = notify.notified() => break,
                    _ = tokio::time::sleep(interval) => match ats.upgrade() {
                        Some(ats) => interval = if let Some(expires_in) = ats.fetch_token().await {
                            expires_in
                        } else {
                            retry_interval
                        },
                        None => break,
                    }
                }
            }
        });
    }

    async fn fetch_token(self: &Arc<Self>) -> Option<Duration> {
        let jwt = self.fetch_jwt().await;
        if jwt.is_none() {
            self.tx.send_replace(false);
            return None;
        }
        let jwt = jwt.unwrap();
        let expires_in = jwt.expires_in;
        self.cached_token.store(Arc::new(jwt.into()));
        self.tx.send_replace(true);
        return Some(expires_in);
    }

    async fn fetch_jwt(self: &Arc<Self>) -> Option<Jwt> {
        let response = match self.automation_api.upgrade()?.fetch_access_jwt(&self.client_id, &self.client_secret).await
        {
            Ok(r) => r,
            Err(err) => {
                error_log!("Failed to get access JWT: {err}");
                return None;
            }
        };

        if response.status_code / 100 != 2 {
            error_log!("Failed to get access JWT: unexpected status code {}", response.status_code);
            return None;
        }

        let mut jwt: Jwt = match response.json() {
            Ok(jwt) => jwt,
            Err(err) => {
                error_log!("Failed to parse access JWT: {err}");
                return None;
            }
        };

        if jwt.expires_in <= TOKEN_EXPIRATION_GAP {
            let expires_in = jwt.expires_in.as_secs();
            error_log!("Access token life time ({expires_in}s) is not long enough to use it");
            return None;
        }
        jwt.expires_in -= TOKEN_EXPIRATION_GAP;
        Some(jwt)
    }
}

// MARK: - Expiring Token

struct ExpiringToken {
    value: String,
    expiration_time: SystemTime,
}

impl ExpiringToken {
    #[inline]
    fn is_valid(&self) -> bool {
        return SystemTime::now() < self.expiration_time;
    }
}

impl Default for ExpiringToken {
    fn default() -> Self {
        ExpiringToken {
            value: "".to_owned(),
            expiration_time: SystemTime::UNIX_EPOCH,
        }
    }
}

impl From<Jwt> for ExpiringToken {
    fn from(jwt: Jwt) -> Self {
        return ExpiringToken {
            value: format!("{BEARER_PREFIX}{}", jwt.access_token),
            expiration_time: SystemTime::now() + jwt.expires_in - TOKEN_EXPIRATION_GAP,
        };
    }
}

// MARK: - Tests

#[cfg(test)]
mod tests {
    use std::{
        any::Any,
        pin::Pin,
        sync::atomic::{AtomicBool, Ordering},
    };

    use crate::{
        assert_delta,
        network::{
            network_error::NetworkError, network_manager::automation_api::MockAutomationApiNetworkManager,
            response::Response,
        },
    };

    use super::*;
    use rstest::*;
    use tokio::time::timeout;

    const CLIENT_ID: &str = "cid";
    const CLIENT_SECRET: &str = "cs";

    async fn wait() {
        tokio::time::sleep(Duration::from_millis(100)).await;
    }

    fn new_ats(
        automation_api: Weak<dyn AutomationApiNetworkManager>,
        runtime_handle: Handle,
    ) -> Arc<AccessTokenSourceImpl> {
        return Arc::new(AccessTokenSourceImpl::new(
            automation_api,
            CLIENT_ID.to_owned(),
            CLIENT_SECRET.to_owned(),
            RETRY_INTERVAL,
            runtime_handle,
        ));
    }

    fn with_bearer_prefix(value: &str) -> String {
        format!("{BEARER_PREFIX}{value}")
    }

    fn make_token(value: &str) -> ExpiringToken {
        ExpiringToken {
            value: with_bearer_prefix(value),
            expiration_time: SystemTime::now() + Duration::from_secs(3600),
        }
    }

    fn make_expired_token(value: &str) -> ExpiringToken {
        ExpiringToken {
            value: with_bearer_prefix(value),
            expiration_time: SystemTime::now() - Duration::from_secs(1),
        }
    }

    fn make_jwt_body(value: (&str, u64)) -> String {
        serde_json::to_string(&Jwt {
            access_token: value.0.to_owned(),
            expires_in: Duration::from_secs(value.1) + TOKEN_EXPIRATION_GAP,
        })
        .unwrap()
    }

    fn make_jwt_response(
        fetched_token: Option<(&str, u64)>,
        code: u16,
    ) -> Pin<Box<std::future::Ready<std::result::Result<Response, NetworkError>>>> {
        Box::pin(std::future::ready(match fetched_token {
            Some(token) => Ok(Response {
                status_code: code,
                body: make_jwt_body(token).into_bytes(),
                headers: Vec::new(),
            }),
            None => Err(NetworkError::Timeout),
        }))
    }

    fn checkpoint_automation_api_mock(automation_api: Arc<dyn AutomationApiNetworkManager>) {
        let automation_api: Arc<dyn Any + Send + Sync> = automation_api;
        let automation_api: Arc<MockAutomationApiNetworkManager> = automation_api.downcast().unwrap();
        let mut automation_api = Arc::into_inner(automation_api).unwrap();
        automation_api.checkpoint();
    }

    #[rstest]
    #[case(make_token("A"), Some(with_bearer_prefix("A")))]
    #[case(make_expired_token("A"), None)]
    #[case(ExpiringToken::default(), None)]
    #[tokio::test]
    async fn test_get_token_now(#[case] cached_token: ExpiringToken, #[case] expected: Option<String>) {
        // arrange
        let mut automation_api = MockAutomationApiNetworkManager::new();
        automation_api.expect_fetch_access_jwt().never();
        let automation_api: Arc<dyn AutomationApiNetworkManager> = Arc::new(automation_api);
        let ats = new_ats(Arc::downgrade(&automation_api), Handle::current());
        ats.cached_token.store(Arc::new(cached_token));
        // act
        let actual = ats.get_token_now();
        // assert
        assert_eq!(expected, actual);
        wait().await;
        checkpoint_automation_api_mock(automation_api);
    }

    #[rstest]
    #[case(make_token("A"), with_bearer_prefix("A"), "", true)]
    #[case(make_token("A"), with_bearer_prefix("B"), "Bearer A", false)]
    #[tokio::test]
    async fn test_discard_token(
        #[case] cached_token: ExpiringToken,
        #[case] token_to_discard: String,
        #[case] expected: &str,
        #[case] expect_fetch: bool,
    ) {
        // arrange
        let response_mutex = Arc::new(std::sync::Mutex::new(()));
        let response_lock = response_mutex.lock().unwrap();
        let mut automation_api = MockAutomationApiNetworkManager::new();
        if expect_fetch {
            let response_mutex = response_mutex.clone();
            automation_api.expect_fetch_access_jwt().once().returning(move |client_id, client_secret| {
                assert_eq!(CLIENT_ID, client_id);
                assert_eq!(CLIENT_SECRET, client_secret);
                drop(response_mutex.lock().unwrap());
                make_jwt_response(Some(("C", 7199)), 200)
            });
        } else {
            automation_api.expect_fetch_access_jwt().never();
        }
        let automation_api: Arc<dyn AutomationApiNetworkManager> = Arc::new(automation_api);
        let expected_cached_token_after_fetch =
            if expect_fetch { with_bearer_prefix("C") } else { cached_token.value.clone() };
        let ats = new_ats(Arc::downgrade(&automation_api), Handle::current());
        ats.cached_token.store(Arc::new(cached_token));
        // act
        ats.discard_token(token_to_discard).await;
        // assert
        assert_eq!(expected, ats.cached_token.load_full().value);
        drop(response_lock);
        wait().await;
        assert_eq!(expected_cached_token_after_fetch, ats.cached_token.load_full().value);
        checkpoint_automation_api_mock(automation_api);
    }

    #[rstest]
    #[case(9999, Some(("A", 1)), 200, 1000, "Bearer A")]
    #[case(100, Some(("A", 10)), 500, 100, "")]
    #[case(100, None, 0, 100, "")]
    #[tokio::test]
    async fn test_polling(
        #[case] retry_interval: u64,
        #[case] fetched_token: Option<(&'static str, u64)>,
        #[case] code: u16,
        #[case] expected_interval: u64,
        #[case] expected_cached_token: &str,
    ) {
        const S: Duration = Duration::from_millis(25);
        // arrange
        let (tx, mut rx) = tokio::sync::mpsc::channel(2);
        let mut automation_api = MockAutomationApiNetworkManager::new();
        automation_api.expect_fetch_access_jwt().times(2).returning(move |client_id, client_secret| {
            assert_eq!(CLIENT_ID, client_id);
            assert_eq!(CLIENT_SECRET, client_secret);
            tx.try_send(SystemTime::now()).unwrap();
            make_jwt_response(fetched_token, code)
        });
        let automation_api: Arc<dyn AutomationApiNetworkManager> = Arc::new(automation_api);
        let ats = Arc::new(AccessTokenSourceImpl::new(
            Arc::downgrade(&automation_api),
            CLIENT_ID.to_owned(),
            CLIENT_SECRET.to_owned(),
            Duration::from_millis(retry_interval),
            Handle::current(),
        ));
        // act
        ats.start_polling();
        // assert
        let start = SystemTime::now();
        let fetch_ts1 = timeout(Duration::from_secs(2), rx.recv()).await.unwrap().unwrap();
        assert_eq!(expected_cached_token, ats.cached_token.load_full().value);
        let fetch_ts2 = timeout(Duration::from_secs(2), rx.recv()).await.unwrap().unwrap();
        assert_eq!(expected_cached_token, ats.cached_token.load_full().value);
        assert_delta!(Duration::ZERO, fetch_ts1.duration_since(start).unwrap(), S);
        assert_delta!(Duration::from_millis(expected_interval), fetch_ts2.duration_since(fetch_ts1).unwrap(), S);
        checkpoint_automation_api_mock(automation_api);
    }

    #[rstest]
    #[tokio::test]
    async fn test_drop(#[values(true, false)] polling_started: bool) {
        // arrange
        let mut automation_api = MockAutomationApiNetworkManager::new();
        if polling_started {
            automation_api
                .expect_fetch_access_jwt()
                .times(0..2)
                .returning(|_, _| make_jwt_response(Some(("AT1", 1)), 200));
        } else {
            automation_api.expect_fetch_access_jwt().never();
        }
        let automation_api: Arc<dyn AutomationApiNetworkManager> = Arc::new(automation_api);
        let ats = Arc::new(AccessTokenSourceImpl::new(
            Arc::downgrade(&automation_api),
            CLIENT_ID.to_owned(),
            CLIENT_SECRET.to_owned(),
            Duration::from_millis(1),
            Handle::current(),
        ));
        if polling_started {
            ats.start_polling();
        }
        // act
        drop(ats);
        // assert
        wait().await;
        if polling_started {
            tokio::time::sleep(Duration::from_secs(1)).await;
        }
        checkpoint_automation_api_mock(automation_api);
    }

    #[rstest]
    #[case(make_token("A"), None, false, Some(with_bearer_prefix("A")), "Bearer A")]
    #[case(make_expired_token("E"), Some(("A", 7199)), true, Some(with_bearer_prefix("A")), "Bearer A")]
    #[case(ExpiringToken::default(), Some(("A", 7199)), true, Some(with_bearer_prefix("A")), "Bearer A")]
    #[case(make_expired_token("E"), None, true, None, "Bearer E")]
    #[case(ExpiringToken::default(), None, true, None, "")]
    #[tokio::test]
    async fn test_get_token(
        #[case] cached_token: ExpiringToken,
        #[case] fetched_token: Option<(&'static str, u64)>,
        #[case] expect_fetch_awaiting: bool,
        #[case] expected: Option<String>,
        #[case] expected_cached_token: &str,
    ) {
        // arrange
        let mut automation_api = MockAutomationApiNetworkManager::new();
        if expect_fetch_awaiting {
            automation_api.expect_fetch_access_jwt().once().returning(move |client_id, client_secret| {
                assert_eq!(CLIENT_ID, client_id);
                assert_eq!(CLIENT_SECRET, client_secret);
                make_jwt_response(fetched_token, 200)
            });
        } else {
            automation_api.expect_fetch_access_jwt().never();
        }
        let automation_api: Arc<dyn AutomationApiNetworkManager> = Arc::new(automation_api);
        let ats = new_ats(Arc::downgrade(&automation_api), Handle::current());
        ats.cached_token.store(Arc::new(cached_token));
        // act
        let actual = if expect_fetch_awaiting {
            let flag = Arc::new(AtomicBool::new(false));
            let flag_clone = flag.clone();
            let ats_clone = ats.clone();
            let handle = tokio::spawn(async move {
                wait().await;
                flag_clone.store(true, Ordering::Relaxed);
                ats_clone.start_polling();
            });
            let actual = timeout(Duration::from_secs(1), ats.get_token()).await.unwrap();
            assert!(flag.load(Ordering::Relaxed));
            handle.await.unwrap();
            actual
        } else {
            timeout(Duration::from_secs(1), ats.get_token()).await.unwrap()
        };
        // assert
        assert_eq!(expected, actual);
        assert_eq!(expected_cached_token, ats.cached_token.load_full().value);
        checkpoint_automation_api_mock(automation_api);
    }
}
