//! Kubernetes Integration
use std::collections::btree_map::BTreeMap;

use chrono::offset::Utc;
use chrono::{DateTime, Duration};
use k8s_openapi::api::core::v1::{ConfigMap, Secret};
use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
use kube::core::params::{ListParams, Patch, PatchParams, PostParams};
use kube::{Api, Client};
use log::{debug, error, info, warn};
use rand::seq::SliceRandom;

use crate::model::k8s::*;

/// The set of characters allowed to appear in bootstrap tokens
const TOKEN_CHARS: [char; 36] = [
    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
    'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '0', '1', '2', '3',
    '4', '5', '6', '7', '8', '9',
];

/// Kubernetes Bootstrap Token
///
/// Bootstrap tokens are typically used to add new nodes to the cluster using
/// `kubeadm join`.  They are bearer tokens consisting of two parts: a token ID
/// and a secret.  For additional information, see [Authenticating with
/// Bootstrap Tokens][0].
///
/// The Dynk8s Provisioner allocates bootstrap tokens for ephemeral nodes.
/// Each token is assigned to a specific EC2 instance; the instance ID is
/// stored in the token Secret's metadata, using the
/// `dynk8s.du5t1n.me/ec2-instance-id` label.
///
/// [0]: https://kubernetes.io/docs/reference/access-authn-authz/bootstrap-tokens/
#[derive(Clone, Debug)]
struct BootstrapToken {
    /// The token ID (generated)
    token_id: String,
    /// The token secret (generated)
    secret: String,
    /// The date and time the token expires
    expiration: DateTime<Utc>,
    /// The EC2 instance ID to which the token is assigned
    instance_id: Option<String>,
}

#[allow(dead_code)]
impl BootstrapToken {
    /// Generate a new token
    ///
    /// Initially, the token is *not* assigned to an EC2 instance.  Use
    /// [`Self::instance_id`] to set the associated instance ID.
    pub fn new() -> Self {
        let mut rng = rand::thread_rng();
        let mut token_id = String::with_capacity(6);
        while token_id.len() < 6 {
            token_id.push(TOKEN_CHARS.choose(&mut rng).unwrap().clone());
        }
        let mut secret = String::with_capacity(16);
        while secret.len() < 16 {
            secret.push(TOKEN_CHARS.choose(&mut rng).unwrap().clone());
        }
        let expiration = Utc::now() + Duration::hours(1);
        Self {
            token_id,
            secret,
            expiration,
            instance_id: None,
        }
    }

    /// Return the expiration date/time
    pub fn expiration(&self) -> &DateTime<Utc> {
        &self.expiration
    }

    /// Set the ID of the EC2 instance associated with the token
    pub fn instance_id(mut self, instance_id: String) -> Self {
        self.instance_id = Some(instance_id);
        self
    }

    /// Return the secret part of the token
    pub fn secret(&self) -> &str {
        &self.secret
    }

    /// Set the expiration date/time from a time-to-live duration
    pub fn set_ttl(mut self, ttl: Duration) -> Self {
        self.expiration = Utc::now() + ttl;
        self
    }

    /// Set the expiration date/time
    pub fn set_expiration(mut self, expiration: DateTime<Utc>) -> Self {
        self.expiration = expiration;
        self
    }

    /// Return the ID part of the token
    pub fn token_id(&self) -> &str {
        &self.token_id
    }

    /// Return the token as a string
    pub fn token(&self) -> String {
        format!("{}.{}", self.token_id, self.secret)
    }
}

impl Into<Secret> for BootstrapToken {
    /// Create a [`Secret`] for the token
    ///
    /// Converting a [`BootstrapToken`] into a [`Secret`] populates the fields
    /// necessary to store the token in Kubernetes.  The `Secret` can be passed
    /// to e.g. [`kube::Api::create`].
    fn into(self) -> Secret {
        let mut data = BTreeMap::<String, String>::new();
        data.insert("token-id".into(), self.token_id.clone());
        data.insert("token-secret".into(), self.secret);
        data.insert("expiration".into(), self.expiration.to_rfc3339());
        data.insert("usage-bootstrap-authentication".into(), "true".into());
        data.insert("usage-bootstrap-signing".into(), "true".into());
        data.insert(
            "auth-extra-groups".into(),
            "system:bootstrappers:kubeadm:default-node-token".into(),
        );
        let mut labels = BTreeMap::<String, String>::new();
        if let Some(instance_id) = self.instance_id {
            labels.insert(
                "dynk8s.du5t1n.me/ec2-instance-id".into(),
                instance_id.into(),
            );
        }
        Secret {
            data: None,
            immutable: None,
            metadata: ObjectMeta {
                annotations: None,
                cluster_name: None,
                creation_timestamp: None,
                deletion_grace_period_seconds: None,
                deletion_timestamp: None,
                finalizers: None,
                generate_name: None,
                generation: None,
                labels: Some(labels),
                managed_fields: None,
                name: Some(format!("bootstrap-token-{}", self.token_id)),
                namespace: None,
                owner_references: None,
                resource_version: None,
                self_link: None,
                uid: None,
            },
            string_data: Some(data),
            type_: Some("bootstrap.kubernetes.io/token".into()),
        }
    }
}

/// Assign an existing WireGuard configuration to the specified EC2 instance
///
/// This function finds the first unused WireGuard client configuration, stored
/// as a Kubernetes Secret resource, and assigns it to the specified EC2
/// instance.  Configs are assigned by setting the
/// `dynk8s.du5t1n.me/ec2-instance-id` label in the Secret resource's metadata.
///
/// Secret resources for WireGuard configuration have a *type* of
/// `dynk8s.du5t1n.me/wireguard-config`.  The Secret's `data` field must
/// contain a `wireguard-config` property, which contains the WireGuard client
/// configuration the node should use.  Configs must be created ahead of time
/// and must refer to working keys already configured on the WireGuard server.
pub async fn assign_wireguard_config(
    instance_id: &str,
) -> Result<(), kube::Error> {
    let client = Client::try_default().await?;
    let secrets: Api<Secret> = Api::default_namespaced(client);
    debug!(
        "Checking for WireGuard configs already assigned to instance {}",
        instance_id
    );
    let lp = ListParams::default()
        .fields("type=dynk8s.du5t1n.me/wireguard-config")
        .labels(&format!("dynk8s.du5t1n.me/ec2-instance-id={}", instance_id));
    let res = secrets.list(&lp).await?;
    if !res.items.is_empty() {
        info!(
            "WireGuard config already assigned to instance {}",
            instance_id
        );
        return Ok(());
    }
    debug!("Looking for available WireGuard configs");
    let lp = ListParams::default()
        .fields("type=dynk8s.du5t1n.me/wireguard-config")
        .labels("dynk8s.du5t1n.me/ec2-instance-id=");
    let res = secrets.list(&lp).await?;
    if res.items.is_empty() {
        error!(
            "No WireGuard config available for instance {}",
            &instance_id
        );
    } else {
        if let Some(name) = &res.items[0].metadata.name {
            let mut labels = BTreeMap::<String, String>::new();
            labels.insert(
                "dynk8s.du5t1n.me/ec2-instance-id".into(),
                instance_id.into(),
            );
            let mut secret = Secret::default();
            secret.metadata.labels = Some(labels);
            let patch = Patch::Apply(&secret);
            let pp = PatchParams::apply(env!("CARGO_PKG_NAME")).force();
            secrets.patch(&name, &pp, &patch).await?;
            info!(
                "Assigned WireGuard config {} to instance {}",
                name, &instance_id
            );
        }
    }
    Ok(())
}

/// Unassign all WireGuard configs from the specified EC2 instance
///
/// This function finds all WireGuard configs, stored as Kubernetes Secret
/// resources, associated with the specified EC2 instance and unassigns them.
/// Unassigned configs have the `dynk8s.du5t1n.me/ec2-instance-id` label set to
/// the empty string.
pub async fn unassign_wireguard_config(
    instance_id: &str,
) -> Result<(), kube::Error> {
    let client = Client::try_default().await?;
    let secrets: Api<Secret> = Api::default_namespaced(client);
    let lp = ListParams::default()
        .fields("type=dynk8s.du5t1n.me/wireguard-config")
        .labels(&format!("dynk8s.du5t1n.me/ec2-instance-id={}", instance_id));
    info!(
        "Unassigning WireGuard configs from instance {}",
        instance_id
    );
    for s in secrets.list(&lp).await? {
        if let Some(name) = &s.metadata.name {
            let mut labels = BTreeMap::<String, String>::new();
            labels
                .insert("dynk8s.du5t1n.me/ec2-instance-id".into(), "".into());
            let mut secret = Secret::default();
            secret.metadata.labels = Some(labels);
            let patch = Patch::Apply(&secret);
            let pp = PatchParams::apply(env!("CARGO_PKG_NAME")).force();
            secrets.patch(&name, &pp, &patch).await?;
            info!(
                "Unassigned WireGuard config {} from instance {}",
                name, &instance_id
            );
        }
    }
    Ok(())
}

/// Retrieve the WireGuard config assigned to the specified EC2 instance
///
/// This function finds the first WireGuard client configuration, stored as a
/// Kubernetes Secret resource, associated with the specified EC2 instance.
///
/// If multiple WireGuard configs are assigned to an EC2 instance, only the
/// first one returned by the Kubernetes list query is returned.
pub async fn get_wireguard_config(
    instance_id: &str,
) -> Result<Option<String>, kube::Error> {
    let client = Client::try_default().await?;
    let secrets: Api<Secret> = Api::default_namespaced(client);
    let lp = ListParams::default()
        .fields("type=dynk8s.du5t1n.me/wireguard-config")
        .labels(&format!("dynk8s.du5t1n.me/ec2-instance-id={}", instance_id));
    for s in secrets.list(&lp).await? {
        if let Some(data) = s.data {
            match data.get("wireguard-config") {
                Some(s) => match String::from_utf8(s.0.clone()) {
                    Ok(s) => return Ok(Some(s)),
                    Err(e) => {
                        error!("Invalid WireGuard configuration: {}", e);
                    }
                },
                None => {
                    error!(concat!(
                        "Invalid WireGuard configuration: ",
                        "missing wireguard-config property"
                    ));
                }
            };
        }
    }
    Ok(None)
}

/// Generate and store a bootstrap token for the specified EC2 instance
///
/// This function generates a new bootstrap token and stores it as a Kubernetes
/// Secret resource.  The token is assigned to the given EC2 instance, and will
/// be provided to that instance when it is ready to join the cluster.
pub async fn create_bootstrap_token<I: AsRef<str>>(
    instance_id: I,
) -> Result<(), kube::Error> {
    let instance_id = instance_id.as_ref();
    info!("Creating bootstrap token for instance {}", instance_id);
    let token = BootstrapToken::new().instance_id(instance_id.into());
    let client = Client::try_default().await?;
    let secrets: Api<Secret> = Api::namespaced(client, "kube-system");
    let pp: PostParams = Default::default();
    let secret = secrets.create(&pp, &token.into()).await?;
    info!("Successfully created secret {:?}", &secret.metadata.name);
    Ok(())
}

/// Delete bootstrap tokens associated withthe specified EC2 instance
pub async fn delete_bootstrap_tokens<I: AsRef<str>>(
    instance_id: I,
) -> Result<(), kube::Error> {
    let instance_id = instance_id.as_ref();
    info!("Deleting bootstrap tokens for instance {}", &instance_id);
    let client = Client::try_default().await?;
    let secrets: Api<Secret> = Api::namespaced(client, "kube-system");
    let lp = ListParams::default()
        .fields("type=bootstrap.kubernetes.io/token")
        .labels(&format!("dynk8s.du5t1n.me/ec2-instance-id={}", &instance_id));
    secrets.delete_collection(&Default::default(), &lp).await?;
    Ok(())
}

/// Get the `kubeadm join` configuration for the specified EC2 instance
///
/// This function creates a kubeconfig file that can be passed to `kubeadm
/// join` to add the specified EC2 instance to the Kubernetes cluster as a
/// worker node.  The cluster configuration is read from the `cluster-info`
/// ConfigMap in the *kube-public* namespace.  The bootstrap token assigned to
/// the instance is included for client authentication.
pub async fn get_kubeconfig<I: AsRef<str>>(
    instance_id: I,
) -> Result<Option<KubeConfig>, kube::Error> {
    let instance_id = instance_id.as_ref();
    let token = match get_bootstrap_token(&instance_id).await {
        Ok(Some(t)) => t,
        Ok(None) => {
            warn!("No bootstrap token assigned to instance {}", &instance_id);
            return Ok(None);
        }
        Err(e) => {
            error!(
                "Could not get bootstrap token for instance {}: {}",
                &instance_id, e
            );
            return Ok(None);
        }
    };
    match get_cluster_info().await? {
        Some(config) => {
            let cluster = Cluster {
                name: "kubernetes".into(),
                cluster: config.clusters[0].cluster.clone(),
            };
            let context = Context {
                name: "kubeadm".into(),
                context: ContextInfo {
                    cluster: "kubernetes".into(),
                    user: "kubeadm".into(),
                },
            };
            let user = User {
                name: "kubeadm".into(),
                user: UserInfo { token: token },
            };
            let mut kubeconfig = KubeConfig::default();
            kubeconfig.clusters = vec![cluster];
            kubeconfig.contexts = Some(vec![context]);
            kubeconfig.current_context = "kubeadm".into();
            kubeconfig.users = Some(vec![user]);
            Ok(Some(kubeconfig))
        }
        None => {
            warn!("No kubeconfig loaded from cluster-info");
            Ok(None)
        }
    }
}

/// Retrieve the bootstrap token assigned to an EC2 instance
async fn get_bootstrap_token<I: AsRef<str>>(
    instance_id: I,
) -> Result<Option<String>, kube::Error> {
    let instance_id = instance_id.as_ref();
    let client = Client::try_default().await?;
    let secrets: Api<Secret> = Api::namespaced(client, "kube-system");
    let lp = ListParams::default()
        .fields("type=bootstrap.kubernetes.io/token")
        .labels(&format!(
            "dynk8s.du5t1n.me/ec2-instance-id={}",
            &instance_id
        ));
    for s in secrets.list(&lp).await? {
        match token_string(&s) {
            Ok(t) => return Ok(Some(t)),
            Err(e) => {
                error!("Invalid bootstrap token: {}", e);
            }
        }
    }
    Ok(None)
}

/// Get cluster information from the ConfigMap
async fn get_cluster_info() -> Result<Option<KubeConfig>, kube::Error> {
    let client = Client::try_default().await?;
    let configmaps: Api<ConfigMap> = Api::namespaced(client, "kube-public");
    let cluster_info = configmaps.get("cluster-info").await?;
    if let Some(data) = cluster_info.data {
        if let Some(config) = data.get("kubeconfig") {
            match serde_yaml::from_str::<KubeConfig>(config) {
                Ok(c) => return Ok(Some(c)),
                Err(e) => {
                    error!(
                        "Could not load kubeconfig from cluster-info: {}",
                        e
                    );
                }
            };
        } else {
            error!("No kubeconfig property found in cluster-info ConfigMap");
        }
    } else {
        error!("No data property found in cluster-info ConfigMap");
    }
    Ok(None)
}

/// Get the string representation of a bootstrap token from a Secret
fn token_string(secret: &Secret) -> Result<String, String> {
    let data = match &secret.data {
        Some(d) => d,
        None => return Err("Missing data property".into()),
    };
    let token_id = match data.get("token-id") {
        Some(s) => match String::from_utf8(s.0.clone()) {
            Ok(s) => s,
            Err(e) => return Err(e.to_string()),
        },
        None => return Err("Missing token-id".into()),
    };
    let secret = match data.get("token-secret") {
        Some(s) => match String::from_utf8(s.0.clone()) {
            Ok(s) => s,
            Err(e) => return Err(e.to_string()),
        },
        None => return Err("Missing token-secret".into()),
    };
    Ok(format!("{}.{}", token_id, secret))
}

#[cfg(test)]
mod test {
    use super::*;
    use k8s_openapi::ByteString;
    use regex::Regex;

    #[test]
    fn test_bootstrap_token_new() {
        let token = BootstrapToken::new();
        let id_re = Regex::new(r"^[a-z0-9]{6}$").unwrap();
        let secret_re = Regex::new(r"^[a-z0-9]{16}$").unwrap();
        let token_re = Regex::new(r"[a-z0-9]{6}\.[a-z0-9]{16}$").unwrap();
        assert!(id_re.is_match(&token.token_id()));
        assert!(secret_re.is_match(&token.secret()));
        assert!(token_re.is_match(&token.token()));
    }

    #[test]
    fn test_bootstrap_token_into_secret() {
        let token = BootstrapToken::new();
        let secret: Secret = token.clone().into();
        let data = secret.string_data.unwrap();
        assert_eq!(
            &secret.metadata.name,
            &Some(format!("bootstrap-token-{}", token.token_id))
        );
        assert_eq!(data.get("token-id").unwrap(), &token.token_id);
        assert_eq!(data.get("token-secret").unwrap(), &token.secret);
        assert_eq!(
            data.get("expiration").unwrap(),
            &token.expiration.to_rfc3339()
        );
        assert!(secret
            .metadata
            .labels
            .unwrap()
            .get("dynk8s.du5t1n.me/ec2-instance-id")
            .is_none());
    }

    #[test]
    fn test_bootstrap_token_into_secret_instance_id() {
        let token =
            BootstrapToken::new().instance_id("i-0a1b2c3d4e5f6f7f8".into());
        let secret: Secret = token.clone().into();
        let data = secret.string_data.unwrap();
        assert_eq!(
            &secret.metadata.name,
            &Some(format!("bootstrap-token-{}", token.token_id))
        );
        assert_eq!(data.get("token-id").unwrap(), &token.token_id);
        assert_eq!(data.get("token-secret").unwrap(), &token.secret);
        assert_eq!(
            data.get("expiration").unwrap(),
            &token.expiration.to_rfc3339()
        );
        assert_eq!(
            secret
                .metadata
                .labels
                .unwrap()
                .get("dynk8s.du5t1n.me/ec2-instance-id")
                .unwrap(),
            "i-0a1b2c3d4e5f6f7f8"
        );
    }

    #[test]
    fn test_token_string() {
        let token = BootstrapToken::new();
        let mut data = BTreeMap::new();
        data.insert("token-id".into(), ByteString(token.token_id().into()));
        data.insert("token-secret".into(), ByteString(token.secret().into()));
        let secret = Secret {
            data: Some(data),
            ..Default::default()
        };
        assert_eq!(token.token(), token_string(&secret).unwrap());
    }
}