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

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

-- Helper function to clear current suggestion
local function clear_suggestion()
    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
    current_suggestion.text = suggestion
    current_suggestion.start_col = start_col
    
    -- Show suggestion in a different color
    vim.api.nvim_buf_add_highlight(0, current_suggestion.namespace_id, 'GeminiSuggestion', line, start_col, start_col + #suggestion)
    vim.api.nvim_buf_set_text(0, line, start_col, line, start_col, {suggestion})
    
    -- 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 = 'overlay',
    })
end

-- Function to accept current suggestion
function M.accept_suggestion()
    if current_suggestion.text and current_suggestion.start_col then
        local line = vim.api.nvim_win_get_cursor(0)[1] - 1
        vim.api.nvim_buf_set_text(0, line, current_suggestion.start_col, line, current_suggestion.start_col, {current_suggestion.text})
        clear_suggestion()
        return true
    end
    return false
end

function M.trigger_completion()
    local cursor = vim.api.nvim_win_get_cursor(0)
    local line = cursor[1]
    local col = cursor[2]
    
    -- Get current buffer content for context
    local lines = vim.api.nvim_buf_get_lines(0, math.max(0, line - 10), line, false)
    local prefix = table.concat(lines, "\n")
    
    -- Get current line up to cursor
    local current_line = vim.api.nvim_get_current_line()
    local before_cursor = string.sub(current_line, 1, col)
    
    -- Construct prompt for Gemini
    local prompt = string.format(
        "Complete this code. Only provide the completion, no explanation:\n%s\n%s",
        prefix,
        before_cursor
    )
    
    -- Get completion from Gemini
    api.get_response(prompt, nil, function(response, error)
        if error then
            vim.notify("Completion error: " .. error, vim.log.levels.ERROR)
            return
        end
        
        if type(response) == "string" then
            -- Get first non-empty line as suggestion
            local suggestion = response:match("^%s*(.-)%s*$")
            if suggestion and #suggestion > 0 then
                vim.schedule(function()
                    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 = '#888888',
        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 })
    
    -- Clear suggestion on cursor move
    vim.api.nvim_create_autocmd({'CursorMovedI', 'CursorMoved'}, {
        callback = function()
            clear_suggestion()
        end
    })
end

return M