refactor: improve code quanity (#196)

- rewrite Repl, remove ReplHandler
- move ReplyStreamHandler to repl/ and rename it to ReplyHandler
- deprecate utils::print_now
- refactor session info

4 files changed, 229 insertions, 350 deletions

diff --git a/src/repl/abort.rs b/src/repl/abort.rs
index f76abb6..af58b35 100644
--- a/src/repl/abort.rs
+++ b/src/repl/abort.rs
@@ -3,15 +3,19 @@ use std::sync::{
     Arc,
 };
 
-pub type SharedAbortSignal = Arc<AbortSignal>;
+pub type AbortSignal = Arc<AbortSignalInner>;
 
-pub struct AbortSignal {
+pub struct AbortSignalInner {
     ctrlc: AtomicBool,
     ctrld: AtomicBool,
 }
 
-impl AbortSignal {
-    pub fn new() -> SharedAbortSignal {
+pub fn create_abort_signal() -> AbortSignal {
+    AbortSignalInner::new()
+}
+
+impl AbortSignalInner {
+    pub fn new() -> AbortSignal {
         Arc::new(Self {
             ctrlc: AtomicBool::new(false),
             ctrld: AtomicBool::new(false),
diff --git a/src/repl/handler.rs b/src/repl/handler.rs
deleted file mode 100644
index 94f91a1..0000000
--- a/src/repl/handler.rs
+++ /dev/null
@@ -1,205 +0,0 @@
-use crate::client::init_client;
-use crate::config::SharedConfig;
-use crate::print_now;
-use crate::render::{render_stream, MarkdownRender};
-use std::fs;
-use std::io::Read;
-
-use super::abort::SharedAbortSignal;
-
-use anyhow::{bail, Context, Result};
-use arboard::Clipboard;
-use crossbeam::channel::Sender;
-use crossbeam::sync::WaitGroup;
-use std::cell::RefCell;
-
-pub enum ReplCmd {
-    Submit(String),
-    Info,
-    RoleInfo,
-    SessionInfo,
-    SetModel(String),
-    SetRole(String),
-    ExitRole,
-    StartSession(Option<String>),
-    ExitSession,
-    Set(String),
-    Copy,
-    ReadFile(String),
-}
-
-pub struct ReplCmdHandler {
-    config: SharedConfig,
-    abort: SharedAbortSignal,
-    clipboard: std::result::Result<RefCell<Clipboard>, arboard::Error>,
-}
-
-impl ReplCmdHandler {
-    pub fn init(config: SharedConfig, abort: SharedAbortSignal) -> Result<Self> {
-        let clipboard = Clipboard::new().map(RefCell::new);
-        Ok(Self {
-            config,
-            abort,
-            clipboard,
-        })
-    }
-
-    pub fn handle(&self, cmd: ReplCmd) -> Result<()> {
-        match cmd {
-            ReplCmd::Submit(input) => {
-                if input.is_empty() {
-                    return Ok(());
-                }
-                self.config.read().maybe_print_send_tokens(&input);
-                let wg = WaitGroup::new();
-                let client = init_client(self.config.clone())?;
-                let ret = render_stream(
-                    &input,
-                    client.as_ref(),
-                    &self.config,
-                    true,
-                    self.abort.clone(),
-                    wg.clone(),
-                );
-                wg.wait();
-                let buffer = ret?;
-                self.config.write().save_message(&input, &buffer)?;
-                if self.config.read().auto_copy {
-                    let _ = self.copy(&buffer);
-                }
-            }
-            ReplCmd::Info => {
-                let output = self.config.read().info()?;
-                print_now!("{}\n\n", output.trim_end());
-            }
-            ReplCmd::SetModel(name) => {
-                self.config.write().set_model(&name)?;
-                print_now!("\n");
-            }
-            ReplCmd::SetRole(name) => {
-                self.config.write().set_role(&name)?;
-                print_now!("\n");
-            }
-            ReplCmd::RoleInfo => {
-                if let Some(role) = &self.config.read().role {
-                    print_now!("{}\n\n", role.info()?);
-                } else {
-                    bail!("No role")
-                }
-            }
-            ReplCmd::ExitRole => {
-                self.config.write().clear_role()?;
-                print_now!("\n");
-            }
-            ReplCmd::StartSession(name) => {
-                self.config.write().start_session(&name)?;
-                print_now!("\n");
-            }
-            ReplCmd::SessionInfo => {
-                if let Some(session) = &self.config.read().session {
-                    let render_options = self.config.read().get_render_options()?;
-                    let mut markdown_render = MarkdownRender::init(render_options)?;
-                    print_now!("{}\n\n", session.render(&mut markdown_render)?);
-                } else {
-                    bail!("No session")
-                }
-            }
-            ReplCmd::ExitSession => {
-                self.config.write().end_session()?;
-                print_now!("\n");
-            }
-            ReplCmd::Set(input) => {
-                self.config.write().update(&input)?;
-                print_now!("\n");
-            }
-            ReplCmd::Copy => {
-                let reply = self
-                    .config
-                    .read()
-                    .last_message
-                    .as_ref()
-                    .map(|v| v.1.clone())
-                    .unwrap_or_default();
-                self.copy(&reply)
-                    .with_context(|| "Failed to copy the last output")?;
-                print_now!("\n");
-            }
-            ReplCmd::ReadFile(file) => {
-                let mut contents = String::new();
-                let mut file = fs::File::open(file).with_context(|| "Unable to open file")?;
-                file.read_to_string(&mut contents)
-                    .with_context(|| "Unable to read file")?;
-                self.handle(ReplCmd::Submit(contents))?;
-            }
-        }
-        Ok(())
-    }
-
-    fn copy(&self, text: &str) -> Result<()> {
-        match self.clipboard.as_ref() {
-            Err(err) => bail!("{}", err),
-            Ok(clip) => {
-                clip.borrow_mut().set_text(text)?;
-                Ok(())
-            }
-        }
-    }
-}
-
-pub struct ReplyStreamHandler {
-    sender: Sender<ReplyStreamEvent>,
-    buffer: String,
-    abort: SharedAbortSignal,
-}
-
-impl ReplyStreamHandler {
-    pub fn new(sender: Sender<ReplyStreamEvent>, abort: SharedAbortSignal) -> Self {
-        Self {
-            sender,
-            abort,
-            buffer: String::new(),
-        }
-    }
-
-    pub fn text(&mut self, text: &str) -> Result<()> {
-        if self.buffer.is_empty() && text == "\n\n" {
-            return Ok(());
-        }
-        self.buffer.push_str(text);
-        let ret = self
-            .sender
-            .send(ReplyStreamEvent::Text(text.to_string()))
-            .with_context(|| "Failed to send StreamEvent:Text");
-        self.safe_ret(ret)?;
-        Ok(())
-    }
-
-    pub fn done(&mut self) -> Result<()> {
-        let ret = self
-            .sender
-            .send(ReplyStreamEvent::Done)
-            .with_context(|| "Failed to send StreamEvent:Done");
-        self.safe_ret(ret)?;
-        Ok(())
-    }
-
-    pub fn get_buffer(&self) -> &str {
-        &self.buffer
-    }
-
-    pub fn get_abort(&self) -> SharedAbortSignal {
-        self.abort.clone()
-    }
-
-    fn safe_ret(&self, ret: Result<()>) -> Result<()> {
-        if ret.is_err() && self.abort.aborted() {
-            return Ok(());
-        }
-        ret
-    }
-}
-
-pub enum ReplyStreamEvent {
-    Text(String),
-    Done,
-}
diff --git a/src/repl/init.rs b/src/repl/init.rs
deleted file mode 100644
index fbbcf22..0000000
--- a/src/repl/init.rs
+++ /dev/null
@@ -1,85 +0,0 @@
-use super::{
-    highlighter::ReplHighlighter, prompt::ReplPrompt, validator::ReplValidator, REPL_COMMANDS,
-};
-
-use crate::config::SharedConfig;
-
-use anyhow::Result;
-use reedline::{
-    default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
-    ColumnarMenu, DefaultCompleter, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
-    ReedlineEvent, ReedlineMenu, Vi,
-};
-
-const MENU_NAME: &str = "completion_menu";
-
-pub struct Repl {
-    pub(crate) editor: Reedline,
-    pub(crate) prompt: ReplPrompt,
-}
-
-impl Repl {
-    pub fn init(config: SharedConfig) -> Result<Self> {
-        let commands: Vec<String> = REPL_COMMANDS
-            .into_iter()
-            .map(|(v, _)| v.to_string())
-            .collect();
-
-        let completer = Self::create_completer(&config, &commands);
-        let highlighter = ReplHighlighter::new(commands, config.clone());
-        let menu = Self::create_menu();
-        let edit_mode: Box<dyn EditMode> = if config.read().keybindings.is_vi() {
-            let mut normal_keybindings = default_vi_normal_keybindings();
-            let mut insert_keybindings = default_vi_insert_keybindings();
-            Self::extra_keybindings(&mut normal_keybindings);
-            Self::extra_keybindings(&mut insert_keybindings);
-            Box::new(Vi::new(insert_keybindings, normal_keybindings))
-        } else {
-            let mut keybindings = default_emacs_keybindings();
-            Self::extra_keybindings(&mut keybindings);
-            Box::new(Emacs::new(keybindings))
-        };
-        let mut editor = Reedline::create()
-            .with_completer(Box::new(completer))
-            .with_highlighter(Box::new(highlighter))
-            .with_menu(menu)
-            .with_edit_mode(edit_mode)
-            .with_quick_completions(true)
-            .with_partial_completions(true)
-            .with_validator(Box::new(ReplValidator))
-            .with_ansi_colors(true);
-        editor.enable_bracketed_paste()?;
-        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.read().repl_completions());
-        let mut completer =
-            DefaultCompleter::with_inclusions(&['.', '-', '_', ':']).set_min_word_len(2);
-        completer.insert(completion.clone());
-        completer
-    }
-
-    fn extra_keybindings(keybindings: &mut Keybindings) {
-        keybindings.add_binding(
-            KeyModifiers::NONE,
-            KeyCode::Tab,
-            ReedlineEvent::UntilFound(vec![
-                ReedlineEvent::Menu(MENU_NAME.to_string()),
-                ReedlineEvent::MenuNext,
-            ]),
-        );
-        keybindings.add_binding(
-            KeyModifiers::CONTROL,
-            KeyCode::Char('s'),
-            ReedlineEvent::Submit,
-        );
-    }
-
-    fn create_menu() -> ReedlineMenu {
-        let completion_menu = ColumnarMenu::default().with_name(MENU_NAME);
-        ReedlineMenu::EngineCompleter(Box::new(completion_menu))
-    }
-}
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index eb99fda..c9c3c73 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -1,24 +1,35 @@
 mod abort;
-mod handler;
 mod highlighter;
-mod init;
 mod prompt;
 mod validator;
 
-pub use self::abort::*;
-pub use self::handler::*;
-pub use self::init::Repl;
+pub use self::abort::{create_abort_signal, AbortSignal};
 
+use self::highlighter::ReplHighlighter;
+use self::prompt::ReplPrompt;
+use self::validator::ReplValidator;
+
+use crate::client::init_client;
 use crate::config::SharedConfig;
-use crate::print_now;
+use crate::render::{render_error, render_stream};
 
-use anyhow::Result;
+use anyhow::{bail, Context, Result};
+use arboard::Clipboard;
+use crossbeam::sync::WaitGroup;
 use fancy_regex::Regex;
 use lazy_static::lazy_static;
 use reedline::Signal;
-use std::rc::Rc;
+use reedline::{
+    default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
+    ColumnarMenu, DefaultCompleter, EditMode, Emacs, KeyCode, KeyModifiers, Keybindings, Reedline,
+    ReedlineEvent, ReedlineMenu, Vi,
+};
+use std::cell::RefCell;
+use std::io::Read;
+
+const MENU_NAME: &str = "completion_menu";
 
-pub const REPL_COMMANDS: [(&str, &str); 14] = [
+const REPL_COMMANDS: [(&str, &str); 14] = [
     (".help", "Print this help message"),
     (".info", "Print system info"),
     (".edit", "Multi-line editing (CTRL+S to finish)"),
@@ -40,128 +51,284 @@ lazy_static! {
     static ref EDIT_RE: Regex = Regex::new(r"^\s*\.edit\s*").unwrap();
 }
 
+pub struct Repl {
+    config: SharedConfig,
+    editor: Reedline,
+    prompt: ReplPrompt,
+    abort: AbortSignal,
+    clipboard: std::result::Result<RefCell<Clipboard>, arboard::Error>,
+}
+
 impl Repl {
-    pub fn run(&mut self, config: SharedConfig) -> Result<()> {
-        let abort = AbortSignal::new();
-        let handler = ReplCmdHandler::init(config, abort.clone())?;
-        print_now!("Welcome to aichat {}\n", env!("CARGO_PKG_VERSION"));
-        print_now!("Type \".help\" for more information.\n");
+    pub fn init(config: SharedConfig) -> Result<Self> {
+        let commands: Vec<String> = REPL_COMMANDS
+            .into_iter()
+            .map(|(v, _)| v.to_string())
+            .collect();
+
+        let completer = Self::create_completer(&config, &commands);
+        let highlighter = ReplHighlighter::new(commands, config.clone());
+        let menu = Self::create_menu();
+        let edit_mode: Box<dyn EditMode> = if config.read().keybindings.is_vi() {
+            let mut normal_keybindings = default_vi_normal_keybindings();
+            let mut insert_keybindings = default_vi_insert_keybindings();
+            Self::extra_keybindings(&mut normal_keybindings);
+            Self::extra_keybindings(&mut insert_keybindings);
+            Box::new(Vi::new(insert_keybindings, normal_keybindings))
+        } else {
+            let mut keybindings = default_emacs_keybindings();
+            Self::extra_keybindings(&mut keybindings);
+            Box::new(Emacs::new(keybindings))
+        };
+        let mut editor = Reedline::create()
+            .with_completer(Box::new(completer))
+            .with_highlighter(Box::new(highlighter))
+            .with_menu(menu)
+            .with_edit_mode(edit_mode)
+            .with_quick_completions(true)
+            .with_partial_completions(true)
+            .with_validator(Box::new(ReplValidator))
+            .with_ansi_colors(true);
+
+        editor.enable_bracketed_paste()?;
+
+        let prompt = ReplPrompt::new(config.clone());
+
+        let abort = create_abort_signal();
+
+        let clipboard = Clipboard::new().map(RefCell::new);
+
+        Ok(Self {
+            config,
+            editor,
+            prompt,
+            clipboard,
+            abort,
+        })
+    }
+
+    pub fn run(&mut self) -> Result<()> {
+        self.banner();
+
         let mut already_ctrlc = false;
-        let handler = Rc::new(handler);
+
         loop {
-            if abort.aborted_ctrld() {
+            if self.abort.aborted_ctrld() {
                 break;
             }
-            if abort.aborted_ctrlc() && !already_ctrlc {
+            if self.abort.aborted_ctrlc() && !already_ctrlc {
                 already_ctrlc = true;
             }
             let sig = self.editor.read_line(&self.prompt);
             match sig {
                 Ok(Signal::Success(line)) => {
                     already_ctrlc = false;
-                    abort.reset();
-                    match self.handle_line(&handler, &line) {
+                    self.abort.reset();
+                    match self.handle(&line) {
                         Ok(quit) => {
                             if quit {
                                 break;
                             }
                         }
                         Err(err) => {
-                            let err = format!("{err:?}");
-                            print_now!("Error: {}\n\n", err.trim());
+                            render_error(err, self.config.read().highlight);
                         }
                     }
                 }
                 Ok(Signal::CtrlC) => {
-                    abort.set_ctrlc();
+                    self.abort.set_ctrlc();
                     if already_ctrlc {
                         break;
                     }
                     already_ctrlc = true;
-                    print_now!("(To exit, press Ctrl+C again or Ctrl+D or type .exit)\n\n");
+                    println!("(To exit, press Ctrl+C again or Ctrl+D or type .exit)\n");
                 }
                 Ok(Signal::CtrlD) => {
-                    abort.set_ctrld();
+                    self.abort.set_ctrld();
                     break;
                 }
                 _ => {}
             }
         }
-        handler.handle(ReplCmd::ExitSession)?;
+        self.handle(".exit session")?;
         Ok(())
     }
 
-    fn handle_line(&mut self, handler: &Rc<ReplCmdHandler>, line: &str) -> Result<bool> {
+    fn handle(&self, line: &str) -> Result<bool> {
         match parse_command(line) {
             Some((cmd, args)) => match cmd {
                 ".help" => {
                     dump_repl_help();
                 }
                 ".info" => match args {
-                    Some("role") => handler.handle(ReplCmd::RoleInfo)?,
-                    Some("session") => handler.handle(ReplCmd::SessionInfo)?,
-                    Some(_) => unknown_command(),
+                    Some("role") => {
+                        let info = self.config.read().role_info()?;
+                        println!("{}", info);
+                    }
+                    Some("session") => {
+                        let info = self.config.read().session_info()?;
+                        println!("{}", info);
+                    }
+                    Some(_) => unknown_command()?,
                     None => {
-                        handler.handle(ReplCmd::Info)?;
+                        let output = self.config.read().sys_info()?;
+                        println!("{}", output);
                     }
                 },
                 ".edit" => {
                     if let Some(text) = args {
-                        handler.handle(ReplCmd::Submit(text.to_string()))?;
+                        self.ask(text)?;
                     }
                 }
                 ".model" => match args {
-                    Some(name) => handler.handle(ReplCmd::SetModel(name.to_string()))?,
-                    None => print_now!("Usage: .model <name>\n\n"),
+                    Some(name) => {
+                        self.config.write().set_model(name)?;
+                    }
+                    None => println!("Usage: .model <name>"),
                 },
                 ".role" => match args {
-                    Some(name) => handler.handle(ReplCmd::SetRole(name.to_string()))?,
-                    None => print_now!("Usage: .role <name>\n\n"),
+                    Some(name) => {
+                        self.config.write().set_role(name)?;
+                    }
+                    None => println!("Usage: .role <name>"),
                 },
                 ".session" => {
-                    handler.handle(ReplCmd::StartSession(args.map(|v| v.to_string())))?;
+                    self.config.write().start_session(args)?;
                 }
                 ".set" => {
-                    handler.handle(ReplCmd::Set(args.unwrap_or_default().to_string()))?;
+                    if let Some(args) = args {
+                        self.config.write().update(args)?;
+                    }
                 }
                 ".copy" => {
-                    handler.handle(ReplCmd::Copy)?;
+                    let config = self.config.read();
+                    self.copy(config.last_reply())
+                        .with_context(|| "Failed to copy the last output")?;
                 }
                 ".read" => match args {
-                    Some(file) => handler.handle(ReplCmd::ReadFile(file.to_string()))?,
-                    None => print_now!("Usage: .read <file name>\n\n"),
+                    Some(file) => {
+                        let mut content = String::new();
+                        let mut file =
+                            std::fs::File::open(file).with_context(|| "Unable to open file")?;
+                        file.read_to_string(&mut content)
+                            .with_context(|| "Unable to read file")?;
+                        self.ask(&content)?;
+                    }
+                    None => println!("Usage: .read <textfile>"),
                 },
                 ".exit" => match args {
-                    Some("role") => handler.handle(ReplCmd::ExitRole)?,
-                    Some("session") => handler.handle(ReplCmd::ExitSession)?,
-                    Some(_) => unknown_command(),
+                    Some("role") => {
+                        self.config.write().clear_role()?;
+                    }
+                    Some("session") => {
+                        self.config.write().end_session()?;
+                    }
+                    Some(_) => unknown_command()?,
                     None => {
                         return Ok(true);
                     }
                 },
-                // deprecated
+                // deprecated this command
                 ".clear" => match args {
                     Some("role") => {
-                        print_now!("Deprecated. Use '.exit role' instead.\n\n");
+                        println!(r#"Deprecated. Use ".exit role" instead."#);
                     }
-                    Some("session") => {
-                        print_now!("Deprecated. Use '.exit session' instead.\n\n");
+                    Some("conversation") => {
+                        println!(r#"Deprecated. Use ".exit session" instead."#);
                     }
-                    _ => unknown_command(),
+                    _ => unknown_command()?,
                 },
-                _ => unknown_command(),
+                _ => unknown_command()?,
             },
             None => {
-                handler.handle(ReplCmd::Submit(line.to_string()))?;
+                self.ask(line)?;
             }
         }
 
+        println!();
+
         Ok(false)
     }
+
+    fn ask(&self, input: &str) -> Result<()> {
+        if input.is_empty() {
+            return Ok(());
+        }
+        self.config.read().maybe_print_send_tokens(input);
+        let wg = WaitGroup::new();
+        let client = init_client(self.config.clone())?;
+        let ret = render_stream(
+            input,
+            client.as_ref(),
+            &self.config,
+            true,
+            self.abort.clone(),
+            wg.clone(),
+        );
+        wg.wait();
+        let buffer = ret?;
+        self.config.write().save_message(input, &buffer)?;
+        if self.config.read().auto_copy {
+            let _ = self.copy(&buffer);
+        }
+        Ok(())
+    }
+
+    fn banner(&self) {
+        let version = env!("CARGO_PKG_VERSION");
+        print!(
+            r#"Welcome to aichat {version}
+Type ".help" for more information.
+"#
+        )
+    }
+
+    fn create_completer(config: &SharedConfig, commands: &[String]) -> DefaultCompleter {
+        let mut completion = commands.to_vec();
+        completion.extend(config.read().repl_completions());
+        let mut completer =
+            DefaultCompleter::with_inclusions(&['.', '-', '_', ':']).set_min_word_len(2);
+        completer.insert(completion.clone());
+        completer
+    }
+
+    fn extra_keybindings(keybindings: &mut Keybindings) {
+        keybindings.add_binding(
+            KeyModifiers::NONE,
+            KeyCode::Tab,
+            ReedlineEvent::UntilFound(vec![
+                ReedlineEvent::Menu(MENU_NAME.to_string()),
+                ReedlineEvent::MenuNext,
+            ]),
+        );
+        keybindings.add_binding(
+            KeyModifiers::CONTROL,
+            KeyCode::Char('s'),
+            ReedlineEvent::Submit,
+        );
+    }
+
+    fn create_menu() -> ReedlineMenu {
+        let completion_menu = ColumnarMenu::default().with_name(MENU_NAME);
+        ReedlineMenu::EngineCompleter(Box::new(completion_menu))
+    }
+
+    fn copy(&self, text: &str) -> Result<()> {
+        if text.is_empty() {
+            bail!("No text")
+        }
+        match self.clipboard.as_ref() {
+            Err(err) => bail!("{}", err),
+            Ok(clip) => {
+                clip.borrow_mut().set_text(text)?;
+                Ok(())
+            }
+        }
+    }
 }
 
-fn unknown_command() {
-    print_now!("Unknown command. Try `.help`.\n\n");
+fn unknown_command() -> Result<()> {
+    bail!(r#"Unknown command. Type ".help" for more information."#);
 }
 
 fn dump_repl_help() {
@@ -170,12 +337,10 @@ fn dump_repl_help() {
         .map(|(name, desc)| format!("{name:<24} {desc}"))
         .collect::<Vec<String>>()
         .join("\n");
-    print_now!(
+    println!(
         r###"{head}
 
-Press Ctrl+C to abort readline, Ctrl+D to exit the REPL
-
-"###,
+Press Ctrl+C to abort readline, Ctrl+D to exit the REPL"###,
     );
 }