" ChatGPT Vim Plugin

if !has('python3')
  echo "Python 3 support is required for ChatGPT plugin"
  finish
endif

if !exists("g:gpt_max_tokens")
  let g:gpt_max_tokens = 2000
endif

if !exists("g:gpt_temperature")
  let g:gpt_temperature = 0.7
endif

if !exists("g:gpt_model")
  let g:gpt_model = 'gpt-4o'
endif

if !exists("g:gpt_lang")
  let g:gpt_lang = 'English'
endif

if !exists("g:gpt_split_direction")
  let g:gpt_split_direction = 'horizontal'
endif

if !exists("g:gpt_split_ratio")
  let g:gpt_split_ratio = 3
endif

if !exists("g:gpt_persona")
  let g:gpt_persona = 'default'
endif

let g:gpt_templates = {
\ 'Rewrite': 'Rewrite this more idiomatically',
\ 'Review': 'Review this code',
\ 'Document': 'Return documentation following language pattern conventions',
\ 'Explain': 'Explain how this works',
\ 'Test': 'Write a test',
\ 'Fix': 'Fix this error',
\}

let g:gpt_keys = keys(g:gpt_templates)

let g:gpt_personas = {
\ "default": '',
\ "bob": 'You are a helpful expert programmer we are working together to solve complex coding challenges, and I need your help. Please make sure to wrap all code blocks in ``` annotate the programming language you are using.',
\}

function! DisplayChatGPTResponse(response, finish_reason, chat_gpt_session_id)  " {{{
  let response = a:response
  let finish_reason = a:finish_reason

  let chat_gpt_session_id = a:chat_gpt_session_id

  if !bufexists(chat_gpt_session_id)
    if g:gpt_split_direction ==# 'vertical'
      silent execute winwidth(0)/g:gpt_split_ratio.'vnew '. chat_gpt_session_id
    else
      silent execute winheight(0)/g:gpt_split_ratio.'new '. chat_gpt_session_id
    endif
    call setbufvar(chat_gpt_session_id, '&buftype', 'nofile')
    call setbufvar(chat_gpt_session_id, '&bufhidden', 'hide')
    call setbufvar(chat_gpt_session_id, '&swapfile', 0)
    setlocal modifiable
    setlocal wrap
    setlocal linebreak
    call setbufvar(chat_gpt_session_id, '&ft', 'markdown')
    call setbufvar(chat_gpt_session_id, '&syntax', 'markdown')
  endif

  if bufwinnr(chat_gpt_session_id) == -1
    if g:gpt_split_direction ==# 'vertical'
      execute winwidth(0)/g:gpt_split_ratio.'vsplit ' . chat_gpt_session_id
    else
      execute winheight(0)/g:gpt_split_ratio.'split ' . chat_gpt_session_id
    endif
  endif

  let last_lines = getbufline(chat_gpt_session_id, '$')
  let last_line = empty(last_lines) ? '' : last_lines[-1]

  let new_lines = substitute(last_line . response, '\n', '\r\n\r', 'g')
  let lines = split(new_lines, '\n')

  let clean_lines = []
  for line in lines
    call add(clean_lines, substitute(line, '\r', '', 'g'))
  endfor

  call setbufline(chat_gpt_session_id, '$', clean_lines)

  execute bufwinnr(chat_gpt_session_id) . 'wincmd w'
  " Move the viewport to the bottom of the buffer
  normal! G
  call cursor('$', 1)

  if finish_reason != ''
    wincmd p
  endif
endfunction

" }}}

function! ChatGPT(prompt, persist) abort  " {{{
  " echo 'prompt: ' . a:prompt
  " echo 'persist: ' . a:persist

  silent py3file ~/.vim/python/gpt.py
endfunction

" }}}

function! SendToChatGPT(prompt, bang) abort  " {{{

	let persist = (a:bang ==# '!') ? 1 : 0

    let save_cursor = getcurpos()
    let [current_line, current_col] = getcurpos()[1:2]
    let save_reg = @@
    let save_regtype = getregtype('@')

    let [line_start, col_start] = getpos("'<")[1:2]
    let [line_end, col_end] = getpos("'>")[1:2]

    if (col_end - col_start > 0 || line_end - line_start > 0) &&
       \ (current_line == line_start && current_col == col_start ||
       \  current_line == line_end && current_col == col_end)

        let current_line_start = line_start
        let current_line_end = line_end

        if current_line_start == line_start && current_line_end == line_end
            execute 'normal! ' . line_start . 'G' . col_start . '|v' . line_end . 'G' . col_end . '|y'
            let snippet = "\n" . '```' . &syntax . "\n" . @@ . "\n" . '```'
        else
            let snippet = ''
        endif
    else
        let snippet = ''
    endif

	if has_key(g:gpt_templates, a:prompt)
		let prompt = g:gpt_templates[a:prompt]
	else
		let prompt = a:prompt
	endif

	let prompt = prompt . snippet

	call ChatGPT(prompt, persist)

    let @@ = save_reg
    call setreg('@', save_reg, save_regtype)

    let curpos = getcurpos()
    call setpos("'<", curpos)
    call setpos("'>", curpos)
    call setpos('.', save_cursor)
endfunction

" }}}

command! -range -bang -nargs=1 Gpt call SendToChatGPT(<q-args>, '<bang>')

for i in range(len(g:gpt_keys))
  execute "command! -range -bang -nargs=0 " . g:gpt_keys[i] . " call SendToChatGPT('" . g:gpt_keys[i] . "', '<bang>')"
endfor

function! SetPersona(persona)  " {{{
    let personas = keys(g:gpt_personas)
    if index(personas, a:persona) != -1
      echo 'Persona set to: ' . a:persona
      let g:gpt_persona = a:persona
    else
      let g:gpt_persona = 'default'
      echo a:persona . ' is not a valid persona. Defaulting to: ' . g:gpt_persona
    end
endfunction

" }}}

command! -nargs=1 Persona call SetPersona(<q-args>)