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 send_udp(payload)
    local uv = vim.uv or vim.loop
    local client = uv.new_udp()
    if client then
        local ok, data = pcall(vim.fn.json_encode, payload)
        if ok then
            client:send(data, "127.0.0.1", 3002, function(err)
                client:close()
            end)
        else
            client:close()
        end
    end
end

function M.prompt_user(prompt_text, options, request_id)
    if options and #options > 0 then
        vim.ui.select(options, { prompt = prompt_text }, function(choice, idx)
            send_udp({
                event = "agent_prompt_response",
                request_id = request_id,
                response = choice,
                index = idx,
                session_id = vim.v.servername
            })
        end)
    else
        vim.ui.input({ prompt = prompt_text }, function(input)
            send_udp({
                event = "agent_prompt_response",
                request_id = request_id,
                response = input,
                session_id = vim.v.servername
            })
        end)
    end
end

local ns_id = vim.api.nvim_create_namespace("gemini_ghost_text")
local active_diffs = {}

function M.preview_diff(diff_id, bufnr, start_row, end_row, new_lines)
    bufnr = (bufnr == 0 or bufnr == nil) and vim.api.nvim_get_current_buf() or bufnr
    local virt_lines = {}
    for _, line in ipairs(new_lines) do
        table.insert(virt_lines, {{line, "DiffAdd"}})
    end
    
    local extmark_id = vim.api.nvim_buf_set_extmark(bufnr, ns_id, start_row, 0, {
        virt_lines = virt_lines,
        virt_lines_above = true,
    })
    
    active_diffs[diff_id] = {
        bufnr = bufnr,
        start_row = start_row,
        end_row = end_row,
        new_lines = new_lines,
        extmark_id = extmark_id
    }
    
    if pcall(require, "notify") then
        require("notify")("Ghost text diff added. Use <leader>aa to accept or <leader>ar to reject.", "info", { title = "Antigravity" })
    else
        print("Ghost text diff added. Use <leader>aa to accept or <leader>ar to reject.")
    end
end

function M.accept_diff()
    local diff_id, diff = next(active_diffs)
    if not diff then
        print("No active diffs to accept.")
        return
    end
    
    vim.api.nvim_buf_set_lines(diff.bufnr, diff.start_row, diff.end_row, false, diff.new_lines)
    vim.api.nvim_buf_del_extmark(diff.bufnr, ns_id, diff.extmark_id)
    active_diffs[diff_id] = nil
    
    send_udp({ event = "agent_diff_accepted", diff_id = diff_id, session_id = vim.v.servername })
    print("Diff accepted.")
end

function M.reject_diff()
    local diff_id, diff = next(active_diffs)
    if not diff then
        print("No active diffs to reject.")
        return
    end
    
    vim.api.nvim_buf_del_extmark(diff.bufnr, ns_id, diff.extmark_id)
    active_diffs[diff_id] = nil
    
    send_udp({ event = "agent_diff_rejected", diff_id = diff_id, session_id = vim.v.servername })
    print("Diff rejected.")
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 cursor = vim.api.nvim_win_get_cursor(0)
        send_udp({
            event = "agent_ask",
            prompt = input,
            file = file,
            line = cursor[1],
            col = cursor[2],
            session_id = vim.v.servername
        })
        M.show_progress("Sent to Antigravity via UDP...")
        vim.defer_fn(M.close_progress, 1500)
    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)
            send_udp({
                event = "agent_ask_selection",
                prompt = input,
                file = file,
                start_row = csrow,
                end_row = cerow,
                session_id = vim.v.servername
            })
            M.show_progress("Sent selection to Antigravity via UDP...")
            vim.defer_fn(M.close_progress, 1500)
        end)
    end, { desc = "Ask Antigravity Agent about selection", noremap = true, silent = true })

    vim.keymap.set('n', '<leader>aa', M.accept_diff, { desc = "Accept Agent Diff", noremap = true, silent = true })
    vim.keymap.set('n', '<leader>ar', M.reject_diff, { desc = "Reject Agent Diff", noremap = true, silent = true })
end

return M