local util = require("util")

local GROUP = vim.api.nvim_create_augroup("ow.lsp.codelens", { clear = true })
local NS = vim.api.nvim_create_namespace("ow.lsp.codelens")
local REFRESH_DEBOUNCE_MS = 200

local M = {}

---@type "eol" | "right_align" | "inline"
M.virt_text_pos = "eol"
M.separator = " | "

---@class ow.lsp.codelens.Row
---@field ready lsp.CodeLens[]  resolved lenses with `command`, ready to paint
---@field pending integer       count of unresolved lenses still in flight

---@class ow.lsp.codelens.State
---@field enabled boolean
---@field attached boolean
---@field generation integer
---@field rows table<integer, ow.lsp.codelens.Row>

---@type table<integer, ow.lsp.codelens.State>
local state_by_buf = {}

---@param buf integer
---@return ow.lsp.codelens.State
local function get_state(buf)
    state_by_buf[buf] = state_by_buf[buf]
        or {
            enabled = false,
            attached = false,
            generation = 0,
            rows = {},
        }
    return state_by_buf[buf]
end

---@param buf integer
---@param row integer
local function render_row(buf, row)
    if not vim.api.nvim_buf_is_valid(buf) then
        return
    end
    local state = state_by_buf[buf]
    if not state or not state.enabled then
        return
    end

    local entry = state.rows[row]
    -- Preserve existing extmark while any lens on the row is still resolving
    -- so the user doesn't see a partial row (e.g. "refs" without "impls").
    if not entry or entry.pending > 0 or #entry.ready == 0 then
        return
    end

    table.sort(entry.ready, function(a, b)
        return a.range.start.character < b.range.start.character
    end)
    local parts = {}
    for i, lens in ipairs(entry.ready) do
        table.insert(parts, { lens.command.title, "LspCodeLens" })
        if i < #entry.ready then
            table.insert(parts, { M.separator, "LspCodeLensSeparator" })
        end
    end
    local col = M.virt_text_pos == "inline"
            and entry.ready[1].range.start.character
        or 0

    -- One extmark per row. Extmarks auto-shift with edits, so multiple from
    -- prior refreshes can end up on the same row; pick the first and drop
    -- the rest, then reuse its id for an atomic update.
    local marks = vim.api.nvim_buf_get_extmarks(
        buf,
        NS,
        { row, 0 },
        { row, -1 },
        { details = true }
    )
    local primary
    for i, mark in ipairs(marks) do
        if i == 1 then
            primary = { id = mark[1], col = mark[3], details = mark[4] }
        else
            vim.api.nvim_buf_del_extmark(buf, NS, mark[1])
        end
    end

    if
        primary
        and primary.col == col
        and primary.details
        and primary.details.virt_text_pos == M.virt_text_pos
        and vim.deep_equal(primary.details.virt_text, parts)
    then
        return
    end

    vim.api.nvim_buf_set_extmark(buf, NS, row, col, {
        id = primary and primary.id or nil,
        virt_text = parts,
        virt_text_pos = M.virt_text_pos,
        hl_mode = "combine",
    })
end

---@param buf integer
local function render_all(buf)
    if not vim.api.nvim_buf_is_valid(buf) then
        return
    end
    local state = state_by_buf[buf]
    if not state or not state.enabled then
        vim.api.nvim_buf_clear_namespace(buf, NS, 0, -1)
        return
    end

    for row in pairs(state.rows) do
        render_row(buf, row)
    end

    -- Drop extmarks on rows that are no longer in state (lens removed).
    for _, mark in ipairs(vim.api.nvim_buf_get_extmarks(buf, NS, 0, -1, {})) do
        if not state.rows[mark[2]] then
            vim.api.nvim_buf_del_extmark(buf, NS, mark[1])
        end
    end
end

---@param buf integer
local function refresh(buf)
    local state = get_state(buf)
    if not state.enabled then
        return
    end

    state.generation = state.generation + 1
    local gen = state.generation
    local params =
        { textDocument = vim.lsp.util.make_text_document_params(buf) }
    local new_rows = {}

    vim.lsp.buf_request_all(
        buf,
        vim.lsp.protocol.Methods.textDocument_codeLens,
        params,
        function(results)
            if state.generation ~= gen then
                return
            end
            for client_id, response in pairs(results) do
                if not response.err and type(response.result) == "table" then
                    local client = vim.lsp.get_client_by_id(client_id)
                    for _, lens in ipairs(response.result) do
                        local row = lens.range.start.line
                        new_rows[row] = new_rows[row]
                            or { ready = {}, pending = 0 }
                        if lens.command then
                            table.insert(new_rows[row].ready, lens)
                        elseif
                            client
                            and client:supports_method(
                                vim.lsp.protocol.Methods.codeLens_resolve
                            )
                        then
                            new_rows[row].pending = new_rows[row].pending + 1
                            client:request(
                                vim.lsp.protocol.Methods.codeLens_resolve,
                                lens,
                                function(_, resolved)
                                    if state.generation ~= gen then
                                        return
                                    end
                                    local entry = new_rows[row]
                                    entry.pending = entry.pending - 1
                                    if type(resolved) == "table" then
                                        table.insert(entry.ready, resolved)
                                    end
                                    render_row(buf, row)
                                end,
                                buf
                            )
                        end
                    end
                end
            end
            -- Swap in the new row map and paint everything once. Rows with
            -- pending resolves keep their prior extmark in place; each
            -- resolve callback above repaints its own row when it arrives.
            state.rows = new_rows
            render_all(buf)
        end
    )
end

local refresh_debounced = util.debounce(refresh, REFRESH_DEBOUNCE_MS)

---@param buf integer
local function attach_buf(buf)
    local state = get_state(buf)
    if state.attached then
        return
    end
    state.attached = true
    vim.api.nvim_buf_attach(buf, false, {
        on_lines = function(_, b)
            local s = state_by_buf[b]
            if not s or not s.enabled then
                if s then
                    s.attached = false
                end
                return true
            end
            refresh_debounced:call(b)
        end,
        on_reload = function(_, b)
            local s = state_by_buf[b]
            if s and s.enabled then
                refresh_debounced:call(b)
            end
        end,
        on_detach = function(_, b)
            local s = state_by_buf[b]
            if s then
                s.attached = false
            end
        end,
    })
end

---@param buf? integer
---@return boolean
function M.is_enabled(buf)
    buf = buf or vim.api.nvim_get_current_buf()
    local state = state_by_buf[buf]
    return state ~= nil and state.enabled
end

---@param value? boolean
---@param buf? integer
function M.enable(value, buf)
    buf = buf or vim.api.nvim_get_current_buf()
    if value == nil then
        value = true
    end
    local state = get_state(buf)
    state.enabled = value
    if value then
        attach_buf(buf)
        refresh(buf)
    else
        refresh_debounced:cancel(buf)
        render_all(buf)
    end
end

---@param buf? integer
function M.toggle(buf)
    M.enable(not M.is_enabled(buf), buf)
end

function M.setup()
    vim.api.nvim_create_autocmd({ "BufEnter", "LspAttach" }, {
        group = GROUP,
        callback = function(ev)
            local state = state_by_buf[ev.buf]
            if state and state.enabled then
                refresh(ev.buf)
            end
        end,
    })

    vim.api.nvim_create_autocmd({ "BufDelete", "BufWipeout" }, {
        group = GROUP,
        callback = function(ev)
            refresh_debounced:cancel(ev.buf)
            state_by_buf[ev.buf] = nil
        end,
    })
end

return M