From da5205ad39b4276c119b81aa4930b8d8ec854889 Mon Sep 17 00:00:00 2001
From: Eric Fu <fuyufjh@gmail.com>
Date: Sun, 28 May 2023 16:23:39 +0800
Subject: [PATCH] fix prompt display under Windows

---
 src/main.rs        |  9 ++++++---
 src/repl_helper.rs | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 3 deletions(-)
 create mode 100644 src/repl_helper.rs

diff --git a/src/main.rs b/src/main.rs
index 486d5ef..3498947 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,14 +3,16 @@ use clap::Parser;
 use console::style;
 use futures::stream::StreamExt;
 use log::{debug, trace};
+use repl_helper::ReplHelper;
 use reqwest::header::{HeaderMap, AUTHORIZATION};
 use reqwest::{Client, RequestBuilder, StatusCode};
 use reqwest_eventsource::{Event, EventSource};
 use rustyline::error::ReadlineError;
-use rustyline::{DefaultEditor, Editor};
+use rustyline::Editor;
 use std::io::Write;
 
 mod model;
+mod repl_helper;
 mod spinner;
 
 use model::*;
@@ -146,7 +148,8 @@ impl Session {
     }
 
     pub async fn run_interactive(&mut self) -> Result<()> {
-        let mut rl = DefaultEditor::new()?;
+        let mut rl = Editor::<repl_helper::ReplHelper, _>::new()?;
+        rl.set_helper(Some(ReplHelper::default()));
 
         // Persist input history in `$HOME/.heygpt_history`
         let history_file = {
@@ -211,7 +214,7 @@ impl Session {
         I: rustyline::history::History,
     {
         loop {
-            let readline = rl.readline(&format!("{} => ", style(role).bold().cyan()));
+            let readline = rl.readline(&format!("{} => ", role));
             match readline {
                 Ok(line) => {
                     if line.is_empty() {
diff --git a/src/repl_helper.rs b/src/repl_helper.rs
new file mode 100644
index 0000000..0412f56
--- /dev/null
+++ b/src/repl_helper.rs
@@ -0,0 +1,44 @@
+use std::borrow::Cow;
+
+use console::style;
+use rustyline::completion::Completer;
+use rustyline::highlight::Highlighter;
+use rustyline::hint::Hinter;
+use rustyline::validate::{ValidationContext, ValidationResult, Validator};
+use rustyline::Helper;
+
+/// The rustyline helper for interactive mode.
+/// Currently it's mostly for highlighting the prompt correctly.
+#[derive(Clone, Debug, Default)]
+pub struct ReplHelper;
+
+impl Helper for ReplHelper {}
+
+impl Validator for ReplHelper {
+    fn validate(&self, _ctx: &mut ValidationContext) -> rustyline::Result<ValidationResult> {
+        Ok(ValidationResult::Valid(None))
+    }
+}
+
+impl Completer for ReplHelper {
+    type Candidate = String;
+}
+
+impl Hinter for ReplHelper {
+    type Hint = String;
+}
+
+impl Highlighter for ReplHelper {
+    fn highlight_prompt<'b, 's: 'b, 'p: 'b>(
+        &'s self,
+        prompt: &'p str,
+        _default: bool,
+    ) -> Cow<'b, str> {
+        if let Some(idx) = prompt.find(" => ") {
+            let (role, rest) = prompt.split_at(idx);
+            Cow::Owned(format!("{}{}", style(role).bold().cyan(), rest))
+        } else {
+            Cow::Borrowed(prompt)
+        }
+    }
+}