chore: optimize code

5 files changed, 46 insertions, 33 deletions

diff --git a/src/client.rs b/src/client.rs
index 698fffc..3f74aa6 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -1,5 +1,5 @@
 use crate::config::SharedConfig;
-use crate::repl::ReplyReceiver;
+use crate::repl::ReplyStreamHandler;
 
 use anyhow::{anyhow, Context, Result};
 use eventsource_stream::Eventsource;
@@ -38,8 +38,7 @@ impl ChatGptClient {
         &self,
         input: &str,
         prompt: Option<String>,
-        receiver: &mut ReplyReceiver,
-        ctrlc: Arc<AtomicBool>,
+        handler: &mut ReplyStreamHandler,
     ) -> Result<()> {
         async fn watch_ctrlc(ctrlc: Arc<AtomicBool>) {
             loop {
@@ -49,14 +48,15 @@ impl ChatGptClient {
                 sleep(Duration::from_millis(100)).await;
             }
         }
+        let ctrlc = handler.get_ctrlc();
         self.runtime.block_on(async {
             tokio::select! {
-                ret = self.acquire_stream_inner(input, prompt, receiver) => {
-                    receiver.done();
+                ret = self.acquire_stream_inner(input, prompt, handler) => {
+                    handler.done();
                     ret
                 }
                 _ = watch_ctrlc(ctrlc.clone()) => {
-                    receiver.done();
+                    handler.done();
                     Ok(())
                  },
                 _ =  tokio::signal::ctrl_c() => {
@@ -86,10 +86,10 @@ impl ChatGptClient {
         &self,
         content: &str,
         prompt: Option<String>,
-        receiver: &mut ReplyReceiver,
+        handler: &mut ReplyStreamHandler,
     ) -> Result<()> {
         if self.config.borrow().dry_run {
-            receiver.text(&combine(content, prompt));
+            handler.text(&combine(content, prompt));
             return Ok(());
         }
         let builder = self.request_builder(content, prompt, true)?;
@@ -113,7 +113,7 @@ impl ChatGptClient {
                         continue;
                     }
                 }
-                receiver.text(text);
+                handler.text(text);
             }
         }
 
diff --git a/src/main.rs b/src/main.rs
index ec127f5..c7ae56e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,7 +19,7 @@ use is_terminal::IsTerminal;
 use anyhow::{anyhow, Result};
 use clap::Parser;
 use render::MarkdownRender;
-use repl::{Repl, ReplCmdHandler};
+use repl::Repl;
 
 fn main() -> Result<()> {
     let cli = Cli::parse();
@@ -73,6 +73,5 @@ fn start_directive(client: ChatGptClient, config: SharedConfig, input: &str) ->
 
 fn start_interactive(client: ChatGptClient, config: SharedConfig) -> Result<()> {
     let mut repl = Repl::init(config.clone())?;
-    let handler = ReplCmdHandler::init(client, config)?;
-    repl.run(handler)
+    repl.run(client, config)
 }
diff --git a/src/render.rs b/src/render.rs
index f44283c..57633a6 100644
--- a/src/render.rs
+++ b/src/render.rs
@@ -1,4 +1,4 @@
-use crate::{repl::RenderStreamEvent, utils::dump};
+use crate::{repl::ReplyStreamEvent, utils::dump};
 use anyhow::Result;
 use crossbeam::channel::Receiver;
 use mdcat::{
@@ -14,7 +14,7 @@ use std::sync::{
 use syntect::parsing::SyntaxSet;
 
 pub fn render_stream(
-    rx: Receiver<RenderStreamEvent>,
+    rx: Receiver<ReplyStreamEvent>,
     ctrlc: Arc<AtomicBool>,
     markdown_render: Arc<MarkdownRender>,
 ) -> Result<()> {
@@ -26,7 +26,7 @@ pub fn render_stream(
         }
         if let Ok(evt) = rx.try_recv() {
             match evt {
-                RenderStreamEvent::Text(text) => {
+                ReplyStreamEvent::Text(text) => {
                     buffer.push_str(&text);
                     if text.contains('\n') {
                         let markdown = markdown_render.render(&buffer)?;
@@ -40,7 +40,7 @@ pub fn render_stream(
                         }
                     }
                 }
-                RenderStreamEvent::Done => {
+                ReplyStreamEvent::Done => {
                     let markdown = markdown_render.render(&buffer)?;
                     let tail = markdown
                         .lines()
diff --git a/src/repl/handler.rs b/src/repl/handler.rs
index 23fa609..21c49a8 100644
--- a/src/repl/handler.rs
+++ b/src/repl/handler.rs
@@ -62,7 +62,7 @@ impl ReplCmdHandler {
                 let prompt = self.config.borrow().get_prompt();
                 let wg = WaitGroup::new();
                 let highlight = self.config.borrow().highlight;
-                let mut receiver = if highlight {
+                let mut stream_handler = if highlight {
                     let (tx, rx) = unbounded();
                     let ctrlc = self.ctrlc.clone();
                     let wg = wg.clone();
@@ -71,19 +71,20 @@ impl ReplCmdHandler {
                         let _ = render::render_stream(rx, ctrlc, render);
                         drop(wg);
                     });
-                    ReplyReceiver::new(Some(tx))
+                    ReplyStreamHandler::new(Some(tx), self.ctrlc.clone())
                 } else {
-                    ReplyReceiver::new(None)
+                    ReplyStreamHandler::new(None, self.ctrlc.clone())
                 };
                 self.client
-                    .acquire_stream(&input, prompt, &mut receiver, self.ctrlc.clone())?;
+                    .acquire_stream(&input, prompt, &mut stream_handler)?;
+                let buffer = stream_handler.get_buffer();
                 self.config.borrow().save_message(
                     self.state.borrow_mut().save_file.as_mut(),
                     &input,
-                    &receiver.output,
+                    buffer,
                 )?;
                 wg.wait();
-                self.state.borrow_mut().reply = receiver.output;
+                self.state.borrow_mut().reply = buffer.to_string();
             }
             ReplCmd::SetRole(name) => {
                 let output = self.config.borrow_mut().change_role(&name);
@@ -118,44 +119,54 @@ impl ReplCmdHandler {
     }
 }
 
-pub struct ReplyReceiver {
-    output: String,
-    sender: Option<Sender<RenderStreamEvent>>,
+pub struct ReplyStreamHandler {
+    sender: Option<Sender<ReplyStreamEvent>>,
+    buffer: String,
+    ctrlc: Arc<AtomicBool>,
 }
 
-impl ReplyReceiver {
-    pub fn new(sender: Option<Sender<RenderStreamEvent>>) -> Self {
+impl ReplyStreamHandler {
+    pub fn new(sender: Option<Sender<ReplyStreamEvent>>, ctrlc: Arc<AtomicBool>) -> Self {
         Self {
-            output: String::new(),
             sender,
+            ctrlc,
+            buffer: String::new(),
         }
     }
 
     pub fn text(&mut self, text: &str) {
         match self.sender.as_ref() {
             Some(tx) => {
-                let _ = tx.send(RenderStreamEvent::Text(text.to_string()));
+                let _ = tx.send(ReplyStreamEvent::Text(text.to_string()));
             }
             None => {
                 dump(text, 0);
             }
         }
-        self.output.push_str(text);
+        self.buffer.push_str(text);
     }
 
     pub fn done(&mut self) {
         match self.sender.as_ref() {
             Some(tx) => {
-                let _ = tx.send(RenderStreamEvent::Done);
+                let _ = tx.send(ReplyStreamEvent::Done);
             }
             None => {
                 dump("", 2);
             }
         }
     }
+
+    pub fn get_buffer(&self) -> &str {
+        &self.buffer
+    }
+
+    pub fn get_ctrlc(&self) -> Arc<AtomicBool> {
+        self.ctrlc.clone()
+    }
 }
 
-pub enum RenderStreamEvent {
+pub enum ReplyStreamEvent {
     Text(String),
     Done,
 }
diff --git a/src/repl/mod.rs b/src/repl/mod.rs
index 6aeb233..30f09fe 100644
--- a/src/repl/mod.rs
+++ b/src/repl/mod.rs
@@ -1,6 +1,8 @@
 mod handler;
 mod init;
 
+use crate::client::ChatGptClient;
+use crate::config::SharedConfig;
 use crate::term;
 use crate::utils::{copy, dump};
 
@@ -32,7 +34,8 @@ pub struct Repl {
 }
 
 impl Repl {
-    pub fn run(&mut self, handler: ReplCmdHandler) -> Result<()> {
+    pub fn run(&mut self, client: ChatGptClient, config: SharedConfig) -> Result<()> {
+        let handler = ReplCmdHandler::init(client, config)?;
         dump(
             format!("Welcome to aichat {}", env!("CARGO_PKG_VERSION")),
             1,