feat: add support for NO_COLOR (#67)

5 files changed, 164 insertions, 47 deletions

diff --git a/src/repl/handler.rs b/src/repl/handler.rs
index a7f4b08..159127c 100644
--- a/src/repl/handler.rs
+++ b/src/repl/handler.rs
@@ -50,50 +50,47 @@ impl ReplCmdHandler {
                     self.reply.borrow_mut().clear();
                     return Ok(());
                 }
-                let highlight = self.config.lock().highlight;
-                let light_theme = self.config.lock().light_theme;
                 let wg = WaitGroup::new();
                 let ret = render_stream(
                     &input,
                     &self.client,
-                    highlight,
-                    light_theme,
+                    self.config.clone(),
                     true,
                     self.abort.clone(),
                     wg.clone(),
                 );
                 wg.wait();
                 let buffer = ret?;
-                self.config.lock().save_message(&input, &buffer)?;
-                self.config.lock().save_conversation(&input, &buffer)?;
+                self.config.read().save_message(&input, &buffer)?;
+                self.config.write().save_conversation(&input, &buffer)?;
                 *self.reply.borrow_mut() = buffer;
             }
             ReplCmd::SetRole(name) => {
-                let output = self.config.lock().change_role(&name)?;
+                let output = self.config.write().change_role(&name)?;
                 print_now!("{}\n\n", output.trim_end());
             }
             ReplCmd::ClearRole => {
-                self.config.lock().clear_role()?;
+                self.config.write().clear_role()?;
                 print_now!("\n");
             }
             ReplCmd::Prompt(prompt) => {
-                self.config.lock().add_prompt(&prompt)?;
+                self.config.write().add_prompt(&prompt)?;
                 print_now!("\n");
             }
             ReplCmd::ViewInfo => {
-                let output = self.config.lock().info()?;
+                let output = self.config.read().info()?;
                 print_now!("{}\n\n", output.trim_end());
             }
             ReplCmd::UpdateConfig(input) => {
-                self.config.lock().update(&input)?;
+                self.config.write().update(&input)?;
                 print_now!("\n");
             }
             ReplCmd::StartConversation => {
-                self.config.lock().start_conversation()?;
+                self.config.write().start_conversation()?;
                 print_now!("\n");
             }
             ReplCmd::EndConversatoin => {
-                self.config.lock().end_conversation();
+                self.config.write().end_conversation();
                 print_now!("\n");
             }
         }
diff --git a/src/repl/highlighter.rs b/src/repl/highlighter.rs
new file mode 100644
index 0000000..593c307
--- /dev/null
+++ b/src/repl/highlighter.rs
@@ -0,0 +1,66 @@
+use crate::config::SharedConfig;
+
+use nu_ansi_term::{Color, Style};
+use reedline::{Highlighter, StyledText};
+
+const MATCH_COLOR: Color = Color::Green;
+
+pub struct ReplHighlighter {
+    external_commands: Vec<String>,
+    config: SharedConfig,
+}
+
+impl ReplHighlighter {
+    /// Construct the default highlighter with a given set of extern commands/keywords to detect and highlight
+    pub fn new(config: SharedConfig, external_commands: Vec<String>) -> ReplHighlighter {
+        Self {
+            config,
+            external_commands,
+        }
+    }
+}
+
+impl Highlighter for ReplHighlighter {
+    fn highlight(&self, line: &str, _cursor: usize) -> StyledText {
+        let mut styled_text = StyledText::new();
+        let color = if self.config.read().light_theme {
+            Color::Black
+        } else {
+            Color::White
+        };
+        let match_color = match self.config.read().highlight {
+            true => MATCH_COLOR,
+            false => color,
+        };
+
+        if self
+            .external_commands
+            .clone()
+            .iter()
+            .any(|x| line.contains(x))
+        {
+            let matches: Vec<&str> = self
+                .external_commands
+                .iter()
+                .filter(|c| line.contains(*c))
+                .map(std::ops::Deref::deref)
+                .collect();
+            let longest_match = matches.iter().fold("".to_string(), |acc, &item| {
+                if item.len() > acc.len() {
+                    item.to_string()
+                } else {
+                    acc
+                }
+            });
+            let buffer_split: Vec<&str> = line.splitn(2, &longest_match).collect();
+
+            styled_text.push((Style::new().fg(color), buffer_split[0].to_string()));
+            styled_text.push((Style::new().fg(match_color), longest_match));
+            styled_text.push((Style::new().fg(color), buffer_split[1].to_string()));
+        } else {
+            styled_text.push((Style::new().fg(color), line.to_string()));
+        }
+
+        styled_text
+    }
+}
diff --git a/src/repl/init.rs b/src/repl/init.rs
index ff75eff..6914cdc 100644
--- a/src/repl/init.rs
+++ b/src/repl/init.rs
@@ -1,22 +1,18 @@
-use super::REPL_COMMANDS;
+use super::{highlighter::ReplHighlighter, prompt::ReplPrompt, REPL_COMMANDS};
 
 use crate::config::{Config, SharedConfig};
 
 use anyhow::{Context, Result};
-use nu_ansi_term::Color;
 use reedline::{
     default_emacs_keybindings, ColumnarMenu, DefaultCompleter, DefaultValidator, Emacs,
-    ExampleHighlighter, FileBackedHistory, KeyCode, KeyModifiers, Keybindings, Reedline,
-    ReedlineEvent, ReedlineMenu,
+    FileBackedHistory, KeyCode, KeyModifiers, Keybindings, Reedline, ReedlineEvent, ReedlineMenu,
 };
 
 const MENU_NAME: &str = "completion_menu";
-const MATCH_COLOR: Color = Color::Green;
-const NEUTRAL_COLOR: Color = Color::White;
-const NEUTRAL_COLOR_LIGHT: Color = Color::Black;
 
 pub struct Repl {
-    pub editor: Reedline,
+    pub(crate) editor: Reedline,
+    pub(crate) prompt: ReplPrompt,
 }
 
 impl Repl {
@@ -27,7 +23,7 @@ impl Repl {
             .collect();
 
         let completer = Self::create_completer(config.clone(), &commands);
-        let highlighter = Self::create_highlighter(config, &commands);
+        let highlighter = ReplHighlighter::new(config.clone(), commands);
         let keybindings = Self::create_keybindings();
         let history = Self::create_history()?;
         let menu = Self::create_menu();
@@ -42,27 +38,18 @@ impl Repl {
             .with_partial_completions(true)
             .with_validator(Box::new(DefaultValidator))
             .with_ansi_colors(true);
-        Ok(Self { editor })
+        let prompt = ReplPrompt::new(config);
+        Ok(Self { editor, prompt })
     }
 
     fn create_completer(config: SharedConfig, commands: &[String]) -> DefaultCompleter {
         let mut completion = commands.to_vec();
-        completion.extend(config.lock().repl_completions());
+        completion.extend(config.read().repl_completions());
         let mut completer = DefaultCompleter::with_inclusions(&['.', '-', '_']).set_min_word_len(2);
         completer.insert(completion.clone());
         completer
     }
 
-    fn create_highlighter(config: SharedConfig, commands: &[String]) -> ExampleHighlighter {
-        let mut highlighter = ExampleHighlighter::new(commands.to_vec());
-        if config.lock().light_theme {
-            highlighter.change_colors(MATCH_COLOR, NEUTRAL_COLOR_LIGHT, NEUTRAL_COLOR_LIGHT);
-        } else {
-            highlighter.change_colors(MATCH_COLOR, NEUTRAL_COLOR, NEUTRAL_COLOR);
-        }
-        highlighter
-    }
-
     fn create_keybindings() -> Keybindings {
         let mut keybindings = default_emacs_keybindings();
         keybindings.add_binding(
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index f40a011..f44a5bf 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -1,12 +1,12 @@
 mod abort;
 mod handler;
+mod highlighter;
 mod init;
 mod prompt;
 
 pub use self::abort::*;
 pub use self::handler::*;
 pub use self::init::Repl;
-use self::prompt::ReplPrompt;
 
 use crate::client::ChatGptClient;
 use crate::config::SharedConfig;
@@ -35,8 +35,7 @@ pub const REPL_COMMANDS: [(&str, &str); 11] = [
 impl Repl {
     pub fn run(&mut self, client: ChatGptClient, config: SharedConfig) -> Result<()> {
         let abort = AbortSignal::new();
-        let handler = ReplCmdHandler::init(client, config.clone(), abort.clone())?;
-        let prompt = ReplPrompt::new(config);
+        let handler = ReplCmdHandler::init(client, config, abort.clone())?;
         print_now!("Welcome to aichat {}\n", env!("CARGO_PKG_VERSION"));
         print_now!("Type \".help\" for more information.\n");
         let mut already_ctrlc = false;
@@ -48,7 +47,7 @@ impl Repl {
             if abort.aborted_ctrlc() && !already_ctrlc {
                 already_ctrlc = true;
             }
-            let sig = self.editor.read_line(&prompt);
+            let sig = self.editor.read_line(&self.prompt);
             match sig {
                 Ok(Signal::Success(line)) => {
                     already_ctrlc = false;
@@ -117,7 +116,8 @@ impl Repl {
                     handler.handle(ReplCmd::ViewInfo)?;
                 }
                 ".set" => {
-                    handler.handle(ReplCmd::UpdateConfig(args.unwrap_or_default().to_string()))?
+                    handler.handle(ReplCmd::UpdateConfig(args.unwrap_or_default().to_string()))?;
+                    self.prompt.sync_config();
                 }
                 ".prompt" => {
                     let text = args.unwrap_or_default().to_string();
diff --git a/src/repl/prompt.rs b/src/repl/prompt.rs
index dd32ff8..1a59ba3 100644
--- a/src/repl/prompt.rs
+++ b/src/repl/prompt.rs
@@ -1,21 +1,74 @@
 use crate::config::SharedConfig;
 
+use crossterm::style::Color;
 use reedline::{Prompt, PromptHistorySearch, PromptHistorySearchStatus};
 use std::borrow::Cow;
 
+const PROMPT_COLOR: Color = Color::Green;
+const PROMPT_MULTILINE_COLOR: nu_ansi_term::Color = nu_ansi_term::Color::LightBlue;
+const INDICATOR_COLOR: Color = Color::Cyan;
+const PROMPT_RIGHT_COLOR: Color = Color::AnsiValue(5);
+
 #[derive(Clone)]
-pub struct ReplPrompt(SharedConfig);
+pub struct ReplPrompt {
+    config: SharedConfig,
+    prompt_color: Color,
+    prompt_multiline_color: nu_ansi_term::Color,
+    indicator_color: Color,
+    prompt_right_color: Color,
+}
 
 impl ReplPrompt {
     pub fn new(config: SharedConfig) -> Self {
-        Self(config)
+        let (prompt_color, prompt_multiline_color, indicator_color, prompt_right_color) =
+            Self::get_colors(config.clone());
+        Self {
+            config,
+            prompt_color,
+            prompt_multiline_color,
+            indicator_color,
+            prompt_right_color,
+        }
+    }
+    pub fn sync_config(&mut self) {
+        let (prompt_color, prompt_multiline_color, indicator_color, prompt_right_color) =
+            Self::get_colors(self.config.clone());
+        self.prompt_color = prompt_color;
+        self.prompt_multiline_color = prompt_multiline_color;
+        self.indicator_color = indicator_color;
+        self.prompt_right_color = prompt_right_color;
+    }
+
+    pub fn get_colors(config: SharedConfig) -> (Color, nu_ansi_term::Color, Color, Color) {
+        let (highlight, light_theme) = config.read().get_render_options();
+        if highlight {
+            (
+                PROMPT_COLOR,
+                PROMPT_MULTILINE_COLOR,
+                INDICATOR_COLOR,
+                PROMPT_RIGHT_COLOR,
+            )
+        } else if light_theme {
+            (
+                Color::Black,
+                nu_ansi_term::Color::Black,
+                Color::Black,
+                Color::Black,
+            )
+        } else {
+            (
+                Color::White,
+                nu_ansi_term::Color::White,
+                Color::White,
+                Color::White,
+            )
+        }
     }
 }
 
 impl Prompt for ReplPrompt {
     fn render_prompt_left(&self) -> Cow<str> {
-        let config = self.0.lock();
-        if let Some(role) = config.role.as_ref() {
+        if let Some(role) = self.config.read().role.as_ref() {
             role.name.to_string().into()
         } else {
             Cow::Borrowed("")
@@ -23,8 +76,7 @@ impl Prompt for ReplPrompt {
     }
 
     fn render_prompt_right(&self) -> Cow<str> {
-        let config = self.0.lock();
-        if let Some(conversation) = config.conversation.as_ref() {
+        if let Some(conversation) = self.config.read().conversation.as_ref() {
             conversation.reamind_tokens().to_string().into()
         } else {
             Cow::Borrowed("")
@@ -32,8 +84,7 @@ impl Prompt for ReplPrompt {
     }
 
     fn render_prompt_indicator(&self, _prompt_mode: reedline::PromptEditMode) -> Cow<str> {
-        let config = self.0.lock();
-        if config.conversation.is_some() {
+        if self.config.read().conversation.is_some() {
             Cow::Borrowed("＄")
         } else {
             Cow::Borrowed("〉")
@@ -59,4 +110,20 @@ impl Prompt for ReplPrompt {
             prefix, history_search.term
         ))
     }
+
+    fn get_prompt_color(&self) -> Color {
+        self.prompt_color
+    }
+    /// Get the default multilince prompt color
+    fn get_prompt_multiline_color(&self) -> nu_ansi_term::Color {
+        self.prompt_multiline_color
+    }
+    /// Get the default indicator color
+    fn get_indicator_color(&self) -> Color {
+        self.indicator_color
+    }
+    /// Get the default right prompt color
+    fn get_prompt_right_color(&self) -> Color {
+        self.prompt_right_color
+    }
 }