" gpt.vim

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

if !exists("g:gpt_openai_api_key")
	echo "g:gpt_openai_api_key is not set. Please set it to your OpenAI API key."
	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_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_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.',
\}

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

let g:gpt_models = {
			\ "gpt-3.5-turbo": 4097,
			\ "gpt-3.5-turbo-16k": 16385,
			\ "gpt-3.5-turbo-1106": 16385,
			\ "gpt-4": 8192,
			\ "gpt-4-turbo": 128000,
			\ "gpt-4-turbo-preview": 128000,
			\ "gpt-4-32k": 32768,
			\ "gpt-4o": 128000,
			\ "gpt-4o-mini": 128000,
\}

function! GptInfo()  " {{{
	echo 'model: ' . g:gpt_model
	echo 'max_tokens: ' . string(g:gpt_max_tokens)
	echo 'temperature: ' . string(g:gpt_temperature)
	echo 'persona: ' . g:gpt_persona
	echo 'language: ' . g:gpt_lang
	echo 'split_direction: ' . g:gpt_split_direction
	echo 'split_ratio: ' . string(g:gpt_split_ratio)
	echo 'personas: ' . join(keys(g:gpt_personas), ', ')
	echo 'models: ' . join(keys(g:gpt_models), ', ')
endfunction

" }}}

function! GptDisplay(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! Gpt(prompt, persist) abort  " {{{
  " echo 'prompt: ' . a:prompt
  " echo 'persist: ' . a:persist
	silent py3file ~/.vim/python/gpt.py
endfunction

" }}}

function! GptSend(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

	let prompt = a:prompt . snippet

	call Gpt(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

" }}}

function! GptPersona(persona=0)  " {{{
	if a:persona == 0
		echo 'Personas: ' . join(keys(g:gpt_personas), ', ')
		return
	endif

	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

" }}}

function! GptModel(model=0)  " {{{
	if a:model == 0
		echo 'Models: ' . join(keys(g:gpt_models), ', ')
		return
	endif

	let models = keys(g:gpt_models)

	if index(models, a:model) != -1
		echo 'model set to: ' . a:model
		let g:gpt_model = a:model
	else
		let g:gpt_model = 'gpt-4o'
		echo a:model . ' is not a valid model. Defaulting to: ' . g:gpt_model
	end
endfunction

" }}}

command! -nargs=0 GptInfo call GptInfo()
command! -nargs=1 -range -bang Gpt call GptSend(<q-args>, '<bang>')
command! -nargs=? GptPersona call GptPersona(<q-args>)
command! -nargs=? GptModel call GptModel(<q-args>)