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 = {}

---@alias ow.lsp.codelens.Position "eol" | "right_align" | "inline" | "above"

---@type ow.lsp.codelens.Position
local position = "above"
local separator = " | "

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

local refresh_debounced = util.debounce(function(buf)
    local state = state_by_buf[buf]
    if state then
        state:refresh()
    end
end, REFRESH_DEBOUNCE_MS)

---@class ow.lsp.codelens.Row
---@field row integer
---@field ready lsp.CodeLens[]
---@field pending integer
local Row = {}
Row.__index = Row

---@param row integer
---@return ow.lsp.codelens.Row
function Row.new(row)
    return setmetatable({ row = row, ready = {}, pending = 0 }, Row)
end

---@param lens lsp.CodeLens
function Row:add(lens)
    table.insert(self.ready, lens)
end

function Row:await()
    self.pending = self.pending + 1
end

---@param resolved? lsp.CodeLens
function Row:resolved(resolved)
    self.pending = self.pending - 1
    if resolved then
        table.insert(self.ready, resolved)
    end
end

---@param buf integer
function Row:render(buf)
    if not vim.api.nvim_buf_is_valid(buf) then
        return
    end
    if self.pending > 0 or #self.ready == 0 then
        return
    end

    table.sort(self.ready, function(a, b)
        return a.range.start.character < b.range.start.character
    end)
    local parts = {}
    for i, lens in ipairs(self.ready) do
        table.insert(parts, { lens.command.title, "LspCodeLens" })
        if i < #self.ready then
            table.insert(parts, { separator, "LspCodeLensSeparator" })
        end
    end

    local col, opts
    if position == "above" then
        local line = vim.api.nvim_buf_get_lines(
            buf,
            self.row,
            self.row + 1,
            false
        )[1] or ""
        local indent = line:match("^%s*") or ""
        local chunks = parts
        if indent ~= "" then
            chunks = { { indent, "LspCodeLensSeparator" } }
            vim.list_extend(chunks, parts)
        end
        col = 0
        opts = {
            virt_lines = { chunks },
            virt_lines_above = true,
            hl_mode = "combine",
        }
    else
        col = position == "inline" and self.ready[1].range.start.character or 0
        opts = {
            virt_text = parts,
            virt_text_pos = position,
            hl_mode = "combine",
        }
    end

    -- 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,
        { self.row, 0 },
        { self.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 then
        local d = primary.details --[[@as vim.api.keyset.extmark_details]]
        local same = opts.virt_lines
                and vim.deep_equal(d.virt_lines, opts.virt_lines)
                and d.virt_lines_above == opts.virt_lines_above
            or not opts.virt_lines
                and d.virt_text_pos == opts.virt_text_pos
                and vim.deep_equal(d.virt_text, opts.virt_text)
        if same then
            return
        end
    end

    opts.id = primary and primary.id or nil
    vim.api.nvim_buf_set_extmark(buf, NS, self.row, col, opts)
end

---@class ow.lsp.codelens.State
---@field buf integer
---@field enabled boolean
---@field attached boolean
---@field generation integer
---@field rows table<integer, ow.lsp.codelens.Row>
local State = {}
State.__index = State

---@param buf integer
---@return ow.lsp.codelens.State
function State.new(buf)
    return setmetatable({
        buf = buf,
        enabled = false,
        attached = false,
        generation = 0,
        rows = {},
    }, State)
end

function State:render()
    if not vim.api.nvim_buf_is_valid(self.buf) then
        return
    end
    if not self.enabled then
        vim.api.nvim_buf_clear_namespace(self.buf, NS, 0, -1)
        return
    end

    for _, row in pairs(self.rows) do
        row:render(self.buf)
    end

    for _, mark in
        ipairs(vim.api.nvim_buf_get_extmarks(self.buf, NS, 0, -1, {}))
    do
        if not self.rows[mark[2]] then
            vim.api.nvim_buf_del_extmark(self.buf, NS, mark[1])
        end
    end
end

function State:refresh()
    if not self.enabled then
        return
    end

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

    local function collect(new_rows, client, lens)
        local row = new_rows[lens.range.start.line]
        if not row then
            row = Row.new(lens.range.start.line)
            new_rows[lens.range.start.line] = row
        end
        if lens.command then
            row:add(lens)
            return
        end
        if
            not client
            or not client:supports_method(
                vim.lsp.protocol.Methods.codeLens_resolve
            )
        then
            return
        end
        row:await()
        client:request(
            vim.lsp.protocol.Methods.codeLens_resolve,
            lens,
            function(_, resolved)
                if self.generation ~= gen then
                    return
                end
                row:resolved(type(resolved) == "table" and resolved or nil)
                row:render(self.buf)
            end,
            self.buf
        )
    end

    vim.lsp.buf_request_all(
        self.buf,
        vim.lsp.protocol.Methods.textDocument_codeLens,
        params,
        function(results)
            if self.generation ~= gen then
                return
            end
            local new_rows = {}
            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
                        collect(new_rows, client, lens)
                    end
                end
            end
            self.rows = new_rows
            self:render()
        end
    )
end

function State:attach()
    if self.attached then
        return
    end
    self.attached = true
    vim.api.nvim_buf_attach(self.buf, false, {
        on_lines = function()
            if not self.enabled then
                self.attached = false
                return true
            end
            refresh_debounced:call(self.buf)
        end,
        on_reload = function()
            if self.enabled then
                refresh_debounced:call(self.buf)
            end
        end,
        on_detach = function()
            self.attached = false
        end,
    })
end

---@param value boolean
function State:enable(value)
    self.enabled = value
    if value then
        self:attach()
        self:refresh()
    else
        refresh_debounced:cancel(self.buf)
        self:render()
    end
end

function State:toggle()
    self:enable(not self.enabled)
end

---@param buf integer
---@return ow.lsp.codelens.State
local function get_state(buf)
    state_by_buf[buf] = state_by_buf[buf] or State.new(buf)
    return state_by_buf[buf]
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)
    if value == nil then
        value = true
    end
    buf = buf or vim.api.nvim_get_current_buf()
    get_state(buf):enable(value)
end

---@param buf? integer
function M.toggle(buf)
    buf = buf or vim.api.nvim_get_current_buf()
    get_state(buf):toggle()
end

---@class ow.lsp.codelens.SetupOpts
---@field position? ow.lsp.codelens.Position
---@field separator? string

---@param opts? ow.lsp.codelens.SetupOpts
function M.setup(opts)
    opts = opts or {}
    if opts.position ~= nil then
        position = opts.position
    end
    if opts.separator ~= nil then
        separator = opts.separator
    end

    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
                state:refresh()
            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