gemini/lua/gemini/completion.lua

local api = require("gemini.api")
local M = {}

-- State for managing current suggestion and debounce timer
local current_suggestion = {
    text = nil,
    start_col = nil,
    namespace_id = vim.api.nvim_create_namespace('gemini_suggestion'),
    timer = nil
}

-- Debug function
local function debug_print(...)
    vim.notify(string.format(...), vim.log.levels.INFO)
end

-- Helper function to clear suggestion
local function clear_suggestion()
    if current_suggestion.timer then
        vim.fn.timer_stop(current_suggestion.timer)
        current_suggestion.timer = nil
    end
    
    if current_suggestion.text then
        vim.api.nvim_buf_clear_namespace(0, current_suggestion.namespace_id, 0, -1)
        current_suggestion.text = nil
        current_suggestion.start_col = nil
    end
end

-- Helper function to show suggestion
local function show_suggestion(suggestion, start_col)
    clear_suggestion()
    
    -- Only use first line of suggestion
    suggestion = suggestion:match("^[^\n]*")
    if suggestion == "" then return end
    
    local line = vim.api.nvim_win_get_cursor(0)[1] - 1
    local line_text = vim.api.nvim_buf_get_lines(0, line, line + 1, true)[1]
    
    -- Get text before cursor on current line
    local prefix = string.sub(line_text, 1, start_col)
    
    -- Only show suggestion if it continues the current text
    if not suggestion:find("^" .. vim.pesc(prefix), 1, true) then
        suggestion = prefix .. suggestion
    end
    
    -- Remove the prefix from the suggestion to avoid duplication
    suggestion = suggestion:sub(#prefix + 1)
    
    if suggestion == "" then return end
    
    current_suggestion.text = suggestion
    current_suggestion.start_col = start_col
    
    -- Make the text virtual (doesn't affect actual buffer content)
    vim.api.nvim_buf_set_extmark(0, current_suggestion.namespace_id, line, start_col, {
        virt_text = {{suggestion, 'GeminiSuggestion'}},
        virt_text_pos = 'inline',
        virt_text_hide = true,
    })
    
    debug_print("Showing suggestion: %s", suggestion)
end

function M.accept_suggestion()
    if not (current_suggestion.text and current_suggestion.start_col) then
        return false
    end

    local line = vim.api.nvim_win_get_cursor(0)[1] - 1
    local col = vim.api.nvim_win_get_cursor(0)[2]
    
    -- Only accept if we're still at or before the suggestion start
    if col < current_suggestion.start_col then
        return false
    end
    
    vim.api.nvim_buf_set_text(
        0,
        line,
        col,
        line,
        col,
        {current_suggestion.text}
    )
    
    -- Move cursor to end of inserted text
    vim.api.nvim_win_set_cursor(0, {line + 1, col + #current_suggestion.text})
    
    clear_suggestion()
    return true
end

function M.trigger_completion()
    -- Clear any existing timer
    if current_suggestion.timer then
        vim.fn.timer_stop(current_suggestion.timer)
    end
    
    debug_print("Triggering completion...")
    
    -- Get current buffer and cursor info
    local cursor = vim.api.nvim_win_get_cursor(0)
    local line = cursor[1]
    local col = cursor[2]
    
    -- Get entire buffer content
    local lines = vim.api.nvim_buf_get_lines(0, 0, -1, false)
    local current_line = lines[line]
    
    -- Don't trigger completion if line is empty or cursor is at the start
    if col == 0 or current_line:match("^%s*$") then
        clear_suggestion()
        return
    end
    
    -- Get text before cursor on current line
    local prefix = string.sub(current_line, 1, col)
    
    -- Don't trigger on certain conditions
    if prefix:match("^%s*$") or -- empty or whitespace
       prefix:match("[%s%-]$") or -- ends with space or dash
       #prefix < 3 then -- too short
        clear_suggestion()
        return
    end
    
    debug_print("Getting completion for: %s", prefix)
    
    -- Construct context from buffer
    local context = {}
    
    -- Add up to 10 previous lines for context
    local start_line = math.max(1, line - 10)
    for i = start_line, line - 1 do
        table.insert(context, lines[i])
    end
    
    -- Add current line
    table.insert(context, current_line)
    
    -- Combine all lines
    local full_context = table.concat(context, "\n")
    
    -- Construct prompt for Gemini
    local prompt = string.format(
        "Complete this code. Return ONLY the completion that would naturally follow, no explanation:\n%s",
        full_context
    )
    
    -- Get completion from Gemini
    api.get_response(prompt, nil, function(response, error)
        if error then 
            debug_print("Completion error: %s", error)
            return 
        end
        
        if type(response) == "string" then
            vim.schedule(function()
                -- Check if cursor position is still valid
                local new_cursor = vim.api.nvim_win_get_cursor(0)
                if new_cursor[1] == line and new_cursor[2] >= col then
                    show_suggestion(response, col)
                end
            end)
        end
    end)
end

-- Setup function to create highlight group and keymaps
function M.setup()
    -- Create highlight group for suggestions
    vim.api.nvim_set_hl(0, 'GeminiSuggestion', {
        fg = '#666666',
        italic = true,
    })
    
    -- Map Tab to accept suggestion or behave normally
    vim.keymap.set('i', '<Tab>', function()
        if not M.accept_suggestion() then
            -- If no suggestion to accept, send regular Tab key
            vim.api.nvim_feedkeys(vim.api.nvim_replace_termcodes('<Tab>', true, true, true), 'n', true)
        end
    end, { expr = false, noremap = true })
    
    -- Set up autocommand for real-time completion
    vim.api.nvim_create_autocmd("TextChangedI", {
        pattern = "*",
        callback = function()
            M.trigger_completion()
        end
    })
    
    -- Clear suggestion on cursor move
    vim.api.nvim_create_autocmd({'CursorMovedI', 'CursorMoved'}, {
        pattern = "*",
        callback = function()
            clear_suggestion()
        end
    })
    
    debug_print("Gemini completion setup complete")
end

return M