local K = require('core.keymap')

local explain_prompt = [[
Explain this Code short and Precise:
$fname
```$ftype
$buf
```
---
]]

local generate_prompt = [[
<|fim_prefix|>$before<|fim_suffix|>$after<|fim_middle|>
]]

return {
    'nomnivore/ollama.nvim',
    config = function()
        local ollama = require('ollama')
        ollama.setup({
            model = 'codegemma:instruct',
            prompts = {
                Ask_About_Code = false,
                Modify_Code = false,
                Raw = false,
                Simplify_Code = false,
                Explain_Code = {
                    model = 'codegemma:instruct',
                    prompt = explain_prompt,
                    input_label = '>',
                    action = require('ollama.actions.factory').create_action({
                        display = true,
                        window = 'vsplit',
                    }),
                },
                Generate_Code = {
                    model = 'codegemma:2b-code',
                    prompt = generate_prompt,
                    input_label = '>',
                    action = 'display',
                    options = {
                        num_predict = 128,
                        temperature = 0.0,
                        top_p = 0.9,
                        stop = { '<|file_separator|>' },
                    },
                },
            },
        })

        vim.keymap.set(
            K.OLLAMA.mode,
            K.OLLAMA.shortcut,
            ollama.prompt,
            { desc = K.OLLAMA.description }
        )
    end,
}