neoai.lua

local ui = require("neoai.ui")
local chat = require("neoai.chat")
local inject = require("neoai.inject")
local ChatHistory = require("neoai.chat.history")
local config = require("neoai.config")
local shortcuts = require("neoai.shortcuts")
local utils = require("neoai.utils")

local M = {}

---Sets the context
---@param line1 integer | nil The first line number to use, nil will use '<
---@param line2 integer | nil The second line number to use, nil will use '>
---@return integer, integer line1 and line2
local set_context = function (line1, line2)
    local buffer = vim.api.nvim_get_current_buf()

    line1 = line1 or vim.fn.line("'<")
    line2 = line2 or vim.fn.line("'>")

    chat.set_context(buffer, line1, line2)

    return line1, line2
end

-- TODO
local function setup_colors()
    local hl_group_name = "NeoAIInput"
    local color = { guifg = "#61afef" }
    vim.api.nvim_command("highlight " .. hl_group_name .. " guifg=" .. color.guifg)
end

---Setup NeoAI given options
---@param options Options | nil The options for the plugin
M.setup = function(options)
    config.setup(options)
    chat.setup_models()
    shortcuts.setup_shortcuts()
end

---Toggles opening and closing neoai window
---@param toggle boolean | nil If true will open GUI and false will close, nil will toggle
---@param prompt string | nil If set then this prompt will be sent to the GUI if toggling on
---@return boolean true if opened and false if closed
M.toggle = function(toggle, prompt)
    local open = (toggle ~= "" and toggle) or (toggle == "" and not ui.is_open())
    if open then
        -- Open
        ui.create_ui()
        if prompt ~= nil then
            ui.send_prompt(prompt)
        end
        return true
    else
        -- Close
        ui.destroy_ui()
        return false
    end
end

---Smart focus, if closed then will open on GUI, if opened and focused then it
---will close GUI and if opened and not focused then it will focus on the GUI.
---@param prompt string The prompt to inject, to inject no prompt just do empty string
M.smart_toggle = function(prompt)
    local send_args = function()
        if not utils.is_empty(prompt) then
            ui.send_prompt(prompt)
        end
    end
    if ui.is_open() then
        if ui.is_focused() then
            M.toggle(false)
        else
            ui.focus()
            send_args()
        end
    else
        M.toggle(true)
        send_args()
    end
end

---Toggles the GUI in Context mode
---@param toggle boolean | nil True will force open, False will force close, nil will toggle
---@param prompt string The prompt to inject into the GUI if any (otherwise specify an empty string)
---@param line1 integer | nil The first line number for context range otherwise will use visual selection
---@param line2 integer | nil The second line number for context range otherwise will use visual selection
---@return boolean True if was opened
M.context_toggle = function (toggle, prompt, line1, line2)
    set_context(line1, line2)
    return M.toggle(toggle, prompt)
end

---Smart toggles context GUI
---@param prompt string The prompt to inject into the GUI if any (otherwise specify an empty string)
---@param line1 integer | nil The first line number for context range otherwise will use visual selection
---@param line2 integer | nil The second line number for context range otherwise will use visual selection
M.context_smart_toggle = function (prompt, line1, line2)
    set_context(line1, line2)
    M.smart_toggle(prompt)
end

---Sends prompt and injects the response straight back into the buffer without
---opening the GUI
---@param prompt string The prompt to send to the AI
---@param strip_function (fun(output: string): string) | nil A function that strips the output
---@param start_line integer | nil The line to start injecting onto (After inserting 2 newlines), nil = current selected line
M.inject = function(prompt, strip_function, start_line)
    chat.new_chat_history()

    strip_function = strip_function or function(x)
        return x
    end

    local current_line = start_line or vim.api.nvim_win_get_cursor(0)[1]
    chat.send_prompt(
    prompt,
    function(txt, _)
        -- Get differences between text
        local txt1 = strip_function(chat.get_current_output())
        local txt2 = strip_function(table.concat({ chat.get_current_output(), txt }, ""))

        inject.append_to_buffer(string.sub(txt2, #txt1 + 1), current_line)
    end,
    false,
    function(_)
        inject.current_line = nil
        vim.notify("NeoAI: Done generating AI response", vim.log.levels.INFO)
    end
    )
end

---Same as inject except uses a context
---@param prompt string The prompt to send to the AI
---@param strip_function (fun(output: string): string) | nil A function that strips the output
---@param line1 integer | nil The first line num in the range if nil will use '<
---@param line2 integer | nil The second line num in the range if nil will use '>
M.context_inject = function(prompt, strip_function, line1, line2)
    line1, line2 = set_context(line1, line2)
    M.inject(prompt, strip_function, line2)
end

local make_prompt = function()
    local ret = ""
    for i, model_obj in ipairs(config.options.models) do
        ret = ret .. i-1 .. "." .. model_obj.name .. "\n"
    end
    ret = ret .. "Choose:"
    return ret
end

M.set_source = function()
    vim.ui.input({ prompt = make_prompt() }, function(text)
        if text == nil or string.len(text) == 0 then return end
        local n = tonumber(text)
        if n ~= nil and n >= 0 and n < #config.options.models then
            config.options.selected_model_index = n
            chat.setup_models()
        end
    end)
end

return M