use std::sync::Arc;

use anyhow::Result;
use log::*;
use serde::{Deserialize, Serialize};
use sqlx::{mysql::MySqlConnectOptions, MySqlPool};

pub const JWT_SPECIFIED_USER: &str = "cloud_ai";

#[derive(Debug, Serialize, Deserialize)]
pub struct JwtUser {
    username: String,
    org_id: i64,
    cluster_id: String,
    token_expired_seconds: u32,
}

impl Default for JwtUser {
    fn default() -> Self {
        let org_id = std::env::var("JWT_ORG_ID")
            .unwrap_or("2130392".into())
            .parse()
            .expect("JWT_ORG_ID is not a number");
        let cluster_id = std::env::var("JWT_CLUSTER_ID").unwrap_or("10651690106009365685".into());
        JwtUser {
            username: JWT_SPECIFIED_USER.into(),
            org_id,
            cluster_id,
            token_expired_seconds: 900,
        }
    }
}

impl JwtUser {
    pub fn new(tenant_token: String) -> Self {
        JwtUser {
            username: format!("{}.{}", tenant_token, JWT_SPECIFIED_USER),
            ..Default::default()
        }
    }
}

#[derive(Default, Debug, Serialize, Deserialize)]
pub struct JwtRequest {
    users: Vec<JwtUser>,
}

impl JwtRequest {
    pub fn add_tenant(mut self, tenant_token: String) -> Self {
        self.users.push(JwtUser::new(tenant_token));
        self
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct JwtUserAuth {
    username: String,
    jwt: String,
    token_expired_seconds: u32,
    org_id: String,
    cluster_id: String,
    success: bool,
    message: String,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct JwtBaseResp {
    // tags: HashMap<>,
    err_code: String,
    err_msg: String,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct JwtResponse {
    jwts: Vec<JwtUserAuth>,
    base_resp: JwtBaseResp,
}

impl JwtResponse {
    pub fn users(&self) -> Result<Vec<JwtUserAuth>> {
        if self.base_resp.err_code == "0" {
            Ok(self.jwts.clone())
        } else {
            Err(anyhow::anyhow!(
                "JWT request failed: {}",
                self.base_resp.err_msg
            ))
        }
    }

    pub fn get_jwt(&self, tenant_token: &str) -> Result<String> {
        let mut msg = String::new();
        let jwt = self
            .users()?
            .iter()
            .find(|u| {
                if !u.success {
                    error!("JWT request failed: {:?}", u);
                    msg.push_str(&u.message);
                }
                u.success && u.username == format!("{}.{}", tenant_token, JWT_SPECIFIED_USER)
            })
            .map(|u| u.jwt.clone());
        if let Some(jwt) = jwt {
            Ok(jwt)
        } else {
            Err(anyhow::anyhow!(
                "JWT not found for tenant {}, error: {}",
                tenant_token,
                msg
            ))
        }
    }
}

pub(crate) struct JwtClient {
    host: String,
    path: String,
    auth_token: String,
}

impl Default for JwtClient {
    fn default() -> Self {
        JwtClient::new_from_env()
    }
}

impl JwtClient {
    pub fn new(host: String, path: String, auth_token: String) -> Self {
        JwtClient {
            host,
            path,
            auth_token,
        }
    }

    fn new_from_env() -> Self {
        let mut host = std::env::var("JWT_HOST").expect("JWT_HOST is not set");
        if host.starts_with("http://") {
            host = host.replace("http://", "https://");
        }
        if !host.starts_with("https://") {
            host = format!("https://{}", host);
        }
        let path = std::env::var("JWT_PATH")
            .unwrap_or("/inner/account-central/api/v1/builtin-user-jwt".to_string());
        let auth_token = std::env::var("JWT_AUTH_TOKEN").expect("JWT_AUTH_TOKEN is not set");
        Self::new(host, path, auth_token)
    }

    pub async fn submit(&self, req: &JwtRequest) -> Result<JwtResponse> {
        let url = format!("{}{}", self.host, self.path);
        let client = reqwest::Client::new();
        let resp = client
            .post(&url)
            .header("Authorization", format!("basic {}", self.auth_token))
            .json(req)
            .send()
            .await?;
        debug!("JWT response: {:?} for request: {:?}", resp, req);
        match resp.json().await {
            Ok(jwt_resp) => Ok(jwt_resp),
            Err(e) => Err(anyhow::anyhow!("Failed to parse JWT response: {}", e)),
        }
    }
}

fn localhost() -> bool {
    std::env::var("LOCALHOST").unwrap_or("false".to_string()) == "true"
}

pub(crate) async fn get_db_connection_by_tenant_token(
    meta_db: Arc<MySqlPool>,
    tenant_token: &str,
) -> Result<Arc<MySqlPool>> {
    if localhost() {
        warn!("Using localhost db");
        return Ok(meta_db.clone());
    }

    let jwt_request = JwtRequest::default().add_tenant(tenant_token.into());
    let jwt_response = JwtClient::default().submit(&jwt_request).await?;
    let dsn = format!(
        "mysql://{}.{}:{}@{}:{}/{}",
        tenant_token,
        JWT_SPECIFIED_USER,
        jwt_response.get_jwt(tenant_token)?,
        meta_db.connect_options().get_host(),
        meta_db.connect_options().get_port(),
        meta_db.connect_options().get_database().unwrap_or("ai"),
    );
    info!("Connecting to db with tenant: {}", tenant_token);
    let opts = dsn
        .parse::<MySqlConnectOptions>()?
        .enable_cleartext_plugin(true);

    let mut retries = std::env::var("TENANT_DB_CONNECT_RETRIES")
        .unwrap_or_else(|_| "10".to_string())
        .parse::<u32>()
        .unwrap_or(10);

    let pool = loop {
        match MySqlPool::connect_with(opts.clone()).await {
            Ok(pool) => break pool,
            Err(e) if retries > 0 => {
                retries -= 1;
                warn!(
                    "Failed to connect to db: {:?} with: {}, retrying... ({}/3)",
                    e,
                    dsn,
                    3 - retries
                );
                tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
            }
            Err(e) => return Err(e.into()),
        }
    };
    Ok(Arc::new(pool))
}

// #[cfg(test)]
// mod tests {
//     use super::*;
//     #[tokio::test]
//     async fn test_jwt_client_submit() {
//         std::env::set_var("JWT_HOST", "https://inner-central.dev.tidbcloud.com");
//         let tenant_token = "tenant123";
//         let req = JwtRequest::default().add_tenant(tenant_token.into());
//         let client = JwtClient::default();
//         let result = client.submit(&req).await.unwrap();
//         assert!(result.get_jwt(tenant_token).is_ok());
//     }
// }