feat: add support for NO_COLOR (#67)

10 files changed, 204 insertions, 80 deletions

diff --git a/src/cli.rs b/src/cli.rs
index e08e9cf..cfc1476 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -6,7 +6,7 @@ pub struct Cli {
     /// Add a GPT prompt
     #[clap(short, long)]
     pub prompt: Option<String>,
-    /// Turn off highlight
+    /// Disable syntax highlightiing
     #[clap(short = 'H', long)]
     pub no_highlight: bool,
     /// No stream output
diff --git a/src/client.rs b/src/client.rs
index 29dc592..7fc2824 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -69,8 +69,8 @@ impl ChatGptClient {
     }
 
     async fn send_message_inner(&self, content: &str) -> Result<String> {
-        if self.config.lock().dry_run {
-            return Ok(self.config.lock().echo_messages(content));
+        if self.config.read().dry_run {
+            return Ok(self.config.read().echo_messages(content));
         }
         let builder = self.request_builder(content, false)?;
         let data: Value = builder.send().await?.json().await?;
@@ -90,8 +90,8 @@ impl ChatGptClient {
         content: &str,
         handler: &mut ReplyStreamHandler,
     ) -> Result<()> {
-        if self.config.lock().dry_run {
-            handler.text(&self.config.lock().echo_messages(content))?;
+        if self.config.read().dry_run {
+            handler.text(&self.config.read().echo_messages(content))?;
             return Ok(());
         }
         let builder = self.request_builder(content, true)?;
@@ -125,7 +125,7 @@ impl ChatGptClient {
 
     fn build_client(&self) -> Result<Client> {
         let mut builder = Client::builder();
-        if let Some(proxy) = self.config.lock().proxy.as_ref() {
+        if let Some(proxy) = self.config.read().proxy.as_ref() {
             builder = builder.proxy(Proxy::all(proxy).with_context(|| "Invalid config.proxy")?);
         }
         let client = builder
@@ -136,13 +136,13 @@ impl ChatGptClient {
     }
 
     fn request_builder(&self, content: &str, stream: bool) -> Result<RequestBuilder> {
-        let messages = self.config.lock().build_messages(content)?;
+        let messages = self.config.read().build_messages(content)?;
         let mut body = json!({
             "model": MODEL,
             "messages": messages,
         });
 
-        if let Some(v) = self.config.lock().get_temperature() {
+        if let Some(v) = self.config.read().get_temperature() {
             body.as_object_mut()
                 .and_then(|m| m.insert("temperature".into(), json!(v)));
         }
@@ -155,7 +155,7 @@ impl ChatGptClient {
         let builder = self
             .build_client()?
             .post(API_URL)
-            .bearer_auth(self.config.lock().get_api_key())
+            .bearer_auth(self.config.read().get_api_key())
             .json(&body);
 
         Ok(builder)
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 96f19db..b0cd6c0 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -10,7 +10,7 @@ use crate::utils::now;
 
 use anyhow::{anyhow, bail, Context, Result};
 use inquire::{Confirm, Text};
-use parking_lot::Mutex;
+use parking_lot::RwLock;
 use serde::Deserialize;
 use std::{
     env,
@@ -85,7 +85,7 @@ impl Default for Config {
     }
 }
 
-pub type SharedConfig = Arc<Mutex<Config>>;
+pub type SharedConfig = Arc<RwLock<Config>>;
 
 impl Config {
     pub fn init(is_interactive: bool) -> Result<Self> {
@@ -284,6 +284,7 @@ impl Config {
             ("highlight", self.highlight.to_string()),
             ("proxy", proxy),
             ("conversation_first", self.conversation_first.to_string()),
+            ("light_theme", self.light_theme.to_string()),
             ("dry_run", self.dry_run.to_string()),
         ];
         let mut output = String::new();
@@ -378,6 +379,10 @@ impl Config {
         Ok(())
     }
 
+    pub fn get_render_options(&self) -> (bool, bool) {
+        (self.highlight, self.light_theme)
+    }
+
     fn open_message_file(&self) -> Result<File> {
         let path = Config::messages_file()?;
         ensure_parent_exists(&path)?;
@@ -417,6 +422,13 @@ impl Config {
         if let Ok(value) = env::var(get_env_name("light_theme")) {
             set_bool(&mut self.light_theme, &value);
         }
+        if let Ok(value) = env::var("NO_COLOR") {
+            let mut no_color = false;
+            set_bool(&mut no_color, &value);
+            if no_color {
+                self.highlight = false;
+            }
+        }
     }
 }
 
diff --git a/src/main.rs b/src/main.rs
index fa0e64e..ffbe5d9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,7 +15,7 @@ use anyhow::{anyhow, Result};
 use clap::Parser;
 use crossbeam::sync::WaitGroup;
 use is_terminal::IsTerminal;
-use parking_lot::Mutex;
+use parking_lot::RwLock;
 use render::{render_stream, MarkdownRender};
 use repl::{AbortSignal, Repl};
 use std::io::{stdin, Read};
@@ -26,10 +26,10 @@ use utils::cl100k_base_singleton;
 fn main() -> Result<()> {
     let cli = Cli::parse();
     let text = cli.text();
-    let config = Arc::new(Mutex::new(Config::init(text.is_none())?));
+    let config = Arc::new(RwLock::new(Config::init(text.is_none())?));
     if cli.list_roles {
         config
-            .lock()
+            .read()
             .roles
             .iter()
             .for_each(|v| println!("{}", v.name));
@@ -38,18 +38,18 @@ fn main() -> Result<()> {
     let role = match &cli.role {
         Some(name) => Some(
             config
-                .lock()
+                .read()
                 .find_role(name)
                 .ok_or_else(|| anyhow!("Unknown role '{name}'"))?,
         ),
         None => None,
     };
-    config.lock().role = role;
+    config.write().role = role;
     if cli.no_highlight {
-        config.lock().highlight = false;
+        config.write().highlight = false;
     }
     if let Some(prompt) = &cli.prompt {
-        config.lock().add_prompt(prompt)?;
+        config.write().add_prompt(prompt)?;
     }
     let no_stream = cli.no_stream;
     let client = ChatGptClient::init(config.clone())?;
@@ -74,9 +74,11 @@ fn start_directive(
     input: &str,
     no_stream: bool,
 ) -> Result<()> {
-    let highlight = config.lock().highlight && stdout().is_terminal();
-    let light_theme = config.lock().light_theme;
+    if !stdout().is_terminal() {
+        config.write().highlight = false;
+    }
     let output = if no_stream {
+        let (highlight, light_theme) = config.read().get_render_options();
         let output = client.send_message(input)?;
         if highlight {
             let mut markdown_render = MarkdownRender::new(light_theme);
@@ -93,24 +95,16 @@ fn start_directive(
             abort_clone.set_ctrlc();
         })
         .expect("Error setting Ctrl-C handler");
-        let output = render_stream(
-            input,
-            &client,
-            highlight,
-            light_theme,
-            false,
-            abort,
-            wg.clone(),
-        )?;
+        let output = render_stream(input, &client, config.clone(), false, abort, wg.clone())?;
         wg.wait();
         output
     };
-    config.lock().save_message(input, &output)
+    config.read().save_message(input, &output)
 }
 
 fn start_interactive(client: ChatGptClient, config: SharedConfig) -> Result<()> {
     cl100k_base_singleton();
-    config.lock().on_repl()?;
+    config.write().on_repl()?;
     let mut repl = Repl::init(config.clone())?;
     repl.run(client, config)
 }
diff --git a/src/render/mod.rs b/src/render/mod.rs
index d200087..5980302 100644
--- a/src/render/mod.rs
+++ b/src/render/mod.rs
@@ -7,6 +7,7 @@ pub use self::markdown::MarkdownRender;
 use self::repl::repl_render_stream;
 
 use crate::client::ChatGptClient;
+use crate::config::SharedConfig;
 use crate::print_now;
 use crate::repl::{ReplyStreamHandler, SharedAbortSignal};
 
@@ -18,12 +19,12 @@ use std::thread::spawn;
 pub fn render_stream(
     input: &str,
     client: &ChatGptClient,
-    highlight: bool,
-    light_theme: bool,
+    config: SharedConfig,
     repl: bool,
     abort: SharedAbortSignal,
     wg: WaitGroup,
 ) -> Result<String> {
+    let (highlight, light_theme) = config.read().get_render_options();
     let mut stream_handler = if highlight {
         let (tx, rx) = unbounded();
         let abort_clone = abort.clone();
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
+    }
 }