1 files changed, 21 insertions, 72 deletions
diff --git a/plugin/vim-ai.vim b/plugin/vim-ai.vim
index c183aa5..0f644b8 100644
--- a/plugin/vim-ai.vim
+++ b/plugin/vim-ai.vim
@@ -2,10 +2,6 @@ let s:plugin_root = expand('<sfile>:p:h:h')
 let s:complete_py = s:plugin_root . "/py/complete.py"
 let s:chat_py = s:plugin_root . "/py/chat.py"
 
-if !exists('g:vim_ai_debug')
-  let g:vim_ai_debug = 0
-endif
-
 function! ScratchWindow()
   below new
   setlocal buftype=nofile
@@ -14,90 +10,43 @@ function! ScratchWindow()
   setlocal ft=aichat
 endfunction
 
-function! AIRun(...) range
-  let lines = trim(join(getline(a:firstline, a:lastline), "\n"))
-  let selection = trim(@*)
-  let is_selection = lines != "" && lines == selection
-  let has_instruction = a:0
-  let prompt = ""
-  if has_instruction
-    if is_selection
-      let prompt = a:1 . ":\n" . lines
-    else
-      let prompt = a:1
-    endif
-  else
-    let prompt = lines
-  endif
+function! MakePrompt(lines, args)
+  let lines = trim(join(a:lines, "\n"))
+  let is_selection = lines != "" && lines == trim(@*)
+  let instruction = trim(get(a:args, 0))
+  let delimiter = instruction != "" && is_selection ? ":\n" : ""
+  let selection = is_selection ? lines : ""
+  let prompt = join([instruction, delimiter, selection], "")
+  return prompt
+endfunction
 
-  if g:vim_ai_debug
-    echo "Prompt:\n" . prompt . "\n"
-  endif
-  echo "Completing..."
-  " WORKAROUND: without sleep is echo on prev line not displayed (when combining with py3)
-  execute 'silent sleep 1m'
+function! AIRun(...) range
+  let prompt = MakePrompt(getline(a:firstline, a:lastline), a:000)
+  set paste
+  execute "normal! " . a:lastline . "Go"
   execute "py3file " . s:complete_py
-  let output = py3eval('output')
-  let output = trim(output)
-
   execute "normal! " . a:lastline . "G"
-  set paste
-  execute "normal! o" . output . "\<Esc>"
   set nopaste
-  execute "normal! " . a:lastline . "G"
 endfunction
 
 function! AIEditRun(...) range
-  let has_instruction = a:0
-  let prompt = trim(join(getline(a:firstline, a:lastline), "\n"))
-  if has_instruction
-    let prompt = a:1 . ":\n" . prompt
-  endif
-
-  let buff_lastline = line('$')
-
-  if g:vim_ai_debug
-    echo "Prompt:\n" . prompt . "\n"
-  endif
-
-  echo "Editing..."
-  " WORKAROUND: without sleep is echo on prev line not displayed (when combining with py3)
-  execute 'silent sleep 1m'
-  execute "py3file " . s:complete_py
-  let output = py3eval('output')
-  let output = trim(output)
-
-  execute a:firstline . ',' . a:lastline . 'd'
-
+  let prompt = MakePrompt(getline(a:firstline, a:lastline), a:000)
   set paste
-  if a:lastline == buff_lastline
-    execute "normal! o" . output . "\<Esc>"
-  else
-    execute "normal! O" . output . "\<Esc>"
-  endif
+  execute "normal! " . a:firstline . "GV" . a:lastline . "Gc"
+  execute "py3file " . s:complete_py
   set nopaste
 endfunction
 
 function! AIChatRun(...) range
-  let lines = trim(join(getline(a:firstline, a:lastline), "\n"))
-  let is_selection = lines != "" && lines == trim(@*)
-  let instruction = a:0 ? trim(a:1) : ""
-
+  let lines = getline(a:firstline, a:lastline)
   set paste
-  if search('^>>> user$', 'nw') == 0
-    " outside of chat window
+  let is_outside_of_chat_window = search('^>>> user$', 'nw') == 0
+  if is_outside_of_chat_window
     call ScratchWindow()
-    let delimiter = instruction != "" && is_selection ? ":\n" : ""
-    let selection = is_selection ? lines : ""
-    let prompt = join([instruction, delimiter, selection], "")
-    execute "normal i>>> user\<Enter>\<Enter>" . prompt
-    if prompt == ""
-      " empty prompt, just opens chat window TODO: handle in python
-      return
-    endif
+    let prompt = MakePrompt(lines, a:000)
+    execute "normal i>>> user\n\n" . prompt
   endif
 
-  echo "Answering..."
   execute "py3file " . s:chat_py
   set nopaste
 endfunction