nethermit 0.1.3

use anyhow::{Context, Result};
use colored::*;
use indicatif::{ProgressBar, ProgressStyle};
use jsonnet::JsonnetVm;
use k8s_openapi::api::core::v1::ConfigMap;
use kube::{
    api::{Api, DynamicObject, GroupVersionKind, Patch, PatchParams, PostParams},
    discovery::{self, Scope},
    Client,
};
use serde_json::Value;
use std::{collections::BTreeMap, path::PathBuf, time::Duration};
use tabled::Tabled;

#[cfg(test)]
use mockall::automock;

pub const SET_LABEL_KEY: &str = "nethermit.set";
pub const SETS_CONFIGMAP_NAME: &str = "nethermit-sets";

#[derive(Debug, serde::Serialize, serde::Deserialize, Clone)]
pub struct ResourceInfo {
    pub api_version: String,
    pub kind: String,
    pub name: String,
    pub namespace: Option<String>,
}

#[derive(Tabled)]
pub struct SetInfo {
    #[tabled(rename = "Set Name")]
    pub name: String,
    #[tabled(rename = "Resources")]
    pub resource_count: usize,
}

#[cfg_attr(test, automock)]
#[async_trait::async_trait]
pub trait KubeClient {
    async fn get_configmap(&self, name: &str, namespace: &str) -> Result<Option<ConfigMap>>;
    async fn create_configmap(&self, cm: &ConfigMap, namespace: &str) -> Result<()>;
    async fn patch_configmap<'a>(
        &'a self,
        name: &'a str,
        namespace: &'a str,
        patch: &'a Patch<&'a Value>,
    ) -> Result<()>;
    async fn delete_configmap(&self, name: &str, namespace: &str) -> Result<()>;
    async fn create_pod(&self, pod: &Value, namespace: &str) -> Result<()>;
    async fn delete_pod(&self, name: &str, namespace: &str) -> Result<()>;
}

pub struct RealKubeClient {
    client: Client,
}

#[async_trait::async_trait]
impl KubeClient for RealKubeClient {
    async fn get_configmap(&self, name: &str, namespace: &str) -> Result<Option<ConfigMap>> {
        let api: Api<ConfigMap> = Api::namespaced(self.client.clone(), namespace);
        match api.get(name).await {
            Ok(cm) => Ok(Some(cm)),
            Err(kube::Error::Api(err)) if err.code == 404 => Ok(None),
            Err(e) => Err(e.into()),
        }
    }

    async fn create_configmap(&self, cm: &ConfigMap, namespace: &str) -> Result<()> {
        let api: Api<ConfigMap> = Api::namespaced(self.client.clone(), namespace);
        api.create(&PostParams::default(), cm).await?;
        Ok(())
    }

    async fn patch_configmap<'a>(
        &'a self,
        name: &'a str,
        namespace: &'a str,
        patch: &'a Patch<&'a Value>,
    ) -> Result<()> {
        let api: Api<ConfigMap> = Api::namespaced(self.client.clone(), namespace);
        api.patch(name, &PatchParams::apply("nethermit"), patch)
            .await?;
        Ok(())
    }

    async fn delete_configmap(&self, name: &str, namespace: &str) -> Result<()> {
        let api: Api<ConfigMap> = Api::namespaced(self.client.clone(), namespace);
        api.delete(name, &Default::default()).await?;
        Ok(())
    }

    async fn create_pod(&self, pod: &Value, namespace: &str) -> Result<()> {
        let api: Api<DynamicObject> = Api::namespaced_with(
            self.client.clone(),
            namespace,
            &kube::core::ApiResource::erase::<k8s_openapi::api::core::v1::Pod>(&()),
        );
        api.create(
            &PostParams::default(),
            &serde_json::from_value(pod.clone())?,
        )
        .await?;
        Ok(())
    }

    async fn delete_pod(&self, name: &str, namespace: &str) -> Result<()> {
        let api: Api<DynamicObject> = Api::namespaced_with(
            self.client.clone(),
            namespace,
            &kube::core::ApiResource::erase::<k8s_openapi::api::core::v1::Pod>(&()),
        );
        api.delete(name, &Default::default()).await?;
        Ok(())
    }
}

impl RealKubeClient {
    pub async fn new() -> Result<Self> {
        Ok(Self {
            client: Client::try_default().await?,
        })
    }
}

pub fn create_spinner(msg: &str) -> ProgressBar {
    let pb = ProgressBar::new_spinner();
    pb.enable_steady_tick(Duration::from_millis(100));
    pb.set_style(
        ProgressStyle::default_spinner()
            .tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈")
            .template("{spinner:.blue} {msg}")
            .unwrap(),
    );
    pb.set_message(msg.to_string());
    pb
}

pub fn format_success(msg: &str) -> String {
    format!("✓ {msg}").green().to_string()
}

pub fn format_error(msg: &str) -> String {
    format!("✗ {msg}").red().to_string()
}

pub fn format_info(msg: &str) -> String {
    format!("ℹ {msg}").blue().to_string()
}

/// Check if a set exists in the ConfigMap
pub async fn set_exists(client: &impl KubeClient, name: &str) -> Result<bool> {
    match client.get_configmap(SETS_CONFIGMAP_NAME, "default").await? {
        Some(cm) => Ok(cm
            .data
            .as_ref()
            .map(|data| data.contains_key(&format!("{name}.resources")))
            .unwrap_or(false)),
        None => Ok(false),
    }
}

/// List all installed sets with resource counts
pub async fn list_sets_info(client: &impl KubeClient) -> Result<Vec<SetInfo>> {
    match client.get_configmap(SETS_CONFIGMAP_NAME, "default").await? {
        Some(cm) => {
            let data = cm.data.unwrap_or_default();
            Ok(data
                .iter()
                .filter(|(k, _)| k.ends_with(".resources"))
                .map(|(k, v)| {
                    let name = k.trim_end_matches(".resources").to_string();
                    let resources: Vec<ResourceInfo> = serde_json::from_str(v).unwrap_or_default();
                    SetInfo {
                        name,
                        resource_count: resources.len(),
                    }
                })
                .collect())
        }
        None => Ok(Vec::new()),
    }
}

/// Add a set to the ConfigMap with its resources
pub async fn add_set(
    client: &impl KubeClient,
    name: &str,
    resources: Vec<ResourceInfo>,
) -> Result<()> {
    // Check if set already exists
    if set_exists(client, name).await? {
        return Err(anyhow::anyhow!("Set '{}' already exists", name));
    }

    // Create or update ConfigMap
    let mut data = BTreeMap::new();
    data.insert(
        format!("{name}.resources"),
        serde_json::to_string(&resources)?,
    );

    let patch_json = serde_json::json!({
        "apiVersion": "v1",
        "kind": "ConfigMap",
        "metadata": {
            "name": SETS_CONFIGMAP_NAME,
        },
        "data": data
    });

    let patch = Patch::Apply(&patch_json);
    client
        .patch_configmap(SETS_CONFIGMAP_NAME, "default", &patch)
        .await?;

    Ok(())
}

/// Remove a set from the ConfigMap
pub async fn remove_set(client: &impl KubeClient, name: &str) -> Result<()> {
    // Check if set exists
    if !set_exists(client, name).await? {
        return Err(anyhow::anyhow!("Set '{}' does not exist", name));
    }

    match client.get_configmap(SETS_CONFIGMAP_NAME, "default").await? {
        Some(cm) => {
            let mut data = cm.data.unwrap_or_default();
            data.remove(&format!("{name}.resources"));

            let patch_json = serde_json::json!({
                "apiVersion": "v1",
                "kind": "ConfigMap",
                "metadata": {
                    "name": SETS_CONFIGMAP_NAME,
                },
                "data": data
            });

            let patch = Patch::Apply(&patch_json);
            client
                .patch_configmap(SETS_CONFIGMAP_NAME, "default", &patch)
                .await?;
            Ok(())
        }
        None => Err(anyhow::anyhow!("ConfigMap not found")),
    }
}

/// Convert Jsonnet content to YAML with optional import paths
pub fn jsonnet_to_yaml(content: &str, import_paths: &[PathBuf]) -> Result<String> {
    let mut vm = JsonnetVm::new();

    // Configure VM settings
    vm.max_stack(100);
    vm.max_trace(Some(100));

    // Add import paths
    for path in import_paths {
        vm.jpath_add(path.to_string_lossy().as_ref());
    }

    // Evaluate Jsonnet content
    let json = vm
        .evaluate_snippet("snippet", content)
        .map_err(|e| anyhow::anyhow!("Failed to evaluate Jsonnet: {}", e))?;

    // Parse JSON
    let value: Value = serde_json::from_str(&json).context("Failed to parse JSON")?;

    // Handle array of resources or single resource
    let yaml = match value {
        Value::Array(resources) => {
            // Convert each resource to YAML and join with separator
            resources
                .into_iter()
                .map(|r| serde_yaml::to_string(&r).context("Failed to convert to YAML"))
                .collect::<Result<Vec<_>>>()?
                .join("\n---\n")
        }
        _ => serde_yaml::to_string(&value).context("Failed to convert to YAML")?,
    };

    Ok(yaml)
}

/// Apply the YAML configuration to a Kubernetes cluster
pub async fn apply_kubernetes(yaml: &str, set_name: Option<&str>) -> Result<()> {
    // Create kubernetes client
    let client = RealKubeClient::new().await?;
    let kube_client = client.client.clone();

    // Split YAML into documents
    let documents: Vec<&str> = yaml
        .split("\n---\n")
        .filter(|s| !s.trim().is_empty())
        .collect();
    let pb = create_spinner("Applying resources");

    let mut resources = Vec::new();

    for (i, doc) in documents.iter().enumerate() {
        // Parse YAML to get type information
        let mut obj: DynamicObject =
            serde_yaml::from_str(doc).context("Failed to parse YAML document")?;
        let type_meta = obj.types.clone().context("Missing type metadata")?;

        // Add set label if provided
        if let Some(name) = set_name {
            let labels = obj.metadata.labels.get_or_insert_with(BTreeMap::new);
            labels.insert(SET_LABEL_KEY.to_string(), name.to_string());
        }

        // Get API resource for this type
        let gvk = GroupVersionKind::try_from(&type_meta)?;
        let (api_resource, caps) = discovery::pinned_kind(&kube_client, &gvk).await?;

        // Create API for this resource type
        let api: Api<DynamicObject> = match caps.scope {
            Scope::Cluster => Api::all_with(kube_client.clone(), &api_resource),
            Scope::Namespaced => {
                let namespace = obj.metadata.namespace.as_deref().unwrap_or("default");
                Api::namespaced_with(kube_client.clone(), namespace, &api_resource)
            }
        };

        // Apply the resource
        let name = obj
            .metadata
            .name
            .as_deref()
            .context("Resource must have a name")?;

        pb.set_message(format!(
            "Applying {} '{}' ({}/{})",
            type_meta.kind,
            name,
            i + 1,
            documents.len()
        ));

        api.patch(name, &PatchParams::apply("nethermit"), &Patch::Apply(&obj))
            .await
            .with_context(|| format!("Failed to apply resource {name}"))?;

        // Store resource info if this is part of a set
        if let Some(_set_name) = set_name {
            resources.push(ResourceInfo {
                api_version: type_meta.api_version,
                kind: type_meta.kind,
                name: name.to_string(),
                namespace: obj.metadata.namespace,
            });
        }
    }

    pb.finish_with_message(format_success(&format!(
        "Applied {} resources",
        documents.len()
    )));

    // If this is a set installation, store the resource information
    if let Some(name) = set_name {
        add_set(&client, name, resources).await?;
    }

    Ok(())
}

/// Delete all resources in a set using stored information
pub async fn delete_set_resources(client: &impl KubeClient, set_name: &str) -> Result<()> {
    // Get the set's resources from ConfigMap
    let cm = match client.get_configmap(SETS_CONFIGMAP_NAME, "default").await? {
        Some(cm) => cm,
        None => return Err(anyhow::anyhow!("ConfigMap not found")),
    };

    let resources: Vec<ResourceInfo> = if let Some(data) = cm.data {
        if let Some(resources_str) = data.get(&format!("{set_name}.resources")) {
            serde_json::from_str(resources_str)?
        } else {
            return Err(anyhow::anyhow!("Set '{}' not found", set_name));
        }
    } else {
        return Err(anyhow::anyhow!("No sets found"));
    };

    let pb = create_spinner("Deleting resources");

    // Delete each resource
    for (i, resource) in resources.iter().enumerate() {
        pb.set_message(format!(
            "Deleting {} '{}' ({}/{})",
            resource.kind,
            resource.name,
            i + 1,
            resources.len()
        ));

        if resource.api_version == "v1" && resource.kind == "Pod" {
            if let Some(ref ns) = resource.namespace {
                client.delete_pod(&resource.name, ns).await?;
            }
        } else {
            // For other resource types, we'll need to add more methods to KubeClient
            return Err(anyhow::anyhow!(
                "Unsupported resource type: {}",
                resource.kind
            ));
        }
    }

    pb.finish_with_message(format_success(&format!(
        "Deleted {} resources",
        resources.len()
    )));

    // Remove the set from the ConfigMap
    remove_set(client, set_name).await?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use mockall::predicate::*;

    #[tokio::test]
    async fn test_set_management() -> Result<()> {
        let mut mock_client = MockKubeClient::new();
        let set_name = "test-set";

        // Initial check if set exists - should return None
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(|_, _| Ok(None));

        // Expect patch_configmap for adding set
        mock_client
            .expect_patch_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"), always())
            .times(1)
            .returning(|_, _, _| Ok(()));

        // Second check if set exists - should return Some with our set
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(move |_, _| {
                let mut data = BTreeMap::new();
                data.insert(
                    format!("{name}.resources", name = set_name),
                    "[]".to_string(),
                );
                Ok(Some(ConfigMap {
                    data: Some(data),
                    ..Default::default()
                }))
            });

        // For list_sets
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(move |_, _| {
                let mut data = BTreeMap::new();
                data.insert(
                    format!("{name}.resources", name = set_name),
                    "[]".to_string(),
                );
                Ok(Some(ConfigMap {
                    data: Some(data),
                    ..Default::default()
                }))
            });

        // For remove_set - check if set exists
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(move |_, _| {
                let mut data = BTreeMap::new();
                data.insert(
                    format!("{name}.resources", name = set_name),
                    "[]".to_string(),
                );
                Ok(Some(ConfigMap {
                    data: Some(data),
                    ..Default::default()
                }))
            });

        // For remove_set - get configmap to update
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(move |_, _| {
                let mut data = BTreeMap::new();
                data.insert(
                    format!("{name}.resources", name = set_name),
                    "[]".to_string(),
                );
                Ok(Some(ConfigMap {
                    data: Some(data),
                    ..Default::default()
                }))
            });

        // For remove_set - patch configmap
        mock_client
            .expect_patch_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"), always())
            .times(1)
            .returning(|_, _, _| Ok(()));

        // Final check if set exists - should return Some with empty data
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(|_, _| {
                Ok(Some(ConfigMap {
                    data: Some(BTreeMap::new()),
                    ..Default::default()
                }))
            });

        // For final list_sets
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(|_, _| {
                Ok(Some(ConfigMap {
                    data: Some(BTreeMap::new()),
                    ..Default::default()
                }))
            });

        // Test adding a new set
        add_set(&mock_client, set_name, Vec::new()).await?;

        // Verify set exists
        assert!(set_exists(&mock_client, set_name).await?);

        // List sets should include our test set
        let sets = list_sets_info(&mock_client).await?;
        assert!(sets.iter().any(|s| s.name == set_name));

        // Test removing the set
        remove_set(&mock_client, set_name).await?;

        // Verify set no longer exists
        assert!(!set_exists(&mock_client, set_name).await?);

        // List should be empty
        let sets = list_sets_info(&mock_client).await?;
        assert!(!sets.iter().any(|s| s.name == set_name));

        Ok(())
    }

    #[tokio::test]
    async fn test_resource_labels() -> Result<()> {
        let set_name = "test-label-set";

        // First check if set exists - should return None
        let mut mock_client = MockKubeClient::new();
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .returning(|_, _| Ok(None));

        // Expect patch_configmap for adding set
        mock_client
            .expect_patch_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"), always())
            .returning(|_, _, _| Ok(()));

        // Create test resources and add to set
        let pod = serde_json::json!({
            "apiVersion": "v1",
            "kind": "Pod",
            "metadata": {
                "name": "test-pod",
                "namespace": "default",
                "labels": {
                    SET_LABEL_KEY: set_name
                }
            },
            "spec": {
                "containers": [{
                    "name": "test",
                    "image": "nginx"
                }]
            }
        });

        // Expect create_pod
        mock_client
            .expect_create_pod()
            .with(always(), eq("default"))
            .returning(|_, _| Ok(()));

        // Create the pod first
        mock_client.create_pod(&pod, "default").await?;

        // For add_set - check if set exists
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .returning(|_, _| Ok(None));

        // Add the pod to our set
        let resources = vec![ResourceInfo {
            api_version: "v1".to_string(),
            kind: "Pod".to_string(),
            name: "test-pod".to_string(),
            namespace: Some("default".to_string()),
        }];
        add_set(&mock_client, set_name, resources).await?;

        Ok(())
    }

    #[tokio::test]
    async fn test_delete_set_resources() -> Result<()> {
        let mut mock_client = MockKubeClient::new();
        let set_name = "test-delete-set";
        let pod_name = "test-delete-pod";

        // Create test resources and add to set
        let pod = serde_json::json!({
            "apiVersion": "v1",
            "kind": "Pod",
            "metadata": {
                "name": pod_name,
                "namespace": "default"
            },
            "spec": {
                "containers": [{
                    "name": "test",
                    "image": "nginx"
                }]
            }
        });

        let resources = vec![ResourceInfo {
            api_version: "v1".to_string(),
            kind: "Pod".to_string(),
            name: pod_name.to_string(),
            namespace: Some("default".to_string()),
        }];

        // For add_set - check if set exists
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(|_, _| Ok(None));

        // For add_set - patch configmap
        mock_client
            .expect_patch_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"), always())
            .times(1)
            .returning(|_, _, _| Ok(()));

        // For create_pod
        mock_client
            .expect_create_pod()
            .with(always(), eq("default"))
            .times(1)
            .returning(|_, _| Ok(()));

        // For delete_set_resources - get configmap
        let resources_clone = resources.clone();
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(move |_, _| {
                let mut data = BTreeMap::new();
                data.insert(
                    format!("{set_name}.resources"),
                    serde_json::to_string(&resources_clone).unwrap(),
                );
                Ok(Some(ConfigMap {
                    data: Some(data),
                    ..Default::default()
                }))
            });

        // For delete_pod
        mock_client
            .expect_delete_pod()
            .with(eq(pod_name), eq("default"))
            .times(1)
            .returning(|_, _| Ok(()));

        // For remove_set - check if set exists
        let resources_clone = resources.clone();
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(move |_, _| {
                let mut data = BTreeMap::new();
                data.insert(
                    format!("{set_name}.resources"),
                    serde_json::to_string(&resources_clone).unwrap(),
                );
                Ok(Some(ConfigMap {
                    data: Some(data),
                    ..Default::default()
                }))
            });

        // For remove_set - get configmap to update
        let resources_clone = resources.clone();
        mock_client
            .expect_get_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"))
            .times(1)
            .returning(move |_, _| {
                let mut data = BTreeMap::new();
                data.insert(
                    format!("{set_name}.resources"),
                    serde_json::to_string(&resources_clone).unwrap(),
                );
                Ok(Some(ConfigMap {
                    data: Some(data),
                    ..Default::default()
                }))
            });

        // For remove_set - patch configmap
        mock_client
            .expect_patch_configmap()
            .with(eq(SETS_CONFIGMAP_NAME), eq("default"), always())
            .times(1)
            .returning(|_, _, _| Ok(()));

        // Create the pod and add it to the set
        mock_client.create_pod(&pod, "default").await?;
        add_set(&mock_client, set_name, resources).await?;

        // Delete the set and its resources
        delete_set_resources(&mock_client, set_name).await?;

        Ok(())
    }
}