local table = require 'std.table'
local icon = require 'std.icon'

-- display inlay hints at the END of the line, instead of
-- mixed with the source code
--
-- stolen from: https://github.com/chrisgrieser/nvim-lsp-endhints/blob/main/lua/lsp-endhints/override-handler.lua

local config = {
	icons = {
		type = icon.type,
		parameter = icon.parameter,
		offspec = icon.offspec,
		unknown = icon.unknown,
	},
	label = {
		truncateAtChars = 20,
		padding = 1,
		marginLeft = 0,
		sameKindSeparator = ', ',
	},
	extmark = {
		priority = 50,
	},
}

local ns = vim.api.nvim_create_namespace  'bugabinga-inlay'

local function changedRefreshHandler( err, result, ctx, _ )
	-- GUARD
	if not vim.lsp.inlay_hint.is_enabled() then return end
	local bufnr = ctx.bufnr or -1
	if not vim.api.nvim_buf_is_valid( bufnr ) then return end
	if not result then return end
	local client = vim.lsp.get_client_by_id( ctx.client_id )
	if not client then return end
	local inlayHintProvider = client.server_capabilities.inlayHintProvider
	if not inlayHintProvider then return end
	if err then
		vim.notify( vim.inspect( err ), vim.log.levels.ERROR )
		return
	end

	local padding = (' '):rep( config.label.padding )
	local marginLeft = (' '):rep( config.label.marginLeft )

	vim.api.nvim_buf_clear_namespace( bufnr, ns, 0, -1 )

	-- Collect all hints for each line, so we can sort them by column in the loop
	-- below. This ensures that the hints are displayed in the correct order.
	local hintLines = vim.iter( result ):fold( {}, function ( acc, hint )
		local lnum = hint.position.line
		local col = hint.position.character

		-- LABEL
		local label = hint.label
		if type( label ) ~= 'string' then
			-- label can be string or `InlayHintLabelPart[]` https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#inlayHint
			label = vim.iter( hint.label ):map( function ( labelPart ) return labelPart.value end ):join  ''
		end
		label = vim.trim( label:gsub( '^:', '' ):gsub( ':$', '' ) )
		if config.label.truncateAtChars and #label > config.label.truncateAtChars then
			label = label:sub( 1, config.label.truncateAtChars ) .. '…'
		end

		-- KIND
		-- 1: type, 2: parameter -- https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#inlayHintKind
		local kind
		if hint.kind == 1 then
			kind = 'type'
		elseif hint.kind == 2 then
			kind = 'parameter'
		elseif hint.kind == nil then
			kind = 'unknown'
		else
			kind = 'offspec'
		end

		if not acc[lnum] then acc[lnum] = {} end
		table.insert( acc[lnum], { label = label, col = col, kind = kind } )
		return acc
	end )

	-- loop: add hints as extmarks for each line
	for lnum, hints in pairs( hintLines ) do
		table.sort( hints, function ( a, b ) return a.col < b.col end )

		-- merge hints of same kind
		-- add icon only when parameter kind is different from the previous one
		local hintsMerged = ''
		for i = 1, #hints do
			local hint = hints[i]
			local lastKind = hints[i - 1] and hints[i - 1].kind
			if lastKind == hint.kind then
				hintsMerged = hintsMerged .. config.label.sameKindSeparator .. hint.label
			else
				local hint_icon = config.icons[hint.kind]
				local pad = i ~= 1 and ' ' or ''
				hintsMerged = hintsMerged .. pad .. hint_icon .. hint.label
			end
		end

		-- add padding & margin
		local virtText = {
			{ padding .. hintsMerged .. padding, 'LspInlayHint' },
		}
		if marginLeft ~= '' then table.insert( virtText, 1, { marginLeft } ) end

		-- add extmark for the line
		vim.api.nvim_buf_set_extmark( bufnr, ns, lnum, 0, {
			virt_text = virtText,
			virt_text_pos = 'eol',
			hl_mode = 'combine', -- ensures highlights are combined with the background, see #2
			strict = false,   -- prevents error on quick buffer changes (e.g. quick series of undos)
			priority = config.extmark.priority,
		} )
	end
end

vim.lsp.handlers['textDocument/inlayHint'] = changedRefreshHandler