local M = {}

local progress_win = nil
local progress_buf = nil
local job_id = nil

function M.show_progress(msg)
    if progress_win and vim.api.nvim_win_is_valid(progress_win) then
        local lines = vim.split(msg, "\n")
        vim.api.nvim_buf_set_lines(progress_buf, -1, -1, false, lines)
        local line_count = vim.api.nvim_buf_line_count(progress_buf)
        vim.api.nvim_win_set_cursor(progress_win, {line_count, 0})
        return
    end

    progress_buf = vim.api.nvim_create_buf(false, true)
    local lines = vim.split(msg, "\n")
    vim.api.nvim_buf_set_lines(progress_buf, 0, -1, false, lines)

    local width = math.min(math.floor(vim.o.columns * 0.4), 80)
    local height = math.min(math.floor(vim.o.lines * 0.4), 20)
    local row = 1
    local col = vim.o.columns - width - 2

    progress_win = vim.api.nvim_open_win(progress_buf, true, {
        relative = 'editor',
        width = width,
        height = height,
        row = row,
        col = col,
        style = 'minimal',
        border = 'rounded',
        title = ' Antigravity AI ',
        title_pos = 'center'
    })
    
    local function close_win()
        if job_id then
            vim.fn.jobstop(job_id)
            job_id = nil
        end
        if progress_win and vim.api.nvim_win_is_valid(progress_win) then
            vim.api.nvim_win_close(progress_win, true)
            progress_win = nil
        end
    end
    vim.keymap.set('n', 'q', close_win, { buffer = progress_buf, noremap = true, silent = true })
    vim.keymap.set('n', '<Esc>', close_win, { buffer = progress_buf, noremap = true, silent = true })
end

function M.close_progress()
    if progress_win and vim.api.nvim_win_is_valid(progress_win) then
        vim.api.nvim_win_close(progress_win, true)
        progress_win = nil
    end
end

local function execute_agent(prompt)
    local cmd = { "agy", "--new-project", "-i", prompt }
    
    M.show_progress("Starting Antigravity...")
    
    -- termopen requires an empty and unmodified buffer
    vim.api.nvim_buf_set_lines(progress_buf, 0, -1, false, {})
    vim.bo[progress_buf].modified = false
    
    -- termopen will attach to the current buffer (which show_progress just focused)
    job_id = vim.fn.termopen(cmd, {
        on_exit = function(_, code)
            job_id = nil
        end
    })
    vim.cmd('startinsert')
end

function M.ask_agent()
    local file = vim.api.nvim_buf_get_name(0)
    vim.ui.input({ prompt = 'Ask Agent: ' }, function(input)
        if not input or input == '' then return end
        
        local final_input = input
        if file and file ~= "" then
            final_input = input .. string.format(" [Context file: %s]", file)
        end
        
        execute_agent(final_input)
    end)
end

function M.setup()
    vim.keymap.set('n', '<leader>ai', M.ask_agent, { desc = "Ask Antigravity Agent", noremap = true, silent = true })
    vim.keymap.set('v', '<leader>ai', function()
        local _, csrow, cscol, _ = unpack(vim.fn.getpos("'<"))
        local _, cerow, cecol, _ = unpack(vim.fn.getpos("'>"))
        
        vim.ui.input({ prompt = 'Ask Agent (about selection): ' }, function(input)
            if not input or input == '' then return end
            
            local file = vim.api.nvim_buf_get_name(0)
            local context = string.format(" [Context: %s lines %d-%d]", file, csrow, cerow)
            local final_input = input .. context
            
            execute_agent(final_input)
        end)
    end, { desc = "Ask Antigravity Agent about selection", noremap = true, silent = true })
end

return M