From 2b5b1928ea727c2d69ca9b8fd247c3f9ef931fe1 Mon Sep 17 00:00:00 2001
From: Anton Kastritskiy <antonk52@meta.com>
Date: Thu, 12 Oct 2023 14:42:24 +0100
Subject: [PATCH] [vim] refine ai completion keymaps and setup

---
 nvim/init.lua                    |  8 ++++
 nvim/lua/antonk52/completion.lua | 65 ++++++++++++++++++++++----------
 2 files changed, 53 insertions(+), 20 deletions(-)

diff --git a/nvim/init.lua b/nvim/init.lua
index dab8ed5..22a7b01 100644
--- a/nvim/init.lua
+++ b/nvim/init.lua
@@ -44,6 +44,14 @@ local plugins = {
                         auto_trigger = true,
                     },
                 })
+                local c = require('copilot.suggestion')
+                require('antonk52.completion').update_ai_completion({
+                    is_visible = c.is_visible,
+                    accept = c.accept,
+                    accept_word = c.accept_word,
+                    accept_line = c.accept_line,
+                    dismiss = c.dismiss,
+                })
             else
                 return print('no copilot at work')
             end
diff --git a/nvim/lua/antonk52/completion.lua b/nvim/lua/antonk52/completion.lua
index 97fa3ef..39df492 100644
--- a/nvim/lua/antonk52/completion.lua
+++ b/nvim/lua/antonk52/completion.lua
@@ -15,11 +15,33 @@ local snippet_sources = {
     { name = 'buffer', keyword_length = 3 },
 }
 
+local noop = function() end
+
+---@class AI_completion
+local AI = {
+    is_visible = function()
+        return false
+    end,
+    accept = noop,
+    accept_word = noop,
+    accept_line = noop,
+    dismiss = noop,
+}
+
+---@param opts AI_completion
+function M.update_ai_completion(opts)
+    AI.is_visible = opts.is_visible
+    AI.accept = opts.accept
+    AI.accept_word = opts.accept_word
+    AI.accept_line = opts.accept_line
+    AI.dismiss = opts.dismiss
+end
+
 function M.setup()
     local mapping = {
         ['<Tab>'] = function(fallback)
-            if vim.env.WORK == nil and require('copilot.suggestion').is_visible() then
-                require('copilot.suggestion').accept()
+            if AI.is_visible() then
+                AI.accept()
             elseif cmp.visible() then
                 cmp.select_next_item()
             else
@@ -33,6 +55,27 @@ function M.setup()
                 fallback()
             end
         end,
+        ['<C-d>'] = function(fallback)
+            if AI.is_visible() then
+                AI.dismiss()
+            else
+                fallback()
+            end
+        end,
+        ['<C-e>'] = function(fallback)
+            if AI.is_visible() then
+                AI.accept_word()
+            else
+                fallback()
+            end
+        end,
+        ['<C-r>'] = function(fallback)
+            if AI.is_visible() then
+                AI.accept_line()
+            else
+                fallback()
+            end
+        end,
         ['<C-y>'] = cmp.mapping.confirm(),
         ['<CR>'] = cmp.mapping.confirm(),
         ['<C-j>'] = cmp.mapping(function(fallback)
@@ -51,24 +94,6 @@ function M.setup()
         end),
     }
 
-    if vim.env.WORK == nil then
-        mapping['<C-w>'] = function()
-            if require('copilot.suggestion').is_visible() then
-                require('copilot.suggestion').accept_word()
-            end
-        end
-        mapping['<C-e>'] = function()
-            if require('copilot.suggestion').is_visible() then
-                require('copilot.suggestion').accept_line()
-            end
-        end
-        mapping['<C-q>'] = function()
-            if require('copilot.suggestion').is_visible() then
-                require('copilot.suggestion').dismiss()
-            end
-        end
-    end
-
     cmp.setup({
         snippet = {
             expand = function(arg)