gemini/lua/gemini/chat.lua

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

-- Store chat state
local state = {
    bufnr = nil,
    winnr = nil,
    current_context = nil,
    conversation_history = {}
}

local function setup_chat_highlighting(bufnr)
    vim.api.nvim_buf_set_option(bufnr, 'syntax', '')
    pcall(vim.treesitter.start, bufnr, "markdown")

    vim.cmd([[
        highlight default GeminiUser guifg=#EBCB8B gui=bold
        highlight default GeminiSeparator guifg=#616E88 gui=bold
        syntax match GeminiUser /^User:.*$/
        syntax match GeminiSeparator /^━━━━━━━━━━━━━━━━━━━━━━━━━━$/
    ]])
end

local function setup_buffer_options(bufnr)
    local options = {
        buftype = 'nofile',
        filetype = 'markdown',
        buflisted = false,
        swapfile = false,
        bufhidden = 'wipe',
        modifiable = false,
    }

    for option, value in pairs(options) do
        vim.api.nvim_buf_set_option(bufnr, option, value)
    end
end

local function setup_buffer_autocmds(bufnr)
    local augroup = vim.api.nvim_create_augroup('GeminiChatBuffer', { clear = true })
    vim.api.nvim_create_autocmd({'BufReadCmd', 'FileReadCmd', 'BufWriteCmd'}, {
        group = augroup,
        buffer = bufnr,
        callback = function()
            vim.notify('File operations are not allowed in the chat window', vim.log.levels.WARN)
            return true
        end
    })
end

local function setup_buffer_keymaps(bufnr)
    local opts = { buffer = bufnr, nowait = true }
    
    -- Disable file operations
    local operations = {'e', 'edit', 'w', 'write', 'sp', 'split', 'vs', 'vsplit',
                       'new', 'vnew', 'read', 'update', 'saveas'}
    for _, op in ipairs(operations) do
        vim.keymap.set('n', ':' .. op, function()
            vim.notify('Operation not allowed in chat window', vim.log.levels.WARN)
        end, opts)
    end

    -- Disable command mode
    vim.keymap.set('n', ':', function()
        vim.notify('Command mode disabled in chat window', vim.log.levels.WARN)
    end, opts)

    -- Set chat-specific keymaps
    local mappings = config.options.mappings
    vim.keymap.set('n', mappings.close, function()
        vim.api.nvim_win_close(state.winnr, true)
        state.winnr = nil
    end, opts)

    vim.keymap.set('n', mappings.return_focus, function()
        vim.cmd('wincmd p')
    end, opts)

    vim.keymap.set('n', mappings.new_query, function()
        require("gemini").prompt_query(state.current_context)
    end, opts)
end

function M.create_window()
    if not state.bufnr or not vim.api.nvim_buf_is_valid(state.bufnr) then
        state.bufnr = vim.api.nvim_create_buf(false, true)
        setup_buffer_options(state.bufnr)
        setup_buffer_autocmds(state.bufnr)
        setup_buffer_keymaps(state.bufnr)
    end

    local win_opts = config.options.window
    if not state.winnr or not vim.api.nvim_win_is_valid(state.winnr) then
        state.winnr = vim.api.nvim_open_win(state.bufnr, true, {
            relative = "editor",
            width = win_opts.width(),
            height = win_opts.height(),
            row = 0,
            col = vim.o.columns - win_opts.width(),
            border = win_opts.border,
            title = win_opts.title,
            title_pos = win_opts.title_pos,
            style = "minimal"
        })

        vim.api.nvim_win_set_option(state.winnr, 'wrap', true)
        vim.api.nvim_win_set_option(state.winnr, 'linebreak', true)
        vim.api.nvim_win_set_option(state.winnr, 'breakindent', true)

        setup_chat_highlighting(state.bufnr)
    end
end

function M.update_content(content, is_new_chat, is_thinking)
    vim.api.nvim_buf_set_option(state.bufnr, 'modifiable', true)
    local lines = vim.split(content, "\n")

    if is_thinking then
        -- For thinking message, just replace/append without affecting history
        vim.api.nvim_buf_set_lines(state.bufnr, 0, -1, false, lines)
    else
        if is_new_chat then
            -- For new messages, show the complete history
            local display_lines = {}
            for i, msg in ipairs(state.conversation_history) do
                if i > 1 then  -- Add separator before messages (except first)
                    table.insert(display_lines, "━━━━━━━━━━━━━━━━━━━━━━━━━━")
                end
                local msg_lines = vim.split(msg.content, "\n")
                for _, line in ipairs(msg_lines) do
                    table.insert(display_lines, line)
                end
            end
            -- Add the new content
            if #state.conversation_history > 0 then
                table.insert(display_lines, "━━━━━━━━━━━━━━━━━━━━━━━━━━")
            end
            for _, line in ipairs(lines) do
                table.insert(display_lines, line)
            end
            vim.api.nvim_buf_set_lines(state.bufnr, 0, -1, false, display_lines)
        else
            -- For subsequent messages in the same chat
            if #state.conversation_history > 0 then
                table.insert(lines, 1, "━━━━━━━━━━━━━━━━━━━━━━━━━━")
            end
            vim.api.nvim_buf_set_lines(state.bufnr, -1, -1, false, lines)
        end
    end

    vim.api.nvim_buf_set_option(state.bufnr, 'modifiable', false)

    -- Scroll to show new content
    local line_count = vim.api.nvim_buf_line_count(state.bufnr)
    local content_lines = #vim.split(content, "\n")
    local start_line = line_count - content_lines + 1
    vim.api.nvim_win_set_cursor(state.winnr, {start_line, 0})
    vim.cmd('normal! zt')
    vim.cmd('wincmd p')
end

function M.set_context(context)
    state.current_context = context
end

function M.clear()
    if state.bufnr and vim.api.nvim_buf_is_valid(state.bufnr) then
        vim.api.nvim_buf_set_lines(state.bufnr, 0, -1, false, {})
    end
    state.conversation_history = {}  -- Clear conversation history
end

-- Add functions to manage conversation history
function M.add_message(role, content)
    table.insert(state.conversation_history, {role = role, content = content})
end

function M.get_conversation_history()
    return state.conversation_history
end

return M