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,
    last_request = 0,
    cache = {},
    debounce_ms = 1000, -- Wait 1 second between requests
    min_chars = 0,      -- Minimum characters before triggering completion
    max_context_lines = 10, -- Maximum number of context lines to send
}

-- Add at the top of the file with other state variables
local completion_cache = {
    last_line = nil,
    last_col = nil,
    suggestions = {},
    timestamp = 0,
    ttl = 30000 -- Cache TTL in milliseconds (30 seconds)
}

-- Add this function
local function get_cached_suggestion(line, col)
    local now = vim.loop.now()
    if completion_cache.last_line == line 
        and completion_cache.last_col == col 
        and (now - completion_cache.timestamp) < completion_cache.ttl then
        return completion_cache.suggestions
    end
    return nil
end

-- Add this function
local function cache_suggestion(line, col, suggestions)
    completion_cache.last_line = line
    completion_cache.last_col = col
    completion_cache.suggestions = suggestions
    completion_cache.timestamp = vim.loop.now()
end

-- Helper function to get visible lines around cursor
local function get_visible_lines()
    local win = vim.api.nvim_get_current_win()
    local cursor_line = vim.api.nvim_win_get_cursor(0)[1]
    local top_line = vim.fn.line('w0')
    local bottom_line = vim.fn.line('w$')
    
    -- Get all visible lines
    local lines = vim.api.nvim_buf_get_lines(0, top_line - 1, bottom_line, false)
    local relative_cursor = cursor_line - top_line + 1
    
    return lines, relative_cursor
end

-- Helper function to get current line indent
local function get_line_indent(line)
    local indent = line:match("^%s+") or ""
    return indent
end

-- 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

-- Check if we should rate limit
local function should_rate_limit()
    local now = vim.loop.now()
    if (now - current_suggestion.last_request) < current_suggestion.debounce_ms then
        return true
    end
    current_suggestion.last_request = now
    return false
end

-- Cache key generation
local function get_cache_key(context)
    return vim.fn.sha256(context)
end

-- Helper function to show suggestion
local function show_suggestion(suggestion, start_col)
    clear_suggestion()
    
    local cursor = vim.api.nvim_win_get_cursor(0)
    local line = cursor[1] - 1
    local current_line = vim.api.nvim_get_current_line()
    
    -- Get current line indent
    local indent = get_line_indent(current_line)
    
    -- Split suggestion into lines and apply indent to all lines except first
    local suggestion_lines = vim.split(suggestion, "\n")
    for i = 2, #suggestion_lines do
        suggestion_lines[i] = indent .. suggestion_lines[i]
    end
    
    if #suggestion_lines == 0 then return end
    
    -- Add support for different suggestion styles
    local config = require("gemini.config")
    local style = config.options.completion_style or "ghost" -- Add this to config

    -- Show ghost text with improved styling
    vim.api.nvim_buf_set_extmark(0, current_suggestion.namespace_id, line, start_col, {
        virt_text = {{suggestion_lines[1], 'GeminiSuggestion'}},
        virt_text_pos = 'overlay',
        hl_mode = 'combine',
        priority = 100, -- Add higher priority to ensure visibility
        right_gravity = false -- Ensure suggestion stays at cursor position
    })
    
    -- Show remaining lines as virtual lines
    if #suggestion_lines > 1 then
        local virt_lines = {}
        for i = 2, #suggestion_lines do
            table.insert(virt_lines, {{suggestion_lines[i], 'GeminiSuggestion'}})
        end
        
        vim.api.nvim_buf_set_extmark(0, current_suggestion.namespace_id, line, start_col, {
            virt_lines = virt_lines,
            virt_lines_above = false,
        })
    end
    
    current_suggestion.text = table.concat(suggestion_lines, "\n")
    current_suggestion.start_col = start_col
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 after the suggestion start
    if col < current_suggestion.start_col then
        return false
    end
    
    -- Get current line text
    local line_text = vim.api.nvim_buf_get_lines(0, line, line + 1, true)[1]
    local text_after_cursor = string.sub(line_text, col + 1)
    
    -- Split suggestion into lines
    local suggestion_lines = vim.split(current_suggestion.text, "\n")
    if #suggestion_lines == 0 then return false end
    
    -- For the first line, only insert the part that doesn't overlap
    local first_line = suggestion_lines[1]
    local common_length = 0
    while common_length < #text_after_cursor and common_length < #first_line do
        if string.sub(text_after_cursor, common_length + 1, common_length + 1) ==
           string.sub(first_line, common_length + 1, common_length + 1) then
            common_length = common_length + 1
        else
            break
        end
    end
    
    -- Only insert the non-overlapping part
    local to_insert = string.sub(first_line, common_length + 1)
    if to_insert ~= "" then
        vim.api.nvim_buf_set_text(
            0,
            line,
            col,
            line,
            col,
            {to_insert}
        )
    end
    
    -- Insert remaining lines below
    if #suggestion_lines > 1 then
        vim.api.nvim_buf_set_lines(
            0,
            line + 1,
            line + 1,
            false,
            vim.list_slice(suggestion_lines, 2)
        )
    end
    
    -- Move cursor to end of inserted text
    local final_line = line + #suggestion_lines - 1
    local final_col
    if final_line == line then
        final_col = col + #to_insert
    else
        final_col = #suggestion_lines[#suggestion_lines]
    end
    vim.api.nvim_win_set_cursor(0, {final_line + 1, final_col})
    
    -- Clear suggestion and cache for the current context
    clear_suggestion()
    current_suggestion.cache = {} -- Clear cache to force new suggestions
    
    -- Schedule next completion after a short delay
    vim.defer_fn(function()
        if vim.b.gemini_completion_enabled then
            M.trigger_completion()
        end
    end, 100)
    
    return true
end

-- Helper function to get relevant context around cursor
local function get_context_around_cursor(lines, current_line, max_lines)
    local context = {}
    local start_line = math.max(1, current_line - math.floor(max_lines/2))
    local end_line = math.min(#lines, current_line + math.floor(max_lines/2))
    
    -- Add file type and cursor position information
    local file_type = vim.bo.filetype
    local cursor_pos = vim.api.nvim_win_get_cursor(0)
    
    -- Add metadata to context
    table.insert(context, string.format("File type: %s", file_type))
    table.insert(context, string.format("Cursor position: line %d, col %d", cursor_pos[1], cursor_pos[2]))
    
    -- Add visible context with line numbers
    for i = start_line, end_line do
        if lines[i] and #lines[i] > 0 then
            table.insert(context, string.format("L%d: %s", i, lines[i]))
        end
    end
    
    return context
end

local function create_debounced_function(fn, wait)
    local timer = nil
    return function(...)
        local args = {...}
        if timer then
            vim.fn.timer_stop(timer)
        end
        timer = vim.fn.timer_start(wait, function()
            fn(unpack(args))
            timer = nil
        end)
    end
end

-- Use the improved debouncing
local trigger_completion_debounced = create_debounced_function(function()
    M.trigger_completion()
end, current_suggestion.debounce_ms)

function M.trigger_completion()
    -- Early exit conditions
    if not vim.b.gemini_completion_enabled then return end
    if vim.fn.pumvisible() ~= 0 then return end
    
    -- Get current line and cursor position
    local cursor = vim.api.nvim_win_get_cursor(0)
    local line = cursor[1] - 1
    local col = cursor[2]
    local current_line = vim.api.nvim_get_current_line()
    
    -- Check filetype exclusions
    local config = require("gemini.config")
    if vim.tbl_contains(config.options.completion.exclude_filetypes, vim.bo.filetype) then
        return
    end
    
    -- Get context
    local visible_lines, relative_cursor = get_visible_lines()
    local context = get_context_around_cursor(visible_lines, relative_cursor, config.options.completion.max_context_lines)
    
    -- Create prompt for API
    local prompt = table.concat(context, "\n") .. "\nPlease complete the following line:\n" .. current_line
    
    -- Make API request
    api.get_response(prompt, nil, function(response, error)
        if error then
            vim.notify("Completion error: " .. error, vim.log.levels.ERROR)
            return
        end
        
        if response then
            -- Extract the completion suggestion
            local suggestion = response:gsub("^%s*(.-)%s*$", "%1") -- Trim whitespace
            
            -- Show the suggestion
            vim.schedule(function()
                -- Only show if cursor position hasn't changed
                local new_cursor = vim.api.nvim_win_get_cursor(0)
                if new_cursor[1] - 1 == line and new_cursor[2] == col then
                    show_suggestion(suggestion, 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,
        blend = 15  -- Makes the ghost text slightly transparent
    })
    
    -- Map Tab to accept suggestion
    vim.keymap.set('i', '<Tab>', function()
        if not M.accept_suggestion() then
            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()
            if vim.b.gemini_completion_enabled then
                M.trigger_completion()
            end
        end
    })
    
    -- Clear suggestion on cursor move
    vim.api.nvim_create_autocmd({'CursorMovedI', 'CursorMoved'}, {
        pattern = "*",
        callback = function()
            clear_suggestion()
        end
    })
end

return M