mc-ingress/src/kube_cache.rs

use std::{collections::HashMap, fmt, sync::Arc, time::Duration};

use async_trait::async_trait;
use futures::StreamExt;
use itertools::Itertools;
use k8s_openapi::api::{apps::v1::Deployment, core::v1::Service};
use kube::{
    api::{ListParams, ObjectList, Patch, PatchParams},
    runtime::{reflector::Lookup, WatchStreamExt},
    Api, Client, ResourceExt,
};
use serde_json::json;
use tokio::task::JoinHandle;

use crate::{
    mc_server::{sanitize_addr, MinecraftAPI, MinecraftServerHandle, ServerDeploymentStatus},
    OpaqueError,
};
const MAIN_LABEL: &str = "tami.moe/minecraft";
const PORT_LABEL: &str = "tami.moe/minecraft-port";
const URL_ANNOTATION: &str = "tami.moe/minecraft-url";

/// This is the layer who is respinsible for caching requests.

// A heads-up:
// It should be also clone-able freely, because it deals with
// the underlying async data access.
#[derive(Debug, Clone)]
pub struct KubeCache {
    // When the store gets copied, it points to the same backing store
    // In other words, this satisfies my restriction for cloning and
    // there is need no to wrap it in Arc<>
    dep_cache: kube::runtime::reflector::Store<Deployment>,
    dep_api: Api<Deployment>,
    srv_api: Api<Service>,
    in_cluster: bool,
}
impl KubeCache {
    /// This initializes the creation of a "kubernetes client"
    /// and if it is not possible returns a None.
    ///
    /// It also returns a JoinHandle which is tied to the api watcher
    /// and if it returns that means that the kubecache is not responsive now.
    ///
    /// # Example:
    /// ```no_run
    /// let (kube_cache, api_task) = kube_cache::KubeCache::create().await.unwrap();
    /// ```
    pub async fn create() -> Option<(KubeCache, JoinHandle<()>)> {
        let in_cluster = match std::env::var("KUBERNETES_SERVICE_HOST") {
            Ok(_) => true,
            Err(_) => false,
        };
        let client = Client::try_default().await.unwrap();

        let dep_api: Api<Deployment> = Api::default_namespaced(client.clone());
        let lp = kube::runtime::watcher::Config::default().labels(MAIN_LABEL);
        let (dep_reader, dep_writer) = kube::runtime::reflector::store();
        let dep_watcher_rf =
            kube::runtime::reflector(dep_writer, kube::runtime::watcher(dep_api.clone(), lp));

        let handle = tokio::spawn(async move {
            // Source: https://docs.rs/kube/latest/kube/runtime/fn.reflector.html
            let infinite_watch = dep_watcher_rf
                .applied_objects()
                .for_each(|_o| std::future::ready(()));
            let _res = infinite_watch.await;
            tracing::error!(
                "deployments watcher ended; this should not happen; (program should exit now)"
            );
        });

        let srv_api: Api<Service> = Api::default_namespaced(client);

        return Some((
            KubeCache {
                dep_cache: dep_reader,
                dep_api: dep_api,
                srv_api,
                in_cluster,
            },
            handle,
        ));
    }
    fn get_dep(&self, name: &str) -> Option<Arc<Deployment>> {
        // dep.name_any() has the caveat of returning "" if it can't query for
        // a name, resulting in unexpected behaviour if that happens and also
        // a query for "" comes as well.
        self.dep_cache.find(|x| x.name_any() == name)
    }
    async fn get_srv(&self, name: &str) -> Result<Service, kube::Error> {
        self.srv_api.get(name).await
    }
    async fn list_srvs(&self) -> ObjectList<Service> {
        // let lp: ListParams = ListParams::default();
        let lp: ListParams = ListParams::default().labels(MAIN_LABEL);
        self.srv_api.list(&lp).await.unwrap()
    }

    pub async fn query_dep(&self, addr: &str, port: &str) -> Option<Arc<Deployment>> {
        self.dep_cache.find(|x| filter_label_value(x, addr, port))
    }

    pub async fn query_srv(&self, addr: &str, port: &str) -> Option<Service> {
        let deploys = self.list_srvs().await;
        deploys.into_iter().find(|x| {
            let in_cluster = match x.spec.as_ref().unwrap().type_.as_ref() {
                Some(t) => t == "ClusterIP",
                None => false,
            };
            let incorrect_type = in_cluster ^ self.in_cluster;
            !incorrect_type && filter_label_value(x, addr, port)
        })
    }

    async fn set_dep_scale(&self, name: &str, num: i32) -> Result<Deployment, kube::Error> {
        let patch = Patch::Merge(json!({"spec":{"replicas": num}}));
        let pp = PatchParams::default();
        self.dep_api.patch(name, &pp, &patch).await
    }
}

#[derive(Clone)]
pub struct McApi {
    cache: KubeCache,
    map: Arc<tokio::sync::Mutex<HashMap<String, JoinHandle<()>>>>,
}

#[async_trait]
impl MinecraftAPI<Server> for McApi {
    #[tracing::instrument(
        name = "MinecraftAPI::query_server",
        level = "info",
        skip(self, addr, port)
    )]
    async fn query_server(&self, addr: &str, port: &str) -> Result<Server, OpaqueError> {
        let addr = sanitize_addr(&addr);

        let deployment = match self.cache.query_dep(&addr, &port).await {
            Some(x) => x,
            None => {
                return Err(OpaqueError::create_with_kind(
                    "Failed to find deployment name by addr",
                    "DepNameLookupFailed",
                ))
            }
        };
        let service = match self.cache.query_srv(&addr, &port).await {
            Some(x) => x,
            None => {
                return Err(OpaqueError::create_with_kind(
                    "Failed to find service by addr and port",
                    "SrvNameLookupFailed",
                ))
            }
        };

        let internal_id = match deployment.metadata.name.clone() {
            Some(x) => x,
            None => {
                return Err(OpaqueError::create_with_kind(
                    "Could not query deployment metadata for name",
                    "InternalIdQueryFailed",
                ));
            }
        };

        tracing::debug!("found kubernetes deployment & service");

        let service_port_spec = service.clone().spec.unwrap().ports.unwrap();
        let port_srv = service_port_spec
            .iter()
            .find(|x| x.name.clone().unwrap() == "mc-router")
            .ok_or(OpaqueError::create(
                "Could not find \"mc-router\" nodePort for server",
            ))?;
        let inter_addr = match self.cache.in_cluster {
            false => {
                let node_port = port_srv
                    .node_port
                    .map(|x| x.to_string())
                    .ok_or(OpaqueError::create("Could not map nodePort to port string"))?;
                format!("localhost:{}", node_port)
            }
            true => {
                let target_port = port_srv.port;
                format!(
                    "{}.default.svc.cluster.local:{}",
                    service
                        .name()
                        .ok_or("Could not get name of ClusterIP service")?,
                    target_port
                )
            }
        };
        return Ok(Server {
            server_addr: addr.to_string(),
            server_port: port.to_string(),
            internal_address: inter_addr,
            internal_id,
            cache: self.cache.clone(),
        });
    }

    fn get_map(&self) -> Arc<tokio::sync::Mutex<HashMap<String, tokio::task::JoinHandle<()>>>> {
        self.map.clone()
    }
}

impl McApi {
    pub async fn create() -> Option<(Self, JoinHandle<()>)> {
        let (kube_cache, kube_task) = KubeCache::create().await?;
        Some((
            Self {
                cache: kube_cache,
                map: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
            },
            kube_task,
        ))
    }
}

#[derive(Clone)]
pub struct Server {
    server_addr: String,
    server_port: String,
    internal_address: String,
    internal_id: String,
    cache: KubeCache,
}
impl fmt::Debug for Server {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("KubeServer")
            .field("internal_id", &self.internal_id)
            .field(
                "join_addr",
                &format!("{}:{}", self.server_addr, self.server_port),
            )
            .finish()
    }
}
#[async_trait::async_trait]
impl MinecraftServerHandle for Server {
    async fn start(&self) -> Result<(), OpaqueError> {
        self.set_scale(1).await.map_err(|e| {
            OpaqueError::create(&format!("failed to set deployment scale: err = {:?}", e))
        })
    }

    async fn stop(&self) -> Result<(), OpaqueError> {
        self.set_scale(0).await.map_err(|e| {
            OpaqueError::create(&format!("failed to set deployment scale: err = {:?}", e))
        })
    }

    #[tracing::instrument(level = "info")]
    async fn query_status(&self) -> Result<crate::mc_server::ServerDeploymentStatus, OpaqueError> {
        let dep = self
            .cache
            .get_dep(&self.internal_id)
            .ok_or_else(|| OpaqueError::create("Failed to get deployment from cache"))?;

        let status = match &dep.status {
            Some(x) => x,
            None => {
                return Err(OpaqueError::create(
                    "failed to get status of deployment for checking replicas",
                ))
            }
        };
        let total_replicas = status
            .replicas
            .unwrap_or_else(|| {
                tracing::trace!("total_replicas failed to get");
                -1
            })
            .clone();
        let available_replicas = status
            .available_replicas
            .unwrap_or_else(|| {
                tracing::trace!("available_replicas failed to get");
                -1
            })
            .clone();
        let ready_replicas = status
            .ready_replicas
            .unwrap_or_else(|| {
                tracing::trace!("ready_replicas failed to get");
                -1
            })
            .clone();
        tracing::debug!("total_replicas: {total_replicas} available_replicas: {available_replicas} ready_replicas : {ready_replicas }");

        if total_replicas > 0 {
            if ready_replicas > 0 {
                return match self.query_server_connectable().await {
                    Ok(x) => Ok(ServerDeploymentStatus::Connectable(x)),
                    Err(_) => Ok(ServerDeploymentStatus::PodOk),
                };
            }
            return Ok(ServerDeploymentStatus::Starting);
        } else {
            return Ok(ServerDeploymentStatus::Offline);
        }
    }

    fn get_internal_address(&self) -> &str {
        self.internal_address.as_str()
    }
    fn get_internal_id(&self) -> &str {
        self.internal_id.as_str()
    }

    fn get_join_addr(&self) -> &str {
        self.server_addr.as_ref()
    }

    fn get_join_port(&self) -> &str {
        self.server_port.as_ref()
    }

    fn get_motd(&self) -> Option<String> {
        let dep = self.cache.get_dep(&self.internal_id)?;

        let all_container_motds = dep
            .spec
            .clone()?
            .template
            .spec?
            .containers
            .iter()
            .map(|cont| match cont.env.clone() {
                Some(es) => es
                    .iter()
                    .filter(|e| e.name.as_str() == "MOTD")
                    .map(|x| x.value.clone())
                    .collect::<Vec<Option<String>>>()
                    .first()?
                    .clone(),
                None => None,
            })
            .collect::<Vec<Option<String>>>();
        all_container_motds.first()?.clone()
    }
}

impl Server {
    async fn set_scale(&self, num: i32) -> Result<(), kube::Error> {
        let _res = self.cache.set_dep_scale(&self.internal_id, num).await?;
        tracing::info!(
            "scaled replicas of {}:{} to {num}",
            self.server_addr,
            self.server_port
        );
        Ok(())
    }
}

fn filter_label_value<R>(res: &R, addr: &str, port: &str) -> bool
where
    R: ResourceExt,
{
    let labels_iter = res.labels().iter();
    let annotations_iter = res.annotations().iter();
    let iters = labels_iter.merge(annotations_iter);

    let mut found_port = false;
    iters
        .filter(|(key, value)| match key.as_str() {
            MAIN_LABEL => value.as_str() == addr,
            PORT_LABEL => {
                found_port = true;
                value.as_str() == port
            }
            URL_ANNOTATION => value
                .lines()
                .find(|str| {
                    if str.find(':').is_some() {
                        str == &format!("{addr}:{port}").as_str()
                    } else {
                        str == &addr
                    }
                })
                .is_some(),
            _ => false,
        })
        .count()
        > 0
}

impl fmt::Debug for ServerDeploymentStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Connectable(_) => write!(f, "Connectable"),
            Self::Starting => write!(f, "Starting"),
            Self::PodOk => write!(f, "PodOk"),
            Self::Offline => write!(f, "Offline"),
            Self::Unavailable(s) => write!(f, "Unavailable ({s})"),
        }
    }
}
impl From<kube::Error> for OpaqueError {
    fn from(value: kube::Error) -> Self {
        OpaqueError::create(value.to_string().as_str())
    }
}