harmony/k3d/src/lib.rs

mod downloadable_asset;
use downloadable_asset::*;

use kube::Client;
use log::{debug, info};
use std::{ffi::OsStr, path::PathBuf};

const K3D_BIN_FILE_NAME: &str = "k3d";

pub struct K3d {
    base_dir: PathBuf,
    cluster_name: Option<String>,
}

impl K3d {
    pub fn new(base_dir: PathBuf, cluster_name: Option<String>) -> Self {
        Self {
            base_dir,
            cluster_name,
        }
    }

    async fn get_binary_for_current_platform(
        &self,
        latest_release: octocrab::models::repos::Release,
    ) -> DownloadableAsset {
        let os = std::env::consts::OS;
        let arch = std::env::consts::ARCH;

        debug!("Detecting platform: OS={}, ARCH={}", os, arch);

        let binary_pattern = match (os, arch) {
            ("linux", "x86") => "k3d-linux-386",
            ("linux", "x86_64") => "k3d-linux-amd64",
            ("linux", "arm") => "k3d-linux-arm",
            ("linux", "aarch64") => "k3d-linux-arm64",
            ("windows", "x86_64") => "k3d-windows-amd64.exe",
            ("macos", "x86_64") => "k3d-darwin-amd64",
            ("macos", "aarch64") => "k3d-darwin-arm64",
            _ => panic!("Unsupported platform: {}-{}", os, arch),
        };

        debug!("Looking for binary matching pattern: {}", binary_pattern);

        let binary_asset = latest_release
            .assets
            .iter()
            .find(|asset| asset.name == binary_pattern)
            .unwrap_or_else(|| panic!("No matching binary found for {}", binary_pattern));

        let binary_url = binary_asset.browser_download_url.clone();

        let checksums_asset = latest_release
            .assets
            .iter()
            .find(|asset| asset.name == "checksums.txt")
            .expect("Checksums file not found in release assets");

        let checksums_url = checksums_asset.browser_download_url.clone();

        let body = reqwest::get(checksums_url)
            .await
            .unwrap()
            .text()
            .await
            .unwrap();

        let checksum = body
            .lines()
            .find_map(|line| {
                if line.ends_with(&binary_pattern) {
                    Some(line.split_whitespace().next().unwrap_or("").to_string())
                } else {
                    None
                }
            })
            .unwrap_or_else(|| panic!("Checksum not found for {}", binary_pattern));

        debug!("Found binary at {} with checksum {}", binary_url, checksum);

        DownloadableAsset {
            url: binary_url,
            file_name: K3D_BIN_FILE_NAME.to_string(),
            checksum,
        }
    }

    pub async fn download_latest_release(&self) -> Result<PathBuf, String> {
        let latest_release = self.get_latest_release_tag().await.unwrap();

        let release_binary = self.get_binary_for_current_platform(latest_release).await;
        debug!("Foudn K3d binary to install : {release_binary:#?}");
        release_binary.download_to_path(self.base_dir.clone()).await
    }

    // TODO : Make sure this will only find actual released versions, no prereleases or test
    // builds
    pub async fn get_latest_release_tag(&self) -> Result<octocrab::models::repos::Release, String> {
        let octo = octocrab::instance();
        let latest_release = octo
            .repos("k3d-io", "k3d")
            .releases()
            .get_latest()
            .await
            .map_err(|e| e.to_string())?;
        debug!("Got k3d releases {latest_release:#?}");

        Ok(latest_release)
    }

    /// Checks if k3d binary exists and is executable
    ///
    /// Verifies that:
    /// 1. The k3d binary exists in the base directory
    /// 2. It has proper executable permissions (on Unix systems)
    /// 3. It responds correctly to a simple command (`k3d --version`)
    pub fn is_installed(&self) -> bool {
        let binary_path = self.get_k3d_binary_path();

        if !binary_path.exists() {
            debug!("K3d binary not found at {:?}", binary_path);
            return false;
        }

        if !self.ensure_binary_executable(&binary_path) {
            return false;
        }

        self.can_execute_binary_check(&binary_path)
    }

    /// Verifies if the specified cluster is already created
    ///
    /// Executes `k3d cluster list <cluster_name>` and checks for a successful response,
    /// indicating that the cluster exists and is registered with k3d.
    pub fn is_cluster_initialized(&self) -> bool {
        let cluster_name = match self.get_cluster_name() {
            Ok(name) => name,
            Err(_) => {
                debug!("Could not get cluster name, can't verify if cluster is initialized");
                return false;
            }
        };

        let binary_path = self.base_dir.join(K3D_BIN_FILE_NAME);
        if !binary_path.exists() {
            return false;
        }

        self.verify_cluster_exists(&binary_path, cluster_name)
    }

    fn get_cluster_name(&self) -> Result<&String, String> {
        match &self.cluster_name {
            Some(name) => Ok(name),
            None => Err("No cluster name available".to_string()),
        }
    }

    /// Creates a new k3d cluster with the specified name
    ///
    /// This method:
    /// 1. Creates a new k3d cluster using `k3d cluster create <cluster_name>`
    /// 2. Waits for the cluster to initialize
    /// 3. Returns a configured Kubernetes client connected to the cluster
    ///
    /// # Returns
    /// - `Ok(Client)` - Successfully created cluster and connected client
    /// - `Err(String)` - Error message detailing what went wrong
    pub async fn initialize_cluster(&self) -> Result<Client, String> {
        let cluster_name = match self.get_cluster_name() {
            Ok(name) => name,
            Err(_) => return Err("Could not get cluster_name, cannot initialize".to_string()),
        };

        debug!("Initializing k3d cluster '{}'", cluster_name);

        self.create_cluster(cluster_name)?;
        self.create_kubernetes_client().await
    }

    fn get_k3d_binary_path(&self) -> PathBuf {
        self.base_dir.join(K3D_BIN_FILE_NAME)
    }

    fn get_k3d_binary(&self) -> Result<PathBuf, String> {
        let path = self.get_k3d_binary_path();
        if !path.exists() {
            return Err(format!("K3d binary not found at {:?}", path));
        }
        Ok(path)
    }

    /// Ensures k3d is installed and the cluster is initialized
    ///
    /// This method provides a complete setup flow:
    /// 1. Checks if k3d is installed, downloads and installs it if needed
    /// 2. Verifies if the specified cluster exists, creates it if not
    /// 3. Returns a Kubernetes client connected to the cluster
    ///
    /// # Returns
    /// - `Ok(Client)` - Successfully ensured k3d and cluster are ready
    /// - `Err(String)` - Error message if any step failed
    pub async fn ensure_installed(&self) -> Result<Client, String> {
        if !self.is_installed() {
            debug!("K3d is not installed, downloading latest release");
            self.download_latest_release()
                .await
                .map_err(|e| format!("Failed to download k3d: {}", e))?;

            if !self.is_installed() {
                return Err("Failed to install k3d properly".to_string());
            }
        }

        let client;
        if !self.is_cluster_initialized() {
            debug!("Cluster is not initialized, initializing now");
            client = self.initialize_cluster().await?;
        } else {
            self.start_cluster().await?;

            debug!("K3d and cluster are already properly set up");
            client = self.create_kubernetes_client().await?;
        }

        self.ensure_k3d_config_is_default(self.get_cluster_name()?)?;
        Ok(client)
    }

    // Private helper methods

    #[cfg(not(target_os = "windows"))]
    fn ensure_binary_executable(&self, binary_path: &PathBuf) -> bool {
        use std::os::unix::fs::PermissionsExt;

        let mut perms = match std::fs::metadata(binary_path) {
            Ok(metadata) => metadata.permissions(),
            Err(e) => {
                debug!("Failed to get binary metadata: {}", e);
                return false;
            }
        };

        perms.set_mode(0o755);

        if let Err(e) = std::fs::set_permissions(binary_path, perms) {
            debug!("Failed to set executable permissions on k3d binary: {}", e);
            return false;
        }

        true
    }

    #[cfg(target_os = "windows")]
    fn ensure_binary_executable(&self, _binary_path: &PathBuf) -> bool {
        // Windows doesn't use executable file permissions
        true
    }

    fn can_execute_binary_check(&self, binary_path: &PathBuf) -> bool {
        match std::process::Command::new(binary_path)
            .arg("--version")
            .output()
        {
            Ok(output) => {
                if output.status.success() {
                    debug!("K3d binary is installed and working");
                    true
                } else {
                    debug!("K3d binary check failed: {:?}", output);
                    false
                }
            }
            Err(e) => {
                debug!("Failed to execute K3d binary: {}", e);
                false
            }
        }
    }

    fn verify_cluster_exists(&self, binary_path: &PathBuf, cluster_name: &str) -> bool {
        match std::process::Command::new(binary_path)
            .args(["cluster", "list", cluster_name, "--no-headers"])
            .output()
        {
            Ok(output) => {
                if output.status.success() && !output.stdout.is_empty() {
                    debug!("Cluster '{}' is initialized", cluster_name);
                    true
                } else {
                    debug!("Cluster '{}' is not initialized", cluster_name);
                    false
                }
            }
            Err(e) => {
                debug!("Failed to check cluster initialization: {}", e);
                false
            }
        }
    }

    pub fn run_k3d_command<I, S>(&self, args: I) -> Result<std::process::Output, String>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<std::ffi::OsStr>,
    {
        let binary_path = self.get_k3d_binary()?;
        self.run_command(binary_path, args)
    }

    pub fn run_command<I, S, C>(&self, cmd: C, args: I) -> Result<std::process::Output, String>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<std::ffi::OsStr>,
        C: AsRef<OsStr>,
    {
        let output = std::process::Command::new(cmd).args(args).output();
        match output {
            Ok(output) => {
                let stderr = String::from_utf8_lossy(&output.stderr);
                debug!("stderr : {}", stderr);
                let stdout = String::from_utf8_lossy(&output.stdout);
                debug!("stdout : {}", stdout);
                Ok(output)
            }
            Err(e) => Err(format!("Failed to execute command: {}", e)),
        }
    }

    fn create_cluster(&self, cluster_name: &str) -> Result<(), String> {
        let output = self.run_k3d_command(["cluster", "create", cluster_name])?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("Failed to create cluster: {}", stderr));
        }

        info!("Successfully created k3d cluster '{}'", cluster_name);
        Ok(())
    }

    fn ensure_k3d_config_is_default(&self, cluster_name: &str) -> Result<(), String> {
        let output = self.run_k3d_command(["kubeconfig", "merge", "-d", cluster_name])?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("Failed to setup k3d kubeconfig : {}", stderr));
        }

        let output = self.run_command(
            "kubectl",
            ["config", "use-context", &format!("k3d-{cluster_name}")],
        )?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!(
                "Failed to switch kubectl context to k3d : {}",
                stderr
            ));
        }
        info!(
            "kubectl is now using 'k3d-{}' as default context",
            cluster_name
        );
        Ok(())
    }

    async fn create_kubernetes_client(&self) -> Result<Client, String> {
        Client::try_default()
            .await
            .map_err(|e| format!("Failed to create Kubernetes client: {}", e))
    }

    pub async fn get_client(&self) -> Result<Client, String> {
        match self.is_cluster_initialized() {
            true => Ok(self.create_kubernetes_client().await?),
            false => Err("Cannot get client! Cluster not initialized yet".to_string()),
        }
    }

    async fn start_cluster(&self) -> Result<(), String> {
        let cluster_name = self.get_cluster_name()?;
        let output = self.run_k3d_command(["cluster", "start", cluster_name])?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("Failed to start cluster: {}", stderr));
        }

        debug!("Successfully started k3d cluster '{}'", cluster_name);
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use regex::Regex;
    use std::path::PathBuf;

    use crate::{K3d, K3D_BIN_FILE_NAME};

    #[tokio::test]
    async fn k3d_latest_release_should_get_latest() {
        let dir = get_clean_test_directory();

        assert!(!dir.join(K3D_BIN_FILE_NAME).exists());

        let k3d = K3d::new(dir.clone(), None);
        let latest_release = k3d.get_latest_release_tag().await.unwrap();

        let tag_regex = Regex::new(r"^v\d+\.\d+\.\d+$").unwrap();
        assert!(tag_regex.is_match(&latest_release.tag_name));
        assert!(!latest_release.tag_name.is_empty());
    }

    #[tokio::test]
    async fn k3d_download_latest_release_should_get_latest_bin() {
        let dir = get_clean_test_directory();

        assert!(!dir.join(K3D_BIN_FILE_NAME).exists());

        let k3d = K3d::new(dir.clone(), None);
        let bin_file_path = k3d.download_latest_release().await.unwrap();
        assert_eq!(bin_file_path, dir.join(K3D_BIN_FILE_NAME));
        assert!(dir.join(K3D_BIN_FILE_NAME).exists());
    }

    fn get_clean_test_directory() -> PathBuf {
        let dir = PathBuf::from("/tmp/harmony-k3d-test-dir");

        if dir.exists() {
            if let Err(e) = std::fs::remove_dir_all(&dir) {
                // TODO sometimes this fails because of the race when running multiple tests at
                // once
                panic!("Failed to clean up test directory: {}", e);
            }
        }

        if let Err(e) = std::fs::create_dir_all(&dir) {
            panic!("Failed to create test directory: {}", e);
        }

        dir
    }
}