use std::{
    borrow::Cow,
    collections::HashSet,
    path::PathBuf,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use crate::utils::StdError;

use crate::async_value::AsyncAccessor;
use crate::data_service_credentials::DataServiceCredentials;
use async_mutex::Mutex as AsyncMutex;
use async_trait::async_trait;
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use serde_json::{from_reader, json, Value};
use thiserror::Error;

#[derive(Deserialize)]
struct AccessTokenResponse {
    refresh_token: String,
    access_token: String,
    expires_in: u64,
}

#[derive(Deserialize)]
struct PollingInfo {
    user_code: String,
    device_code: String,
    verification_uri_complete: String,
}

#[derive(Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AuthToken {
    refresh_token: String,
    access_token: String,
    expiry: u64,
}

fn now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("Expected system time to be greater than UNIX_EPOCH")
        .as_secs()
}

impl From<AccessTokenResponse> for AuthToken {
    fn from(
        AccessTokenResponse {
            access_token,
            refresh_token,
            expires_in,
        }: AccessTokenResponse,
    ) -> Self {
        Self {
            access_token,
            refresh_token,
            expiry: expires_in + now_secs(),
        }
    }
}

#[derive(Debug)]
pub struct IdpInfo {
    pub uri: String,
    pub client_id: String,
    pub audience: String,
}

pub struct Auth0DeviceCodeCredentials {
    pub(crate) token_path: PathBuf,
    pub(crate) idp_info: IdpInfo,
    pub(crate) cached_token: AsyncMutex<Option<AuthToken>>,
    // Optional workspace to validate
    pub(crate) workspace: Option<String>,
}

#[derive(Error, Debug)]
pub enum DeviceCodeError {
    #[error("RefreshTokenFailed: {0}")]
    RefreshTokenFailed(String),
    #[error("AcquireNewTokenFailed: {0}")]
    AcquireNewTokenFailed(String),
    #[error("PollFailed: {0}")]
    PollFailed(String),
    #[error("InvalidAccessToken: {0}")]
    InvalidAccessToken(String),
    #[error("PersistTokenError: {0}")]
    PersistTokenError(String),
}

const EXPIRY_GRACE_PERIOD_SECONDS: u64 = 120;

async fn sleep(duration: &Duration) {
    cfg_if::cfg_if! {
        if #[cfg(feature = "tokio")] {
            tokio::time::sleep(*duration).await
        } else {
            std::thread::sleep(*duration)
        }
    }
}

impl Auth0DeviceCodeCredentials {
    pub fn new(token_path: PathBuf, idp_info: IdpInfo, workspace: Option<String>) -> Self {
        let cached_token: Option<AuthToken> = std::fs::read_to_string(&token_path)
            .ok()
            .and_then(|x| serde_json::from_str(&x).ok());

        Self {
            token_path,
            idp_info,
            workspace,
            cached_token: AsyncMutex::new(cached_token),
        }
    }

    /// Get the required IDP scopes
    fn idp_scopes(&self) -> Vec<Cow<str>> {
        lazy_static! {
            static ref DEFAULT_SCOPES: Vec<Cow<'static, str>> = vec![
                "collection.create".into(),
                "collection.delete".into(),
                "collection.info".into(),
                "collection.list".into(),
                "document.put".into(),
                "document.delete".into(),
                "document.get".into(),
                "document.query".into(),
                "offline_access".into(),
            ];
        };

        let mut scopes = DEFAULT_SCOPES.clone();

        if let Some(workspace) = &self.workspace {
            scopes.push(format!("ws:{}", workspace).into());
        }

        scopes
    }

    /// Validate that JWT string is a valid access token with the correct scopes and claims without
    /// checking the header or the signature.
    fn validate_token(&self, token: &str) -> Result<(), DeviceCodeError> {
        let mut tokens = token.split('.');

        let _header = tokens
            .next()
            .ok_or_else(|| DeviceCodeError::InvalidAccessToken("Missing token header".into()))?;
        let payload = tokens
            .next()
            .ok_or_else(|| DeviceCodeError::InvalidAccessToken("Missing token payload".into()))?;
        let _signature = tokens
            .next()
            .ok_or_else(|| DeviceCodeError::InvalidAccessToken("Missing token signature".into()))?;

        let json_payload = base64::decode(payload).map_err(|e| {
            DeviceCodeError::InvalidAccessToken(format!("Failed to decode base64: {}", e))
        })?;

        let payload_value: Value = from_reader(&json_payload[..]).map_err(|e| {
            DeviceCodeError::InvalidAccessToken(format!(
                "Failed to load json from body string: {}",
                e
            ))
        })?;

        let recieved_scopes = payload_value
            .get("scope")
            .ok_or_else(|| {
                DeviceCodeError::InvalidAccessToken("Could not find scope value in token".into())
            })?
            .as_str()
            .ok_or_else(|| {
                DeviceCodeError::InvalidAccessToken("Scope value was not a string".into())
            })?
            .split_whitespace()
            .collect::<HashSet<_>>();

        let missing_scopes = self
            .idp_scopes()
            .into_iter()
            .filter(|s| !recieved_scopes.contains(s.as_ref()))
            .collect::<Vec<_>>();

        if !missing_scopes.is_empty() {
            return Err(DeviceCodeError::InvalidAccessToken(format!(
                "Payload was missing scopes: {}",
                missing_scopes.join(", ")
            )));
        }

        // Check for the AWS KMS principal tags if a workspace has been specified.
        if let Some(workspace) = &self.workspace {
            let workspace_tags: Vec<String> = serde_json::from_value(
                payload_value
                    .get("https://aws.amazon.com/tags")
                    .ok_or_else(|| {
                        DeviceCodeError::InvalidAccessToken(
                            "Token does not contain KMS federation claim".into(),
                        )
                    })?
                    .get("principal_tags")
                    .ok_or_else(|| {
                        DeviceCodeError::InvalidAccessToken(
                            "Token does not contain KMS federation claim key principal_tags".into(),
                        )
                    })?
                    .get("workspace")
                    .ok_or_else(|| {
                        DeviceCodeError::InvalidAccessToken(
                            "Token does not contain KMS federation claim key workspace".into(),
                        )
                    })?
                    .clone(),
            )
            .map_err(|_| {
                DeviceCodeError::InvalidAccessToken(
                    "Expected KMS federated claim key workspace to be array of strings".into(),
                )
            })?;

            if !workspace_tags.contains(&format!("ws:{}", workspace)) {
                return Err(DeviceCodeError::InvalidAccessToken(
                    "Auth token KMS federated claim did not include workspace scope".into(),
                ));
            }
        }

        Ok(())
    }

    /// Poll the IDP until it returns an access token or fails
    async fn poll_for_access_token(
        &self,
        client: &reqwest::Client,
        polling_info: &PollingInfo,
    ) -> Result<AuthToken, DeviceCodeError> {
        let mut interval = Duration::from_secs(5);

        loop {
            let response = client
                .post(format!("{}/oauth/token", self.idp_info.uri))
                .json(&json!({
                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
                    "device_code": polling_info.device_code,
                    "client_id": self.idp_info.client_id
                }))
                .send()
                .await
                .map_err(|e| {
                    DeviceCodeError::PollFailed(format!("Failed to poll for device code: {}", e))
                })?;

            match response.status().as_u16() {
                200 => {
                    let response_body: serde_json::Value = response.json().await.map_err(|e| {
                        DeviceCodeError::PollFailed(format!(
                            "Failed to load json body from response: {}",
                            e
                        ))
                    })?;

                    let token: AccessTokenResponse = serde_json::from_value(response_body.clone())
                        .map_err(|_| {
                            DeviceCodeError::PollFailed(format!(
                                "Failed to parse access token response. Recieved {}",
                                response_body
                            ))
                        })?;

                    return Ok(token.into());
                }
                403 => {
                    #[derive(Deserialize)]
                    struct PendingResponse {
                        error: String,
                        error_description: Option<String>,
                    }

                    let PendingResponse {
                        error,
                        error_description,
                    } = response.json().await.map_err(|e| {
                        DeviceCodeError::PollFailed(format!(
                            "Failed to parse pending response: {}",
                            e
                        ))
                    })?;

                    match error.as_str() {
                        "authorization_pending" => {}
                        "slow_down" => {
                            interval += Duration::from_secs(5);
                        }
                        "invalid_grant" => {
                            return Err(DeviceCodeError::PollFailed(format!(
                                "Device code authentication failed: {}",
                                error_description
                                    .unwrap_or_else(|| "error_description undefined".into())
                            )))
                        }
                        e => {
                            return Err(DeviceCodeError::PollFailed(format!(
                                "Unexpected error code in response body: {}",
                                e
                            )))
                        }
                    }

                    sleep(&interval).await
                }
                code => {
                    return Err(DeviceCodeError::PollFailed(format!(
                        "Unexpected response code: {}",
                        code
                    )));
                }
            }
        }
    }

    /// Show a prompt in the terminal and open a browser (if available) with an authentication link for generating
    /// an access token.
    fn prompt_user(polling_info: &PollingInfo) {
        if open::that(&polling_info.verification_uri_complete).is_err() {
            println!("Failed to open web browser. Please manually click the link in the following message.")
        }

        let user_code = &polling_info.user_code;
        let code_len = user_code.len();

        println!();
        println!("### ACTION REQUIRED ###");
        println!();
        println!(
            "Visit {} to complete authentication by following the below steps:",
            polling_info.verification_uri_complete
        );
        println!();
        println!("1. Verify that this code matches the code in your browser");
        println!();
        println!("             +------{}------+", "-".repeat(code_len));
        println!("             |      {}      |", " ".repeat(code_len));
        println!("             |      {}      |", user_code);
        println!("             |      {}      |", " ".repeat(code_len));
        println!("             +------{}------+", "-".repeat(code_len));
        println!();
        println!("2. If the codes match, click on the confirm button in the browser");
        println!();
        println!("Waiting for authentication...");
    }

    async fn acquire_new_token_from_user_prompt(&self) -> Result<AuthToken, DeviceCodeError> {
        let client = reqwest::Client::new();

        let info_response = client
            .post(format!("{}/oauth/device/code", self.idp_info.uri))
            .json(&json!({
                "audience": self.idp_info.audience,
                "client_id": self.idp_info.client_id,
                "scope": self.idp_scopes().join(" ")
            }))
            .send()
            .await
            .map_err(|e| {
                DeviceCodeError::AcquireNewTokenFailed(format!(
                    "Failed to get auth0 device code: {}",
                    e
                ))
            })?;

        if info_response.status() != 200 {
            return Err(DeviceCodeError::AcquireNewTokenFailed(format!(
                "Failed to get auth0 device code: status: {}",
                info_response.status()
            )));
        }

        let polling_info: PollingInfo = info_response.json().await.map_err(|e| {
            DeviceCodeError::AcquireNewTokenFailed(format!(
                "Failed to parse polling info json response: {}",
                e
            ))
        })?;

        Self::prompt_user(&polling_info);

        self.poll_for_access_token(&client, &polling_info).await
    }

    async fn try_acquire_token_from_refresh_token(
        &self,
        cached_token: &AuthToken,
    ) -> Result<Option<AuthToken>, DeviceCodeError> {
        let client = reqwest::Client::new();

        let response = client
            .post(format!("{}/oauth/token", self.idp_info.uri))
            .json(&json!({
                "grant_type": "refresh_token",
                "refresh_token": cached_token.refresh_token,
                "client_id": self.idp_info.client_id,
                "scope": self.idp_scopes().join(" ")
            }))
            .send()
            .await
            .map_err(|e| {
                DeviceCodeError::RefreshTokenFailed(format!(
                    "Failed to redeem refresh token: {}",
                    e
                ))
            })?;

        if response.status() == 200 {
            let response: AccessTokenResponse = response.json().await.map_err(|e| {
                DeviceCodeError::RefreshTokenFailed(format!("Failed to parse json response: {}", e))
            })?;

            Ok(Some(response.into()))
        } else {
            Ok(None)
        }
    }

    fn is_token_expired(token: &AuthToken) -> bool {
        (now_secs() + EXPIRY_GRACE_PERIOD_SECONDS) > token.expiry
    }
}

#[async_trait(?Send)]
impl AsyncAccessor for Auth0DeviceCodeCredentials {
    type Value = DataServiceCredentials;
    type Error = StdError;

    async fn get_value(&self) -> Result<Self::Value, Self::Error> {
        // To make sure that we're not refreshing the credentials twice - take a look at the very
        // beginning of the process. If another `get_value` call has happened this lock
        // will make sure we use the cached value instead of creating a new one.
        let mut cache_guard = self.cached_token.lock().await;

        // Check to see if we can make a new token from the cache
        let new_token_from_cache = if let Some(cached_token) = &mut *cache_guard {
            // If the token hasn't expired yet just return with it immediately
            if !Self::is_token_expired(cached_token) {
                return Ok(DataServiceCredentials::new(&cached_token.access_token));
            }

            // The cached token has expired so try and get a new one from the refresh token.
            // If that can't be done this method will return None.
            self.try_acquire_token_from_refresh_token(cached_token)
                .await?
        } else {
            None
        };

        let new_token = match new_token_from_cache {
            // If a new token was generated from the previous step we're safe to use that. If not
            // we'll need to ask the user to allow us to get a new one.
            Some(t) => t,
            None => self.acquire_new_token_from_user_prompt().await?,
        };

        self.validate_token(&new_token.access_token)?;

        std::fs::write(
            &self.token_path,
            serde_json::to_string_pretty(&new_token).map_err(|e| {
                DeviceCodeError::PersistTokenError(format!(
                    "Failed to encode token as string: {}",
                    e
                ))
            })?,
        )
        .map_err(|e| {
            DeviceCodeError::PersistTokenError(format!(
                "Failed to write access token to disk: {}",
                e
            ))
        })?;

        let creds = DataServiceCredentials::new(&new_token.access_token);

        // Store the new token inside the cache
        cache_guard.replace(new_token);

        Ok(creds)
    }
}

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

    fn create_test_credentials() -> Auth0DeviceCodeCredentials {
        Auth0DeviceCodeCredentials::new(
            "test".into(),
            IdpInfo {
                uri: "test".into(),
                client_id: "test".into(),
                audience: "test".into(),
            },
            Some("test-workspace".into()),
        )
    }

    fn encode_jwt(data: impl Serialize) -> String {
        jsonwebtoken::encode(
            &Default::default(),
            &data,
            &EncodingKey::from_secret("secret".as_bytes()),
        )
        .expect("Failed to encode jwt")
    }

    #[test]
    fn test_fail_when_base_scope_not_included() {
        let creds = create_test_credentials();

        let scopes = [
            "collection.info",
            "collection.list",
            "document.put",
            "document.delete",
            "document.get",
            "document.query",
            "offline_access",
            "ws:test-workspace",
        ]
        .join(" ");

        let encoded = encode_jwt(serde_json::json!({ "scope": scopes }));

        let err = creds
            .validate_token(&encoded)
            .expect_err("Expected validation to fail");

        assert_eq!(
            err.to_string(),
            "InvalidAccessToken: Payload was missing scopes: collection.create, collection.delete"
        );
    }

    #[test]
    fn test_fail_when_workspace_scope_not_included() {
        let creds = create_test_credentials();

        let all_base_scopes = [
            "collection.create",
            "collection.delete",
            "collection.info",
            "collection.list",
            "document.put",
            "document.delete",
            "document.get",
            "document.query",
            "offline_access",
        ]
        .join(" ");

        let encoded = encode_jwt(serde_json::json!({ "scope": all_base_scopes }));

        let err = creds
            .validate_token(&encoded)
            .expect_err("Expected validation to fail");

        assert_eq!(
            err.to_string(),
            "InvalidAccessToken: Payload was missing scopes: ws:test-workspace"
        );
    }

    #[test]
    fn test_pass_for_base_scopes_when_workspace_is_none() {
        let mut creds = create_test_credentials();
        creds.workspace = None;

        let all_base_scopes = [
            "collection.create",
            "collection.delete",
            "collection.info",
            "collection.list",
            "document.put",
            "document.delete",
            "document.get",
            "document.query",
            "offline_access",
        ]
        .join(" ");

        let encoded = encode_jwt(serde_json::json!({ "scope": all_base_scopes }));

        creds
            .validate_token(&encoded)
            .expect("Expected validation to pass");
    }

    #[test]
    fn test_fail_with_no_kms_claim() {
        let creds = create_test_credentials();

        let all_scopes = [
            "collection.create",
            "collection.delete",
            "collection.info",
            "collection.list",
            "document.put",
            "document.delete",
            "document.get",
            "document.query",
            "offline_access",
            "ws:test-workspace",
        ]
        .join(" ");

        let encoded = encode_jwt(serde_json::json!({ "scope": all_scopes }));

        let err = creds
            .validate_token(&encoded)
            .expect_err("Expected validation to fail");

        assert_eq!(
            err.to_string(),
            "InvalidAccessToken: Token does not contain KMS federation claim"
        );
    }

    #[test]
    fn test_pass_when_everything_is_set() {
        let creds = create_test_credentials();

        let all_scopes = [
            "collection.create",
            "collection.delete",
            "collection.info",
            "collection.list",
            "document.put",
            "document.delete",
            "document.get",
            "document.query",
            "offline_access",
            "ws:test-workspace",
        ]
        .join(" ");

        let encoded = encode_jwt(serde_json::json!({
            "scope": all_scopes,
            "https://aws.amazon.com/tags": {
                "principal_tags": {
                    "workspace": [ "ws:test-workspace" ]
                }
            }
        }));

        creds
            .validate_token(&encoded)
            .expect("Expected validation to pass");
    }
}