treeflow 0.2.1

CLI tool for simplified Git worktree management to speed up switching contexts when working collaboratively.
Documentation
use crate::utils::errors::CustomError;
use crate::core::project::Project;
use crate::utils::string_extensions::StringExtensions;
use crate::utils::errors::CustomError::{DuplicateProjectRepository, DuplicateProjectWorktrees, DuplicateWorkTypeName, DuplicateWorkTypePrefix};
use clap::Command;
use clap_complete::CompleteEnv;
use directories::ProjectDirs;
use serde::{Deserialize, Serialize};
use std::env;
use std::fs::File;
use std::io::Write;
use std::io::{ErrorKind, Read};
use std::path::{Path, PathBuf};
use path_clean::PathClean;
use crate::utils::Shell;

#[derive(Serialize)]
#[derive(Deserialize)]
#[derive(Clone)]
pub struct WorkType {
    pub type_name: String,
    pub prefix: String,
}

#[derive(Serialize)]
#[derive(Deserialize)]
pub struct Config {
    pub work_types: Vec<WorkType>,
    pub projects: Vec<Project>,
}

#[derive(Clone)]
pub struct ConfigPaths {
    pub directory: PathBuf,
    pub file: PathBuf,
}

/// Get config directory from environment variable TREEFLOW_CONFIG_DIR or default location
pub fn get_config_dir() -> Result<ConfigPaths, CustomError> {
    let config_dir_buf = match env::var("TREEFLOW_CONFIG_DIR") {
        Ok(config_dir_str) => PathBuf::from(config_dir_str),
        Err(_) => ProjectDirs::from("", "", "treeflow").ok_or(CustomError::Custom("Could not find project directories".to_string()))?.config_dir().to_path_buf()
    };

    Ok(ConfigPaths {
        directory: config_dir_buf.clone(),
        file: config_dir_buf.join("config.toml")
    })
}


impl Config {
    fn new() -> Self {
        Config { work_types: vec![], projects: vec![] }
    }

    pub fn remove_project(self, repository: &PathBuf) -> Result<Config, CustomError> {
        let mut config = self;
        let current_dir = env::current_dir().map_err(|error| CustomError::IoError(error))?.clean();
        let repo_path = current_dir.join(repository).as_path().clean();

        config.projects.retain(|project| project.repository != repo_path);
        Ok(config)
    }

    fn validate_project_add(&self, absolute_repo_path: &PathBuf, absolute_worktrees_path: &PathBuf) -> Result<(), CustomError> {
        // Check for duplicate repository path
        if self.projects.iter().any(|p| p.repository == *absolute_repo_path) {
            return Err(DuplicateProjectRepository(absolute_repo_path.clone()));
        }

        // Check for duplicate worktrees path
        if self.projects.iter().any(|p| p.worktrees == *absolute_worktrees_path) {
            return Err(DuplicateProjectWorktrees(absolute_worktrees_path.clone()));
        }

        Ok(())
    }

    pub fn add_project(self, repository: &PathBuf, worktrees_dir: Option<&PathBuf>) -> Result<Config, CustomError> {
        let current_dir = env::current_dir().map_err(|error| CustomError::IoError(error))?.clean();

        let repo_path = repository.clean();
        let absolute_repo_path = current_dir.join(&repo_path);

        let mut repo_directory_name = absolute_repo_path.iter().last().ok_or(CustomError::Custom("Could not get repo directory from".to_string()))?.to_owned();
        repo_directory_name.push("_wt");
        let absolute_worktrees_path = worktrees_dir
            .map(PathBuf::from)
            .map(|path| current_dir.join(path))
            .unwrap_or(PathBuf::from(absolute_repo_path.parent().ok_or(CustomError::Custom("No parent of repo.".to_string()))?.join(Path::new(&repo_directory_name))));
        let absolute_worktrees_path = absolute_worktrees_path.clean();

        self.validate_project_add(&absolute_repo_path, &absolute_worktrees_path)?;

        let mut config = self;
        config.projects.push(Project { repository: PathBuf::from(absolute_repo_path), worktrees: PathBuf::from(absolute_worktrees_path) });
        Ok(config)
    }

    fn validate_work_type_add(&self, type_name: &str, prefix: &str) -> Result<(), CustomError> {
        // Check for duplicate type name
        if self.work_types.iter().any(|wt| wt.type_name == type_name) {
            return Err(DuplicateWorkTypeName(type_name.to_string()));
        }

        // Check for duplicate prefix
        if self.work_types.iter().any(|wt| wt.prefix == prefix) {
            return Err(DuplicateWorkTypePrefix(prefix.to_string()));
        }

        Ok(())
    }

    pub fn add_work_type(self, type_name: String, prefix: String) -> Result<Config, CustomError> {
        self.validate_work_type_add(&type_name, &prefix)?;

        let mut config = self;
        config.work_types.push(WorkType { type_name, prefix });
        Ok(config)
    }

    pub fn remove_work_type(self, type_name: &str) -> Result<Config, CustomError> {
        let mut config = self;
        config.work_types.retain(|work_type| work_type.type_name != type_name);
        Ok(config)
    }

    pub fn print_work_types(&self) -> Result<(), CustomError> {
        for work_type in &self.work_types {
            println!("{}    {}", work_type.type_name, work_type.prefix);
        }
        Ok(())
    }

    pub fn current_project(&self) -> Result<Project, CustomError> {
        let current_dir = env::current_dir().map_err(|error| CustomError::IoError(error))?.clean();
        self
            .projects
            .iter()
            .find(|&project| current_dir.starts_with(&project.repository) || current_dir.starts_with(&project.worktrees))
            .map(|project| project.clone())
            .ok_or(CustomError::ProjectNotFound(current_dir))
    }

    pub fn work_type_from_name(&self, work_name: &str) -> Result<WorkType, CustomError> {
        self
            .work_types
            .iter()
            .find(|&work_type| work_name == work_type.type_name)
            .map(|work_type| work_type.clone())
            .ok_or(CustomError::WorkTypeNotFound(work_name.to_string()))
    }

    pub fn print_projects(&self) -> Result<(), CustomError> {
        if self.projects.is_empty() {
            println!("<<No projects configured>>");
        } else {
            for project in &self.projects {
                println!("{}", project.repository.display());
            }
        }
        Ok(())
    }

    pub fn print(&self) -> Result<(), CustomError> {
        let toml = toml::to_string(self)?;
        let output = toml.trim_cr_end() + "\n";
        print!("{}", output);
        Ok(())
    }

    pub fn print_initialisation<T>(&self, shell: &Shell, enable_shorthands: bool, command_factory: T) -> Result<(), CustomError>
        where T: Fn() -> Command
    {
        let exe_path = env::current_exe()?;
        let treeflow_path = exe_path.to_str().unwrap();

        match shell {
            Shell::Bash => {
                println!("#!/usr/bin/env sh");
                println!();
                println!("{}", include_str!("../templates/bash/base.sh").replace("::treeflow_path::", treeflow_path));

                if enable_shorthands {
                    // Include static aliases
                    println!("{}", include_str!("../templates/bash/static-shorthands.sh"));

                    // Include dynamic work-type aliases
                    let template = include_str!("../templates/bash/dynamic-shorthands.sh");
                    for work_type in &self.work_types {
                        println!("{}", template.replace("::work_type::", work_type.type_name.as_str()))
                    }
                }

                println!("# Dynamic completion generated by clap_complete");
                env::set_var("COMPLETE", "bash");
                CompleteEnv::with_factory(command_factory).complete();

                Ok(())
            }
            Shell::Zsh => {
                println!("#!/usr/bin/env zsh");
                println!();
                println!("{}", include_str!("../templates/zsh/base.zsh").replace("::treeflow_path::", treeflow_path));

                if enable_shorthands {
                    // Include static aliases
                    println!("{}", include_str!("../templates/zsh/static-shorthands.zsh"));

                    // Include dynamic work-type aliases
                    let template = include_str!("../templates/zsh/dynamic-shorthands.zsh");
                    for work_type in &self.work_types {
                        println!("{}", template.replace("::work_type::", work_type.type_name.as_str()))
                    }
                }

                println!("# Dynamic completion generated by clap_complete");
                env::set_var("COMPLETE", "zsh");
                CompleteEnv::with_factory(command_factory).complete();

                Ok(())
            }
        }
    }

    pub fn load(config_paths: &ConfigPaths) -> Result<Self, CustomError> {
        let config_file = config_paths.clone().file;
        match File::options().read(true).open(config_file) {
            Ok(mut file) => {
                let mut buffer = String::new();

                file.read_to_string(&mut buffer)?;
                let config = toml::from_str(&buffer.to_string()).map_err(|e| CustomError::DeserialisationError(e));
                config
            }
            Err(error) => {
                if error.kind() == ErrorKind::NotFound { Ok(Config::new()) } else { Err(CustomError::IoError(error)) }
            }
        }
    }

    pub fn save(&self, config_paths: ConfigPaths) -> Result<(), CustomError> {
        let toml = toml::to_string(self)?;
        let _ = std::fs::create_dir_all(config_paths.directory)?;
        let mut file = File::options().create(true).write(true).truncate(true).open(config_paths.file)?;
        Ok(write!(file, "{}", toml)?)
    }
}