feat: support textwrap (#171)

* feat: support textwrap

* improve cmd_render_stream

* done cmd_render_stream

* add `-w` alias to `--wrap`

* done repl_render_stream

* add `config.wrap_code`

* remove cached config.wrap_width

* fix unxpected duplicate lines on kitty

* refactor markdown render

* improve render

* fix test

9 files changed, 387 insertions, 174 deletions

diff --git a/src/cli.rs b/src/cli.rs
index 183a965..91c48e8 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -7,7 +7,7 @@ pub struct Cli {
     /// List all models
     #[clap(long)]
     pub list_models: bool,
-    /// Choose a llm model
+    /// Choose a LLM model
     #[clap(short, long)]
     pub model: Option<String>,
     /// List all roles
@@ -25,6 +25,9 @@ pub struct Cli {
     /// Initiate or continue a session
     #[clap(short = 's', long)]
     pub session: Option<Option<String>>,
+    /// Specify the text-wrapping mode (no*, auto, <max-width>)
+    #[clap(short = 'w', long)]
+    pub wrap: Option<String>,
     /// Print information
     #[clap(long)]
     pub info: bool,
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 922fb0a..42521cc 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -9,6 +9,7 @@ use self::session::{Session, TEMP_SESSION_NAME};
 use crate::client::openai::{OpenAIClient, OpenAIConfig};
 use crate::client::{all_clients, create_client_config, list_models, ClientConfig, ModelInfo};
 use crate::config::message::num_tokens_from_messages;
+use crate::render::RenderOptions;
 use crate::utils::{get_env_name, now};
 
 use anyhow::{anyhow, bail, Context, Result};
@@ -56,6 +57,10 @@ pub struct Config {
     pub dry_run: bool,
     /// If set true, use light theme
     pub light_theme: bool,
+    /// Specify the text-wrapping mode (no*, auto, <max-width>)
+    pub wrap: Option<String>,
+    /// Whethter wrap code block
+    pub wrap_code: bool,
     /// Automatically copy the last output to the clipboard
     pub auto_copy: bool,
     /// REPL keybindings, possible values: emacs (default), vi
@@ -86,6 +91,8 @@ impl Default for Config {
             highlight: true,
             dry_run: false,
             light_theme: false,
+            wrap: None,
+            wrap_code: false,
             auto_copy: false,
             keybindings: Default::default(),
             clients: vec![ClientConfig::OpenAI(OpenAIConfig::default())],
@@ -125,6 +132,9 @@ impl Config {
         if let Some(name) = config.model.clone() {
             config.set_model(&name)?;
         }
+        if let Some(wrap) = config.wrap.clone() {
+            config.set_wrap(&wrap)?;
+        }
 
         config.merge_env_vars();
         config.load_roles()?;
@@ -163,6 +173,10 @@ impl Config {
     pub fn save_message(&mut self, input: &str, output: &str) -> Result<()> {
         self.last_message = Some((input.to_string(), output.to_string()));
 
+        if self.dry_run {
+            return Ok(());
+        }
+
         if let Some(session) = self.session.as_mut() {
             session.add_message(input, output)?;
             return Ok(());
@@ -295,6 +309,20 @@ impl Config {
         Ok(messages)
     }
 
+    pub fn set_wrap(&mut self, value: &str) -> Result<()> {
+        if value == "no" {
+            self.wrap = None;
+        } else if value == "auto" {
+            self.wrap = Some(value.into());
+        } else {
+            value
+                .parse::<u16>()
+                .map_err(|_| anyhow!("Invalid wrap value"))?;
+            self.wrap = Some(value.into())
+        }
+        Ok(())
+    }
+
     pub fn set_model(&mut self, value: &str) -> Result<()> {
         let models = list_models(self);
         let mut model_info = None;
@@ -333,6 +361,10 @@ impl Config {
         let temperature = self
             .temperature
             .map_or_else(|| String::from("-"), |v| v.to_string());
+        let wrap = self
+            .wrap
+            .clone()
+            .map_or_else(|| String::from("no"), |v| v.to_string());
         let items = vec![
             ("config_file", path_info(&Self::config_file()?)),
             ("roles_file", path_info(&Self::roles_file()?)),
@@ -343,6 +375,8 @@ impl Config {
             ("save", self.save.to_string()),
             ("highlight", self.highlight.to_string()),
             ("light_theme", self.light_theme.to_string()),
+            ("wrap", wrap),
+            ("wrap_code", self.wrap_code.to_string()),
             ("dry_run", self.dry_run.to_string()),
             ("keybindings", self.keybindings.stringify().into()),
         ];
@@ -500,8 +534,13 @@ impl Config {
         }
     }
 
-    pub const fn get_render_options(&self) -> (bool, bool) {
-        (self.highlight, self.light_theme)
+    pub fn get_render_options(&self) -> RenderOptions {
+        RenderOptions::new(
+            self.highlight,
+            self.light_theme,
+            self.wrap.clone(),
+            self.wrap_code,
+        )
     }
 
     pub fn maybe_print_send_tokens(&self, input: &str) {
diff --git a/src/main.rs b/src/main.rs
index d2660a8..f551e9c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -47,6 +47,9 @@ fn main() -> Result<()> {
         println!("{sessions}");
         exit(0);
     }
+    if let Some(wrap) = &cli.wrap {
+        config.write().set_wrap(wrap)?;
+    }
     if cli.light_theme {
         config.write().light_theme = true;
     }
@@ -115,14 +118,10 @@ fn start_directive(
     }
     config.read().maybe_print_send_tokens(input);
     let output = if no_stream {
-        let (highlight, light_theme) = config.read().get_render_options();
+        let render_options = config.read().get_render_options();
         let output = client.send_message(input)?;
-        if highlight {
-            let mut markdown_render = MarkdownRender::new(light_theme);
-            println!("{}", markdown_render.render_block(&output).trim());
-        } else {
-            println!("{}", output.trim());
-        }
+        let mut markdown_render = MarkdownRender::init(render_options)?;
+        println!("{}", markdown_render.render(&output).trim());
         output
     } else {
         let wg = WaitGroup::new();
diff --git a/src/render/cmd.rs b/src/render/cmd.rs
index 527c194..fb45b0a 100644
--- a/src/render/cmd.rs
+++ b/src/render/cmd.rs
@@ -5,15 +5,16 @@ use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
 
 use anyhow::Result;
 use crossbeam::channel::Receiver;
+use textwrap::core::display_width;
 
 #[allow(clippy::unnecessary_wraps, clippy::module_name_repetitions)]
 pub fn cmd_render_stream(
     rx: &Receiver<ReplyStreamEvent>,
-    light_theme: bool,
+    render: &mut MarkdownRender,
     abort: &SharedAbortSignal,
 ) -> Result<()> {
     let mut buffer = String::new();
-    let mut markdown_render = MarkdownRender::new(light_theme);
+    let mut col = 0;
     loop {
         if abort.aborted() {
             return Ok(());
@@ -23,28 +24,39 @@ pub fn cmd_render_stream(
                 ReplyStreamEvent::Text(text) => {
                     if text.contains('\n') {
                         let text = format!("{buffer}{text}");
-                        let mut lines: Vec<&str> = text.split('\n').collect();
-                        buffer = lines.pop().unwrap_or_default().to_string();
-                        let output = lines.join("\n");
-                        print_now!("{}\n", markdown_render.render_block(&output));
+                        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..]);
+                        col = 0;
                     } else {
                         buffer = format!("{buffer}{text}");
-                        if !(markdown_render.is_code_block()
-                            || buffer.len() < 60
+                        if !(render.is_code()
+                            || buffer.len() < 40
                             || buffer.starts_with('#')
                             || buffer.starts_with('>')
                             || buffer.starts_with('|'))
                         {
-                            if let Some((output, remain)) = split_line(&buffer) {
-                                print_now!("{}", markdown_render.render_line(&output));
+                            if let Some((head, remain)) = split_line_sematic(&buffer) {
                                 buffer = remain;
+                                let input = format!("{}{head}", spaces(col));
+                                let output = render.render(&input);
+                                let output = &output[col..];
+                                let (_, tail) = split_line_tail(output);
+                                if output.contains('\n') {
+                                    col = display_width(tail);
+                                } else {
+                                    col += display_width(output);
+                                }
+                                print_now!("{}", output);
                             }
                         }
                     }
                 }
                 ReplyStreamEvent::Done => {
-                    let output = markdown_render.render_block(&buffer);
-                    print_now!("{}\n", output.trim_end());
+                    let input = format!("{}{buffer}", spaces(col));
+                    print_now!("{}\n", render.render(&input));
                     break;
                 }
             }
@@ -53,9 +65,9 @@ pub fn cmd_render_stream(
     Ok(())
 }
 
-fn split_line(line: &str) -> Option<(String, String)> {
+fn split_line_sematic(text: &str) -> Option<(String, String)> {
     let mut balance: Vec<Kind> = Vec::new();
-    let chars: Vec<char> = line.chars().collect();
+    let chars: Vec<char> = text.chars().collect();
     let mut index = 0;
     let len = chars.len();
     while index < len - 1 {
@@ -82,6 +94,18 @@ fn split_line(line: &str) -> Option<(String, String)> {
     None
 }
 
+pub(crate) fn split_line_tail(text: &str) -> (&str, &str) {
+    if let Some((head, tail)) = text.rsplit_once('\n') {
+        (head, tail)
+    } else {
+        ("", text)
+    }
+}
+
+fn spaces(n: usize) -> String {
+    " ".repeat(n)
+}
+
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
 enum Kind {
     ParentheseStart,
@@ -176,12 +200,12 @@ mod tests {
     macro_rules! assert_split_line {
         ($a:literal, $b:literal, true) => {
             assert_eq!(
-                split_line(&format!("{}{}", $a, $b)),
+                split_line_sematic(&format!("{}{}", $a, $b)),
                 Some(($a.into(), $b.into()))
             );
         };
         ($a:literal, $b:literal, false) => {
-            assert_eq!(split_line(&format!("{}{}", $a, $b)), None);
+            assert_eq!(split_line_sematic(&format!("{}{}", $a, $b)), None);
         };
     }
 
diff --git a/src/render/markdown.rs b/src/render/markdown.rs
index 66bb601..4866f3d 100644
--- a/src/render/markdown.rs
+++ b/src/render/markdown.rs
@@ -1,4 +1,6 @@
+use anyhow::{anyhow, Context, Result};
 use crossterm::style::{Color, Stylize};
+use crossterm::terminal;
 use lazy_static::lazy_static;
 use std::collections::HashMap;
 use syntect::highlighting::{Color as SyntectColor, FontStyle, Style, Theme};
@@ -23,115 +25,175 @@ lazy_static! {
 
 #[allow(clippy::module_name_repetitions)]
 pub struct MarkdownRender {
+    options: RenderOptions,
     syntax_set: SyntaxSet,
-    md_theme: Theme,
-    code_color: Color,
+    md_theme: Option<Theme>,
+    code_color: Option<Color>,
     md_syntax: SyntaxReference,
     code_syntax: Option<SyntaxReference>,
     prev_line_type: LineType,
+    wrap_width: Option<u16>,
 }
 
 impl MarkdownRender {
-    pub fn new(light_theme: bool) -> Self {
-        let syntax_set: SyntaxSet =
-            bincode::deserialize_from(SYNTAXES).expect("invalid syntaxes binary");
-        let md_theme: Theme = if light_theme {
-            bincode::deserialize_from(MD_THEME_LIGHT).expect("invalid theme binary")
-        } else {
-            bincode::deserialize_from(MD_THEME).expect("invalid theme binary")
+    pub fn init(options: RenderOptions) -> Result<Self> {
+        let syntax_set: SyntaxSet = bincode::deserialize_from(SYNTAXES)
+            .with_context(|| "MarkdownRender: invalid syntaxes binary")?;
+
+        let md_theme: Option<Theme> = match (options.highlight, options.light_theme) {
+            (false, _) => None,
+            (true, false) => Some(
+                bincode::deserialize_from(MD_THEME)
+                    .with_context(|| "MarkdownRender: invalid theme binary")?,
+            ),
+            (true, true) => Some(
+                bincode::deserialize_from(MD_THEME_LIGHT)
+                    .expect("MarkdownRender: invalid theme binary"),
+            ),
         };
-        let code_color = get_code_color(&md_theme);
+        let code_color = md_theme.as_ref().map(get_code_color);
         let md_syntax = syntax_set.find_syntax_by_extension("md").unwrap().clone();
         let line_type = LineType::Normal;
-        Self {
+        let wrap_width = match options.wrap.as_deref() {
+            None => None,
+            Some(value) => match terminal::size() {
+                Ok((columns, _)) => {
+                    if value == "auto" {
+                        Some(columns)
+                    } else {
+                        let value = value
+                            .parse::<u16>()
+                            .map_err(|_| anyhow!("Invalid wrap value"))?;
+                        Some(columns.min(value))
+                    }
+                }
+                Err(_) => None,
+            },
+        };
+        Ok(Self {
             syntax_set,
             md_theme,
             code_color,
             md_syntax,
             code_syntax: None,
             prev_line_type: line_type,
-        }
+            wrap_width,
+            options,
+        })
+    }
+
+    pub(crate) const fn is_code(&self) -> bool {
+        matches!(
+            self.prev_line_type,
+            LineType::CodeBegin | LineType::CodeInner
+        )
     }
 
-    pub fn render_block(&mut self, src: &str) -> String {
-        src.split('\n')
-            .map(|line| {
-                self.render_line_impl(line)
-                    .unwrap_or_else(|| line.to_string())
-            })
+    pub fn render(&mut self, text: &str) -> String {
+        text.split('\n')
+            .map(|line| self.render_line_mut(line))
             .collect::<Vec<String>>()
             .join("\n")
     }
 
     pub fn render_line(&self, line: &str) -> String {
-        let output = if self.is_code_block() && detect_code_block(line).is_none() {
-            self.render_code_line(line)
+        let (_, code_syntax, is_code) = self.check_line(line);
+        if is_code {
+            self.highlint_code_line(line, &code_syntax)
         } else {
-            self.render_line_inner(line, &self.md_syntax)
-        };
-        output.unwrap_or_else(|| line.to_string())
+            self.highligh_line(line, &self.md_syntax, false)
+        }
     }
 
-    pub const fn is_code_block(&self) -> bool {
-        matches!(
-            self.prev_line_type,
-            LineType::CodeBegin | LineType::CodeInner
-        )
+    fn render_line_mut(&mut self, line: &str) -> String {
+        let (line_type, code_syntax, is_code) = self.check_line(line);
+        let output = if is_code {
+            self.highlint_code_line(line, &code_syntax)
+        } else {
+            self.highligh_line(line, &self.md_syntax, false)
+        };
+        self.prev_line_type = line_type;
+        self.code_syntax = code_syntax;
+        output
     }
 
-    fn render_line_impl(&mut self, line: &str) -> Option<String> {
+    fn check_line(&self, line: &str) -> (LineType, Option<SyntaxReference>, bool) {
+        let mut line_type = self.prev_line_type;
+        let mut code_syntax = self.code_syntax.clone();
+        let mut is_code = false;
         if let Some(lang) = detect_code_block(line) {
-            match self.prev_line_type {
+            match line_type {
                 LineType::Normal | LineType::CodeEnd => {
-                    self.prev_line_type = LineType::CodeBegin;
-                    self.code_syntax = if lang.is_empty() {
+                    line_type = LineType::CodeBegin;
+                    code_syntax = if lang.is_empty() {
                         None
                     } else {
                         self.find_syntax(&lang).cloned()
                     };
                 }
                 LineType::CodeBegin | LineType::CodeInner => {
-                    self.prev_line_type = LineType::CodeEnd;
-                    self.code_syntax = None;
+                    line_type = LineType::CodeEnd;
+                    code_syntax = None;
                 }
             }
-            self.render_line_inner(line, &self.md_syntax)
         } else {
-            match self.prev_line_type {
-                LineType::Normal => self.render_line_inner(line, &self.md_syntax),
+            match line_type {
+                LineType::Normal => {}
                 LineType::CodeEnd => {
-                    self.prev_line_type = LineType::Normal;
-                    self.render_line_inner(line, &self.md_syntax)
+                    line_type = LineType::Normal;
                 }
                 LineType::CodeBegin => {
-                    if self.code_syntax.is_none() {
+                    if code_syntax.is_none() {
                         if let Some(syntax) = self.syntax_set.find_syntax_by_first_line(line) {
-                            self.code_syntax = Some(syntax.clone());
+                            code_syntax = Some(syntax.clone());
                         }
                     }
-                    self.prev_line_type = LineType::CodeInner;
-                    self.render_code_line(line)
+                    line_type = LineType::CodeInner;
+                    is_code = true;
+                }
+                LineType::CodeInner => {
+                    is_code = true;
                 }
-                LineType::CodeInner => self.render_code_line(line),
             }
         }
+        (line_type, code_syntax, is_code)
     }
 
-    fn render_line_inner(&self, line: &str, syntax: &SyntaxReference) -> Option<String> {
+    fn highligh_line(&self, line: &str, syntax: &SyntaxReference, is_code: bool) -> String {
         let ws: String = line.chars().take_while(|c| c.is_whitespace()).collect();
-        let trimed_line = &line[ws.len()..];
-        let mut highlighter = HighlightLines::new(syntax, &self.md_theme);
-        let ranges = highlighter
-            .highlight_line(trimed_line, &self.syntax_set)
-            .ok()?;
-        Some(format!("{ws}{}", as_terminal_escaped(&ranges)))
-    }
-
-    fn render_code_line(&self, line: &str) -> Option<String> {
-        self.code_syntax.as_ref().map_or_else(
-            || Some(format!("{}", line.with(self.code_color))),
-            |syntax| self.render_line_inner(line, syntax),
-        )
+        let trimed_line: &str = &line[ws.len()..];
+        let mut line_highlighted = None;
+        if let Some(theme) = &self.md_theme {
+            let mut highlighter = HighlightLines::new(syntax, theme);
+            if let Ok(ranges) = highlighter.highlight_line(trimed_line, &self.syntax_set) {
+                line_highlighted = Some(format!("{ws}{}", as_terminal_escaped(&ranges)))
+            }
+        }
+        let line = line_highlighted.unwrap_or_else(|| line.into());
+        self.wrap_line(line, is_code)
+    }
+
+    fn highlint_code_line(&self, line: &str, code_syntax: &Option<SyntaxReference>) -> String {
+        if let Some(syntax) = code_syntax {
+            self.highligh_line(line, syntax, true)
+        } else {
+            let line = match self.code_color {
+                Some(color) => line.with(color).to_string(),
+                None => line.to_string(),
+            };
+            self.wrap_line(line, true)
+        }
+    }
+
+    fn wrap_line(&self, line: String, is_code: bool) -> String {
+        if let Some(width) = self.wrap_width {
+            if is_code && !self.options.wrap_code {
+                return line;
+            }
+            textwrap::wrap(&line, width as usize).join("\n")
+        } else {
+            line
+        }
     }
 
     fn find_syntax(&self, lang: &str) -> Option<&SyntaxReference> {
@@ -146,6 +208,30 @@ impl MarkdownRender {
     }
 }
 
+#[derive(Debug, Clone, Default)]
+pub struct RenderOptions {
+    pub highlight: bool,
+    pub light_theme: bool,
+    pub wrap: Option<String>,
+    pub wrap_code: bool,
+}
+
+impl RenderOptions {
+    pub(crate) fn new(
+        highlight: bool,
+        light_theme: bool,
+        wrap: Option<String>,
+        wrap_code: bool,
+    ) -> Self {
+        Self {
+            highlight,
+            light_theme,
+            wrap,
+            wrap_code,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum LineType {
     Normal,
@@ -222,6 +308,44 @@ fn get_code_color(theme: &Theme) -> Color {
 mod tests {
     use super::*;
 
+    const TEXT: &str = r#"
+To unzip a file in Rust, you can use the `zip` crate. Here's an example code that shows how to unzip a file:
+
+```rust
+use std::fs::File;
+
+fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box<dyn std::error::Error>> {
+    todo!()
+}
+```
+"#;
+    const TEXT_NO_WRAP_CODE: &str = r#"
+To unzip a file in Rust, you can use the `zip` crate. Here's an example code
+that shows how to unzip a file:
+
+```rust
+use std::fs::File;
+
+fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box<dyn std::error::Error>> {
+    todo!()
+}
+```
+"#;
+
+    const TEXT_WRAP_ALL: &str = r#"
+To unzip a file in Rust, you can use the `zip` crate. Here's an example code
+that shows how to unzip a file:
+
+```rust
+use std::fs::File;
+
+fn unzip_file(path: &str, output_dir: &str) -> Result<(), Box<dyn
+std::error::Error>> {
+    todo!()
+}
+```
+"#;
+
     #[test]
     fn test_assets() {
         let syntax_set: SyntaxSet =
@@ -233,7 +357,37 @@ mod tests {
 
     #[test]
     fn test_render() {
-        let render = MarkdownRender::new(true);
+        let options = RenderOptions::default();
+        let render = MarkdownRender::init(options).unwrap();
         assert!(render.find_syntax("csharp").is_some());
     }
+
+    #[test]
+    fn no_theme() {
+        let options = RenderOptions::default();
+        let mut render = MarkdownRender::init(options).unwrap();
+        let output = render.render(TEXT);
+        assert_eq!(TEXT, output);
+    }
+
+    #[test]
+    fn no_wrap_code() {
+        let options = RenderOptions::default();
+        let mut render = MarkdownRender::init(options).unwrap();
+        render.wrap_width = Some(80);
+        let output = render.render(TEXT);
+        assert_eq!(TEXT_NO_WRAP_CODE, output);
+    }
+
+    #[test]
+    fn wrap_all() {
+        let options = RenderOptions {
+            wrap_code: true,
+            ..Default::default()
+        };
+        let mut render = MarkdownRender::init(options).unwrap();
+        render.wrap_width = Some(80);
+        let output = render.render(TEXT);
+        assert_eq!(TEXT_WRAP_ALL, output);
+    }
 }
diff --git a/src/render/mod.rs b/src/render/mod.rs
index e7c6228..9d7b06a 100644
--- a/src/render/mod.rs
+++ b/src/render/mod.rs
@@ -4,7 +4,7 @@ mod repl;
 
 use self::cmd::cmd_render_stream;
 #[allow(clippy::module_name_repetitions)]
-pub use self::markdown::MarkdownRender;
+pub use self::markdown::{MarkdownRender, RenderOptions};
 use self::repl::repl_render_stream;
 
 use crate::client::Client;
@@ -26,26 +26,27 @@ pub fn render_stream(
     abort: SharedAbortSignal,
     wg: WaitGroup,
 ) -> Result<String> {
-    let (highlight, light_theme) = config.read().get_render_options();
-    let mut stream_handler = if highlight {
+    let render_options = config.read().get_render_options();
+    let mut stream_handler = {
         let (tx, rx) = unbounded();
         let abort_clone = abort.clone();
         spawn(move || {
-            let err = if repl {
-                repl_render_stream(&rx, light_theme, &abort)
-            } else {
-                cmd_render_stream(&rx, light_theme, &abort)
+            let run = move || {
+                if repl {
+                    let mut render = MarkdownRender::init(render_options)?;
+                    repl_render_stream(&rx, &mut render, &abort)
+                } else {
+                    let mut render = MarkdownRender::init(render_options)?;
+                    cmd_render_stream(&rx, &mut render, &abort)
+                }
             };
-            if let Err(err) = err {
+            if let Err(err) = run() {
                 let err = format!("{err:?}");
                 print_now!("{}\n\n", err.trim());
             }
             drop(wg);
         });
-        ReplyStreamHandler::new(Some(tx), repl, abort_clone)
-    } else {
-        drop(wg);
-        ReplyStreamHandler::new(None, repl, abort)
+        ReplyStreamHandler::new(tx, abort_clone)
     };
     client.send_message_streaming(input, &mut stream_handler)?;
     let buffer = stream_handler.get_buffer();
diff --git a/src/render/repl.rs b/src/render/repl.rs
index eb50750..c411670 100644
--- a/src/render/repl.rs
+++ b/src/render/repl.rs
@@ -1,4 +1,4 @@
-use super::MarkdownRender;
+use super::{cmd::split_line_tail, MarkdownRender};
 
 use crate::repl::{ReplyStreamEvent, SharedAbortSignal};
 
@@ -14,18 +14,18 @@ use std::{
     io::{self, Stdout, Write},
     time::{Duration, Instant},
 };
-use unicode_width::UnicodeWidthStr;
+use textwrap::core::display_width;
 
 #[allow(clippy::module_name_repetitions)]
 pub fn repl_render_stream(
     rx: &Receiver<ReplyStreamEvent>,
-    light_theme: bool,
+    render: &mut MarkdownRender,
     abort: &SharedAbortSignal,
 ) -> Result<()> {
     enable_raw_mode()?;
     let mut stdout = io::stdout();
 
-    let ret = repl_render_stream_inner(rx, light_theme, abort, &mut stdout);
+    let ret = repl_render_stream_inner(rx, render, abort, &mut stdout);
 
     disable_raw_mode()?;
 
@@ -34,15 +34,16 @@ pub fn repl_render_stream(
 
 fn repl_render_stream_inner(
     rx: &Receiver<ReplyStreamEvent>,
-    light_theme: bool,
+    render: &mut MarkdownRender,
     abort: &SharedAbortSignal,
     writer: &mut Stdout,
 ) -> Result<()> {
     let mut last_tick = Instant::now();
     let tick_rate = Duration::from_millis(50);
     let mut buffer = String::new();
-    let mut markdown_render = MarkdownRender::new(light_theme);
     let columns = terminal::size()?.0;
+
+    let mut clear_rows = 0;
     loop {
         if abort.aborted() {
             return Ok(());
@@ -51,44 +52,45 @@ fn repl_render_stream_inner(
         if let Ok(evt) = rx.try_recv() {
             match evt {
                 ReplyStreamEvent::Text(text) => {
-                    if !buffer.is_empty() {
-                        let buffer_width = buffer.width() as u16;
-                        let need_rows = (buffer_width + columns - 1) / columns;
-                        let (col, row) = cursor::position()?;
-
-                        if row + 1 >= need_rows {
-                            if col == 0 {
-                                queue!(writer, cursor::MoveTo(0, row - need_rows))?;
-                            } else {
-                                queue!(writer, cursor::MoveTo(0, row + 1 - need_rows))?;
-                            }
-                        } else {
-                            queue!(
-                                writer,
-                                terminal::ScrollUp(need_rows - 1 - row),
-                                cursor::MoveTo(0, 0)
-                            )?;
-                        }
+                    let (col, mut row) = cursor::position()?;
+
+                    // fix unxpected duplicate lines on kitty, see https://github.com/sigoden/aichat/issues/105
+                    if col == 0 && row > 0 && display_width(&buffer) == columns as usize {
+                        row -= 1;
+                    }
+
+                    if row + 1 >= clear_rows {
+                        queue!(writer, cursor::MoveTo(0, row - clear_rows))?;
+                    } else {
+                        let scroll_rows = clear_rows - row - 1;
+                        queue!(
+                            writer,
+                            terminal::ScrollUp(scroll_rows),
+                            cursor::MoveTo(0, 0),
+                        )?;
                     }
 
                     if text.contains('\n') {
                         let text = format!("{buffer}{text}");
-                        let mut lines: Vec<&str> = text.split('\n').collect();
-                        buffer = lines.pop().unwrap_or_default().to_string();
-                        let output = markdown_render.render_block(&lines.join("\n"));
-                        for line in output.split('\n') {
-                            queue!(
-                                writer,
-                                style::Print(line),
-                                style::Print("\n"),
-                                cursor::MoveLeft(columns),
-                            )?;
-                        }
+                        let (head, tail) = split_line_tail(&text);
+                        buffer = tail.to_string();
+                        let output = render.render(head);
+                        print_block(writer, &output, columns)?;
                         queue!(writer, style::Print(&buffer),)?;
+                        clear_rows = 0;
                     } else {
                         buffer = format!("{buffer}{text}");
-                        let output = markdown_render.render_line(&buffer);
-                        queue!(writer, style::Print(&output))?;
+                        let output = render.render_line(&buffer);
+                        if output.contains('\n') {
+                            let (head, tail) = split_line_tail(&output);
+                            clear_rows = print_block(writer, head, columns)?;
+                            queue!(writer, style::Print(&tail),)?;
+                        } else {
+                            queue!(writer, style::Print(&output))?;
+                            let buffer_width = display_width(&output) as u16;
+                            let need_rows = (buffer_width + columns - 1) / columns;
+                            clear_rows = need_rows.saturating_sub(1);
+                        }
                     }
 
                     writer.flush()?;
@@ -128,3 +130,17 @@ fn repl_render_stream_inner(
     }
     Ok(())
 }
+
+fn print_block(writer: &mut Stdout, text: &str, columns: u16) -> Result<u16> {
+    let mut num = 0;
+    for line in text.split('\n') {
+        queue!(
+            writer,
+            style::Print(line),
+            style::Print("\n"),
+            cursor::MoveLeft(columns),
+        )?;
+        num += 1;
+    }
+    Ok(num)
+}
diff --git a/src/repl/handler.rs b/src/repl/handler.rs
index fd157dc..35a57da 100644
--- a/src/repl/handler.rs
+++ b/src/repl/handler.rs
@@ -115,9 +115,9 @@ impl ReplCmdHandler {
             }
             ReplCmd::ReadFile(file) => {
                 let mut contents = String::new();
-                let mut file = fs::File::open(file).expect("Unable to open file");
+                let mut file = fs::File::open(file).with_context(|| "Unable to open file")?;
                 file.read_to_string(&mut contents)
-                    .expect("Unable to read file");
+                    .with_context(|| "Unable to read file")?;
                 self.handle(ReplCmd::Submit(contents))?;
             }
         }
@@ -137,23 +137,17 @@ impl ReplCmdHandler {
 
 #[allow(clippy::module_name_repetitions)]
 pub struct ReplyStreamHandler {
-    sender: Option<Sender<ReplyStreamEvent>>,
+    sender: Sender<ReplyStreamEvent>,
     buffer: String,
     abort: SharedAbortSignal,
-    repl: bool,
 }
 
 impl ReplyStreamHandler {
-    pub fn new(
-        sender: Option<Sender<ReplyStreamEvent>>,
-        repl: bool,
-        abort: SharedAbortSignal,
-    ) -> Self {
+    pub fn new(sender: Sender<ReplyStreamEvent>, abort: SharedAbortSignal) -> Self {
         Self {
             sender,
             abort,
             buffer: String::new(),
-            repl,
         }
     }
 
@@ -162,37 +156,20 @@ impl ReplyStreamHandler {
             return Ok(());
         }
         self.buffer.push_str(text);
-        match self.sender.as_ref() {
-            Some(tx) => {
-                let ret = tx
-                    .send(ReplyStreamEvent::Text(text.to_string()))
-                    .with_context(|| "Failed to send StreamEvent:Text");
-                self.safe_ret(ret)?;
-            }
-            None => {
-                print_now!("{}", 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<()> {
-        if let Some(tx) = self.sender.as_ref() {
-            let ret = tx
-                .send(ReplyStreamEvent::Done)
-                .with_context(|| "Failed to send StreamEvent:Done");
-            self.safe_ret(ret)?;
-        } else {
-            if !self.buffer.ends_with('\n') {
-                print_now!("\n");
-            }
-            if self.repl {
-                print_now!("\n");
-                if cfg!(macos) {
-                    print_now!("\n");
-                }
-            }
-        }
+        let ret = self
+            .sender
+            .send(ReplyStreamEvent::Done)
+            .with_context(|| "Failed to send StreamEvent:Done");
+        self.safe_ret(ret)?;
         Ok(())
     }
 
diff --git a/src/repl/prompt.rs b/src/repl/prompt.rs
index be3763b..f48ce64 100644
--- a/src/repl/prompt.rs
+++ b/src/repl/prompt.rs
@@ -41,15 +41,15 @@ impl ReplPrompt {
     }
 
     pub fn get_colors(config: &SharedConfig) -> (Color, nu_ansi_term::Color, Color, Color) {
-        let (highlight, light_theme) = config.read().get_render_options();
-        if highlight {
+        let render_options = config.read().get_render_options();
+        if render_options.highlight {
             (
                 PROMPT_COLOR,
                 PROMPT_MULTILINE_COLOR,
                 INDICATOR_COLOR,
                 PROMPT_RIGHT_COLOR,
             )
-        } else if light_theme {
+        } else if render_options.light_theme {
             (
                 Color::Black,
                 nu_ansi_term::Color::Black,