local run_formatter = function(text)
	vim.notify(text)
	vim.notify("----")
	local result = {}
	require("plenary.job"):new({
		command = "sqlfluff",
		cwd = "/usr/bin",
		args = { "format", "-" },
		writer = text,
		on_stdout = function(_, line)
			table.insert(result, line)
		end,
	}):sync()

	-- add the surrounding r#"..."# back
	for _, line in ipairs(result) do
		vim.notify(line)
	end
	-- table.insert(result, "\"#")
	return result
end

local embedded_sql = vim.treesitter.query.parse(
	"rust",
	[[
		(macro_invocation
			(scoped_identifier
				path: (identifier) @path (#eq? @path "sqlx")
				name: (identifier) @name (#any-of? @name "query" "query_scalar" "query_as"))

			(token_tree
				(raw_string_literal) @sql)
			(#offset! @sql 0 3 0 -2)
		)
	]]
)

local get_root = function(bufnr)
	local parser = vim.treesitter.get_parser(bufnr, "rust", {})
	local tree = parser:parse()[1]
	return tree:root()
end

local format = function(bufnr)
	bufnr = bufnr or vim.api.nvim_get_current_buf()

	if vim.api.nvim_get_option_value("filetype", { buf = bufnr }) ~= "rust" then
		vim.notify("SQL format an only be used in Rust")
		return
	end

	local root = get_root(bufnr)

	local changes = {}
	for id, node in embedded_sql:iter_captures(root, bufnr, 0, -1) do
		local name = embedded_sql.captures[id]
		if name == "sql" then
			-- {start row, start col, end row, end col }
			local range = { node:range() }
			local indentation = string.rep(" ", range[2])
			-- run the formatter on the node text
			local formatted = run_formatter(vim.treesitter.get_node_text(node, bufnr))
			-- add indentation
			for idx, line in ipairs(formatted) do
				formatted[idx] = indentation .. line
			end
			-- add changes in reverse order
			table.insert(changes, 1, {
				start = range[1],
				final = range[3],
				formatted = formatted,
			})
		end
	end

	for _, change in ipairs(changes) do
		vim.api.nvim_buf_set_lines(bufnr, change.start, change.final, false, change.formatted)
	end
end

vim.api.nvim_create_user_command("SqlFormat", function()
	format()
end, {})