local M = {
	results = {}
}

local ns = vim.api.nvim_create_namespace("live-tests")

local test_function_query_string = [[
(
 (function_declaration
  name: (identifier) @name
  parameters:
    (parameter_list
     (parameter_declaration
      name: (identifier)
      type: (pointer_type
          (qualified_type
           package: (package_identifier) @_package_name
           name: (type_identifier) @_type_name)))))

 (#eq? @_package_name "testing")
 (#eq? @_type_name "T")
 (#eq? @name "%s")
)
]]

local find_test_line = function(go_bufnr, name)
	local formatted = string.format(test_function_query_string, name)
	local query = vim.treesitter.query.parse("go", formatted)
	local parser = vim.treesitter.get_parser(go_bufnr, "go", {})
	local tree = parser:parse()[1]
	local root = tree:root()

	for id, node in query:iter_captures(root, go_bufnr, 0, -1) do
		if id == 1 then
			local range = { node:range() }
			return range[1]
		end
	end
end

local function clear(table)
	for k in pairs(table) do
		table[k] = nil
	end
end

vim.api.nvim_create_user_command('GoTest', function()
	M.goRunTests()
end, {})

vim.api.nvim_create_user_command('GoPrintResults', function()
	print(vim.inspect(M.results))
end, {})

vim.api.nvim_create_user_command('GoSetTestMarks', function()
	M.goSetAllMarks()
end, {})

vim.api.nvim_create_user_command('GoClearResults', function()
	clear(M.results)
end, {})


vim.api.nvim_create_user_command('GoClearTestMarks', function()
	local currbuf = vim.api.nvim_get_current_buf()
	vim.api.nvim_buf_clear_namespace(currbuf, ns, 0, -1)
end, {})

local errored = false
function M.goRunTests()
	-- TODO: replace jobstart/jobwait with system()
	vim.fn.jobstart({ "go", "test", "./...", "-json" }, {
		stdout_buffered = true,
		on_stdout = function(_, data)
			if errored == true then
				errored = false
				return
			end

			if not data then
				return
			end

			clear(M.results)

			local collected_messages = {}
			for _, line in ipairs(data) do
				if line == "" then goto continue end
				local decoded = vim.json.decode(line)

				if decoded.Test == nil then
					goto continue
				end
				if decoded.Action == 'pass' or decoded.Action == 'fail' then
					M.results[decoded.Test] = { Action = decoded.Action }
				end
				if decoded.Action == 'output' then
					local testname = decoded.Test
					local a = ""
					if collected_messages[testname] == nil then
						a = ""
					else
						a = collected_messages[testname]
					end
					collected_messages[decoded.Test] = a .. decoded.Output
				end
				::continue::
			end
			-- Merge collected messages into results
			for test, _ in pairs(M.results) do
				M.results[test].Message = collected_messages[test]
			end
		end,
		on_stderr = function(_, data)
			for _, value in pairs(data) do
				if value and value ~= "" then
					errored = true
					print('Error while running tests: ', value)
				end
			end
		end,
		on_exit = function()
			M.goSetAllMarks()
		end
	})
end

function M.goSetAllMarks()
	local files = vim.fn.system({ "find", "-type", "f", "-name", "*_test.go" })

	local sep = "\n"
	local filetable = {}
	for str in string.gmatch(files, "([^" .. sep .. "]+)") do
		table.insert(filetable, str)
	end

	for _, file in ipairs(filetable) do
		vim.fn.bufadd(file)
		vim.fn.bufload(file)
		local bufno = vim.fn.bufnr(file)
		M.goSetMarks(bufno)
	end
end

function M.goSetMarks(currbuf)
	vim.api.nvim_buf_clear_namespace(currbuf, ns, 0, -1)
	vim.diagnostic.set(ns, currbuf, {}, {})
	local text = {}
	local failed = {}

	for test, val in pairs(M.results) do
		local linenum = find_test_line(currbuf, test)
		if linenum == nil then goto continue end

		if val.Action == 'pass' then
			text = { "✔" }
			vim.api.nvim_buf_set_extmark(currbuf, ns, linenum, 0, {
				virt_text = { text },
			})
		elseif val.Action == 'fail' then
			table.insert(failed, {
				bufnr = currbuf,
				lnum = linenum,
				col = 0,
				severity = vim.diagnostic.severity.ERROR,
				source = "go-test",
				message = "Test Failed\n\n" .. val.Message,
				user_data = {},
			})
		end

		::continue::
	end
	vim.diagnostic.set(ns, currbuf, failed, {})
end

return M