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

3 files changed, 93 insertions, 26 deletions

diff --git a/src/render/cmd.rs b/src/render/cmd.rs
index 115170b..e09bd71 100644
--- a/src/render/cmd.rs
+++ b/src/render/cmd.rs
@@ -1,7 +1,6 @@
-use super::MarkdownRender;
+use super::{MarkdownRender, ReplyEvent};
 
-use crate::print_now;
-use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
+use crate::repl::AbortSignal;
 use crate::utils::{spaces, split_line_sematic, split_line_tail};
 
 use anyhow::Result;
@@ -9,9 +8,9 @@ use crossbeam::channel::Receiver;
 use textwrap::core::display_width;
 
 pub fn cmd_render_stream(
-    rx: &Receiver<ReplyStreamEvent>,
+    rx: &Receiver<ReplyEvent>,
     render: &mut MarkdownRender,
-    abort: &SharedAbortSignal,
+    abort: &AbortSignal,
 ) -> Result<()> {
     let mut buffer = String::new();
     let mut col = 0;
@@ -21,14 +20,14 @@ pub fn cmd_render_stream(
         }
         if let Ok(evt) = rx.try_recv() {
             match evt {
-                ReplyStreamEvent::Text(text) => {
+                ReplyEvent::Text(text) => {
                     if text.contains('\n') {
                         let text = format!("{buffer}{text}");
                         let (head, tail) = split_line_tail(&text);
                         buffer = tail.to_string();
                         let input = format!("{}{head}", spaces(col));
                         let output = render.render(&input);
-                        print_now!("{}\n", &output[col..]);
+                        println!("{}", &output[col..]);
                         col = 0;
                     } else {
                         buffer = format!("{buffer}{text}");
@@ -51,14 +50,14 @@ pub fn cmd_render_stream(
                                         col += display_width(output);
                                     }
                                 }
-                                print_now!("{}", output);
+                                print!("{}", output);
                             }
                         }
                     }
                 }
-                ReplyStreamEvent::Done => {
+                ReplyEvent::Done => {
                     let input = format!("{}{buffer}", spaces(col));
-                    print_now!("{}\n", render.render(&input));
+                    println!("{}", render.render(&input));
                     break;
                 }
             }
diff --git a/src/render/mod.rs b/src/render/mod.rs
index 03b67ea..27f9e81 100644
--- a/src/render/mod.rs
+++ b/src/render/mod.rs
@@ -8,12 +8,12 @@ use self::repl::repl_render_stream;
 
 use crate::client::Client;
 use crate::config::SharedConfig;
-use crate::print_now;
-use crate::repl::{ReplyStreamHandler, SharedAbortSignal};
+use crate::repl::AbortSignal;
 
-use anyhow::Result;
-use crossbeam::channel::unbounded;
+use anyhow::{Context, Result};
+use crossbeam::channel::{unbounded, Sender};
 use crossbeam::sync::WaitGroup;
+use nu_ansi_term::{Color, Style};
 use std::thread::spawn;
 
 pub fn render_stream(
@@ -21,13 +21,14 @@ pub fn render_stream(
     client: &dyn Client,
     config: &SharedConfig,
     repl: bool,
-    abort: SharedAbortSignal,
+    abort: AbortSignal,
     wg: WaitGroup,
 ) -> Result<String> {
     let render_options = config.read().get_render_options()?;
     let mut stream_handler = {
         let (tx, rx) = unbounded();
         let abort_clone = abort.clone();
+        let highlight = config.read().highlight;
         spawn(move || {
             let run = move || {
                 if repl {
@@ -39,14 +40,81 @@ pub fn render_stream(
                 }
             };
             if let Err(err) = run() {
-                let err = format!("{err:?}");
-                print_now!("\n{}\n\n", err.trim());
+                render_error(err, highlight);
             }
             drop(wg);
         });
-        ReplyStreamHandler::new(tx, abort_clone)
+        ReplyHandler::new(tx, abort_clone)
     };
     client.send_message_streaming(input, &mut stream_handler)?;
     let buffer = stream_handler.get_buffer();
     Ok(buffer.to_string())
 }
+
+pub fn render_error(err: anyhow::Error, highlight: bool) {
+    let err = format!("{err:?}\n");
+    if highlight {
+        let style = Style::new().fg(Color::Red);
+        println!("{}", style.paint(err.trim()));
+    } else {
+        println!("{}", err.trim());
+    }
+}
+
+pub struct ReplyHandler {
+    sender: Sender<ReplyEvent>,
+    buffer: String,
+    abort: AbortSignal,
+}
+
+impl ReplyHandler {
+    pub fn new(sender: Sender<ReplyEvent>, abort: AbortSignal) -> 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(ReplyEvent::Text(text.to_string()))
+            .with_context(|| "Failed to send ReplyEvent:Text");
+        self.safe_ret(ret)?;
+        Ok(())
+    }
+
+    pub fn done(&mut self) -> Result<()> {
+        let ret = self
+            .sender
+            .send(ReplyEvent::Done)
+            .with_context(|| "Failed to send ReplyEvent::Done");
+        self.safe_ret(ret)?;
+        Ok(())
+    }
+
+    pub fn get_buffer(&self) -> &str {
+        &self.buffer
+    }
+
+    pub fn get_abort(&self) -> AbortSignal {
+        self.abort.clone()
+    }
+
+    fn safe_ret(&self, ret: Result<()>) -> Result<()> {
+        if ret.is_err() && self.abort.aborted() {
+            return Ok(());
+        }
+        ret
+    }
+}
+
+pub enum ReplyEvent {
+    Text(String),
+    Done,
+}
diff --git a/src/render/repl.rs b/src/render/repl.rs
index 7363923..efad0e7 100644
--- a/src/render/repl.rs
+++ b/src/render/repl.rs
@@ -1,6 +1,6 @@
-use super::MarkdownRender;
+use super::{MarkdownRender, ReplyEvent};
 
-use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
+use crate::repl::AbortSignal;
 use crate::utils::split_line_tail;
 
 use anyhow::Result;
@@ -18,9 +18,9 @@ use std::{
 use textwrap::core::display_width;
 
 pub fn repl_render_stream(
-    rx: &Receiver<ReplyStreamEvent>,
+    rx: &Receiver<ReplyEvent>,
     render: &mut MarkdownRender,
-    abort: &SharedAbortSignal,
+    abort: &AbortSignal,
 ) -> Result<()> {
     enable_raw_mode()?;
     let mut stdout = io::stdout();
@@ -33,9 +33,9 @@ pub fn repl_render_stream(
 }
 
 fn repl_render_stream_inner(
-    rx: &Receiver<ReplyStreamEvent>,
+    rx: &Receiver<ReplyEvent>,
     render: &mut MarkdownRender,
-    abort: &SharedAbortSignal,
+    abort: &AbortSignal,
     writer: &mut Stdout,
 ) -> Result<()> {
     let mut last_tick = Instant::now();
@@ -51,7 +51,7 @@ fn repl_render_stream_inner(
 
         if let Ok(evt) = rx.try_recv() {
             match evt {
-                ReplyStreamEvent::Text(text) => {
+                ReplyEvent::Text(text) => {
                     let (col, mut row) = cursor::position()?;
 
                     // fix unexpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105
@@ -95,7 +95,7 @@ fn repl_render_stream_inner(
 
                     writer.flush()?;
                 }
-                ReplyStreamEvent::Done => {
+                ReplyEvent::Done => {
                     #[cfg(target_os = "windows")]
                     let eol = "\n\n";
                     #[cfg(not(target_os = "windows"))]