From f4160ff85b61e41a599e85f82d266878d35de47d Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Tue, 24 Oct 2023 19:28:16 +0800
Subject: [PATCH] feat: suppport vi keybindings (#148)

---
 src/config/mod.rs |  4 ++++
 src/repl/init.rs  | 28 +++++++++++++++++++++-------
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/src/config/mod.rs b/src/config/mod.rs
index f29ece7..3fc3f5b 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -74,6 +74,8 @@ pub struct Config {
     pub connect_timeout: usize,
     /// Automatically copy the last output to the clipboard
     pub auto_copy: bool,
+    /// Use vi keybindings, overriding the default Emacs keybindings
+    pub vi_keybindings: bool,
     /// Predefined roles
     #[serde(skip)]
     pub roles: Vec<Role>,
@@ -102,6 +104,7 @@ impl Default for Config {
             light_theme: false,
             connect_timeout: 10,
             auto_copy: false,
+            vi_keybindings: false,
             roles: vec![],
             role: None,
             conversation: None,
@@ -353,6 +356,7 @@ impl Config {
             ("light_theme", self.light_theme.to_string()),
             ("connect_timeout", self.connect_timeout.to_string()),
             ("dry_run", self.dry_run.to_string()),
+            ("vi_keybindings", self.vi_keybindings.to_string()),
         ];
         let mut output = String::new();
         for (name, value) in items {
diff --git a/src/repl/init.rs b/src/repl/init.rs
index 089e855..1204fa7 100644
--- a/src/repl/init.rs
+++ b/src/repl/init.rs
@@ -6,8 +6,9 @@ use crate::config::{Config, SharedConfig};
 
 use anyhow::{Context, Result};
 use reedline::{
-    default_emacs_keybindings, ColumnarMenu, DefaultCompleter, Emacs, FileBackedHistory, KeyCode,
-    KeyModifiers, Keybindings, Reedline, ReedlineEvent, ReedlineMenu,
+    default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
+    ColumnarMenu, DefaultCompleter, EditMode, Emacs, FileBackedHistory, KeyCode, KeyModifiers,
+    Keybindings, Reedline, ReedlineEvent, ReedlineMenu, Vi,
 };
 
 const MENU_NAME: &str = "completion_menu";
@@ -26,10 +27,22 @@ impl Repl {
 
         let completer = Self::create_completer(&config, &commands);
         let highlighter = ReplHighlighter::new(config.clone(), commands);
-        let keybindings = Self::create_keybindings();
         let history = Self::create_history()?;
         let menu = Self::create_menu();
-        let edit_mode = Box::new(Emacs::new(keybindings));
+        let edit_mode: Box<dyn EditMode> = if config.read().vi_keybindings {
+            let mut normal_keybindings = default_vi_normal_keybindings();
+            let mut insert_keybindings = default_vi_insert_keybindings();
+            Self::add_menu_keybindings(&mut normal_keybindings);
+            Self::add_menu_keybindings(&mut insert_keybindings);
+            Self::add_clear_keybindings(&mut normal_keybindings);
+            Self::add_clear_keybindings(&mut insert_keybindings);
+            Box::new(Vi::new(insert_keybindings, normal_keybindings))
+        } else {
+            let mut keybindings = default_emacs_keybindings();
+            Self::add_menu_keybindings(&mut keybindings);
+            Self::add_clear_keybindings(&mut keybindings);
+            Box::new(Emacs::new(keybindings))
+        };
         let mut editor = Reedline::create()
             .with_completer(Box::new(completer))
             .with_highlighter(Box::new(highlighter))
@@ -54,8 +67,7 @@ impl Repl {
         completer
     }
 
-    fn create_keybindings() -> Keybindings {
-        let mut keybindings = default_emacs_keybindings();
+    fn add_menu_keybindings(keybindings: &mut Keybindings) {
         keybindings.add_binding(
             KeyModifiers::NONE,
             KeyCode::Tab,
@@ -64,12 +76,14 @@ impl Repl {
                 ReedlineEvent::MenuNext,
             ]),
         );
+    }
+
+    fn add_clear_keybindings(keybindings: &mut Keybindings) {
         keybindings.add_binding(
             KeyModifiers::CONTROL,
             KeyCode::Char('l'),
             ReedlineEvent::ExecuteHostCommand(".clear screen".into()),
         );
-        keybindings
     }
 
     fn create_menu() -> ReedlineMenu {