From dffd117f5aba76c9114b8d9078f4240ad1da6c4e Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Tue, 26 Mar 2024 18:26:20 +0100
Subject: [PATCH] Implement shell setup for posix

---
 lib/system/env/mod.rs   |  4 ++
 lib/system/env/shell.rs | 37 +++++++++++++++++
 lib/system/env/unix.rs  | 90 +++++++++++++++++++++++++++++++++++++----
 3 files changed, 123 insertions(+), 8 deletions(-)
 create mode 100644 lib/system/env/shell.rs

diff --git a/lib/system/env/mod.rs b/lib/system/env/mod.rs
index b26cfbb..4a468a8 100644
--- a/lib/system/env/mod.rs
+++ b/lib/system/env/mod.rs
@@ -1,5 +1,7 @@
 use crate::{result::RokitResult, storage::Home};
 
+mod shell;
+
 #[cfg(unix)]
 mod unix;
 
@@ -8,6 +10,8 @@ mod windows;
 
 /**
     Tries to add the Rokit binaries directory to the system PATH.
+
+    Returns `true` if the directory was added to the PATH, `false` otherwise.
 */
 pub async fn add_to_path(home: &Home) -> RokitResult<bool> {
     #[cfg(unix)]
diff --git a/lib/system/env/shell.rs b/lib/system/env/shell.rs
new file mode 100644
index 0000000..49791b9
--- /dev/null
+++ b/lib/system/env/shell.rs
@@ -0,0 +1,37 @@
+use std::env::var;
+
+#[derive(Debug, Clone, Copy)]
+pub enum Shell {
+    Posix,
+    Bash,
+    Zsh,
+}
+
+impl Shell {
+    pub const ALL: [Self; 3] = [Self::Posix, Self::Bash, Self::Zsh];
+
+    pub const fn name(&self) -> &'static str {
+        match self {
+            Self::Posix => "sh",
+            Self::Bash => "bash",
+            Self::Zsh => "zsh",
+        }
+    }
+
+    pub const fn env_file_path(&self) -> &'static str {
+        match self {
+            Self::Posix => ".profile",
+            Self::Bash => ".bashrc",
+            Self::Zsh => ".zshenv",
+        }
+    }
+
+    pub fn env_file_should_create_if_nonexistent(&self) -> bool {
+        // Create a new shell env file for the user if we are
+        // confident that this is the shell that they are using
+        var("SHELL").map_or(false, |current_shell| {
+            // Detect /bin/sh, /bin/bash, /bin/zsh, etc
+            current_shell.ends_with(&format!("/{}", self.name()))
+        })
+    }
+}
diff --git a/lib/system/env/unix.rs b/lib/system/env/unix.rs
index 1c0a398..691d142 100644
--- a/lib/system/env/unix.rs
+++ b/lib/system/env/unix.rs
@@ -1,24 +1,98 @@
-use tokio::fs::write;
+use std::path::PathBuf;
+
+use futures::{stream::FuturesUnordered, TryStreamExt};
+use tokio::{
+    fs::{read_to_string, write},
+    io::ErrorKind,
+};
 
 use crate::{
     result::{RokitError, RokitResult},
     storage::Home,
 };
 
+use super::shell::Shell;
+
 const ENV_SHELL_FILE_PATH: &str = "env";
 const ENV_SHELL_SCRIPT: &str = include_str!("./env.sh");
 
 pub async fn add_to_path(home: &Home) -> RokitResult<bool> {
-    // Write our shell script to the known location
+    // Find our binaries dir and try to format it as "$HOME/.rokit/bin"
     let bin_dir = home.path().join("bin");
+    let bin_dir_str = bin_dir.to_str().ok_or(RokitError::InvalidUtf8)?;
+    let bin_dir_in_home = replace_home_path_with_var(bin_dir_str);
+
+    // Do the same for the shell script path - "$HOME/.rokit/env"
     let file_path = home.path().join(ENV_SHELL_FILE_PATH);
-    let file_contents = ENV_SHELL_SCRIPT.replace(
-        "{rokit_bin_path}",
-        bin_dir.to_str().ok_or(RokitError::InvalidUtf8)?,
-    );
+    let file_path_str = file_path.to_str().ok_or(RokitError::InvalidUtf8)?;
+    let file_path_in_home = replace_home_path_with_var(file_path_str);
+
+    // Write our shell init script to the known location
+    let file_contents = ENV_SHELL_SCRIPT.replace("{rokit_bin_path}", &bin_dir_in_home);
     write(file_path, file_contents).await?;
 
-    // TODO: Add the path to known shell profile(s)
+    // Add the path to known shell profiles
+    let added_any = if let Some(home_dir) = dirs::home_dir() {
+        let futs = Shell::ALL
+            .iter()
+            .map(|shell| {
+                let shell_env_path = home_dir.join(shell.env_file_path());
+                let shell_should_create = shell.env_file_should_create_if_nonexistent();
+                append_to_shell_file(
+                    shell_env_path,
+                    format!(". \"{file_path_in_home}\""),
+                    shell_should_create,
+                )
+            })
+            .collect::<FuturesUnordered<_>>();
+        futs.try_collect::<Vec<_>>()
+            .await?
+            .iter()
+            .any(|added| *added)
+    } else {
+        false
+    };
+
+    Ok(added_any)
+}
+
+async fn append_to_shell_file(
+    file_path: PathBuf,
+    line_to_append: String,
+    create_if_nonexistent: bool,
+) -> RokitResult<bool> {
+    let mut file_contents = match read_to_string(&file_path).await {
+        Ok(contents) => contents,
+        Err(e) if e.kind() == ErrorKind::NotFound && create_if_nonexistent => String::new(),
+        Err(e) => return Err(e.into()),
+    };
+
+    if file_contents.contains(&line_to_append) {
+        return Ok(false);
+    }
+
+    // NOTE: Make sure we put the new contents on their own
+    // line and not conflicting with any existing command(s)
+    if !file_contents.ends_with('\n') {
+        file_contents.push('\n');
+    }
+
+    file_contents.push_str(&line_to_append);
+    file_contents.push('\n');
+
+    write(file_path, file_contents).await?;
+
+    Ok(true)
+}
 
-    Ok(false)
+fn replace_home_path_with_var(path: &str) -> String {
+    let home_dir = match dirs::home_dir() {
+        Some(home_dir) => home_dir,
+        None => return path.to_string(),
+    };
+    let home_dir_str = match home_dir.to_str() {
+        Some(home_dir_str) => home_dir_str,
+        None => return path.to_string(),
+    };
+    path.replace(home_dir_str, "$HOME")
 }