local Row = require("lsp.codelens.row")
local log = require("log")

---@class ow.lsp.codelens.RefreshTask
---@field session ow.lsp.codelens.Session
---@field aborted boolean
---@field private cancels fun()[]
local RefreshTask = {}
RefreshTask.__index = RefreshTask

---@param session ow.lsp.codelens.Session
---@return ow.lsp.codelens.RefreshTask
function RefreshTask.new(session)
    return setmetatable({
        session = session,
        aborted = false,
        cancels = {},
    }, RefreshTask)
end

---@param cancel fun()
function RefreshTask:track(cancel)
    table.insert(self.cancels, cancel)
end

function RefreshTask:abort()
    self.aborted = true
    for _, cancel in ipairs(self.cancels) do
        pcall(cancel)
    end
    self.cancels = {}
end

---@param client vim.lsp.Client
---@param lens lsp.CodeLens
---@param row ow.lsp.codelens.Row
function RefreshTask:resolve(client, lens, row)
    local _, req_id = client:request(
        "codeLens/resolve",
        lens,
        ---@param err lsp.ResponseError?
        ---@param resolved lsp.CodeLens?
        ---@param ctx lsp.HandlerContext
        function(err, resolved, ctx)
            if self.aborted then
                return
            end
            if err then
                log.warning(
                    "client %d: %s failed: %s",
                    client.id,
                    ctx.method,
                    err.message
                )
            end
            row:resolve(resolved)
            row:render(self.session.buf)
        end,
        self.session.buf
    )
    if req_id then
        self:track(function()
            client:cancel_request(req_id)
        end)
    end
end

---@param rows table<integer, ow.lsp.codelens.Row>
---@param client vim.lsp.Client?
---@param lens lsp.CodeLens
function RefreshTask:process_lens(rows, client, lens)
    local row = rows[lens.range.start.line]
    if not row then
        row = Row.new(lens.range.start.line)
        rows[lens.range.start.line] = row
    end
    if lens.command then
        row:add(lens)
        return
    end
    if not client or not client:supports_method("codeLens/resolve") then
        return
    end
    row:expect()
    self:resolve(client, lens, row)
end

---@param results table<integer, ow.lsp.CodeLensResponse>
function RefreshTask:process_results(results)
    local session = self.session
    local new_rows = {}
    for client_id, result in pairs(results) do
        if result.err then
            log.warning(
                "client %d: %s failed: %s",
                client_id,
                result.context.method,
                result.err.message
            )
        end
        if not result.err and type(result.result) == "table" then
            local client = vim.lsp.get_client_by_id(client_id)
            for _, lens in ipairs(result.result) do
                self:process_lens(new_rows, client, lens)
            end
        end
    end
    session.rows = new_rows
    session:render()
end

function RefreshTask:run()
    local params = {
        textDocument = vim.lsp.util.make_text_document_params(self.session.buf),
    }
    local cancel = vim.lsp.buf_request_all(
        self.session.buf,
        "textDocument/codeLens",
        params,
        function(results)
            if self.aborted then
                return
            end
            self:process_results(results)
        end
    )
    self:track(cancel)
end

return RefreshTask