accept additional arguments

1 files changed, 20 insertions, 9 deletions

diff --git a/vim-ai.vim b/vim-ai.vim
index f98fe1c..e6751de 100644
--- a/vim-ai.vim
+++ b/vim-ai.vim
@@ -1,19 +1,30 @@
 let g:openaiToken = system("cat ~/.config/openai.token")
 
-function! AIRun() range
-  execute a:firstline . ',' . a:lastline . 'd'
+function! AIRun(...) range
+  let selection = getline(a:firstline, a:lastline)
+  if a:0
+    let instruction = a:1 . ":"
+    call insert(selection, instruction, 0)
+  endif
+
+  let buff_lastline = line('$')
 
-  let selection = @*
-  call writefile(split(selection, "\n"), "/tmp/vim-ai.temp")
+  call writefile(selection, "/tmp/vim-ai.temp")
 
-  echo "Working..."
+  echo "Completing..."
   let output = system("cat /tmp/vim-ai.temp | openai complete - -t " . g:openaiToken)
+  let output = trim(output)
+
+  execute a:firstline . ',' . a:lastline . 'd'
+
+  if a:lastline == buff_lastline
+    execute "normal! o" . output . "\<Esc>"
+  else
+    execute "normal! O" . output . "\<Esc>"
+  endif
 
-  call feedkeys("i")
-  call feedkeys(output)
-  call feedkeys("\<Esc>")
 endfunction
 
-command! -range AI <line1>,<line2>call AIRun()
+command! -range -nargs=? AI <line1>,<line2>call AIRun(<f-args>)
 nnoremap <leader>o :call AIRun()<CR>
 vnoremap <leader>o :call AIRun()<CR>