local git = require("neo-tree.git")
local parser = require("neo-tree.git.parser")
local git_cmd = require("neo-tree.git.cmd")
local utils = require("neo-tree.utils")

local diff_base = nil
local status_cache = {}
local orig_find_status

local function find_worktree_root()
	local root = git.find_worktree_info(vim.fn.getcwd())
	if not root then
		local buf = vim.api.nvim_buf_get_name(0)
		if buf ~= "" then
			root = git.find_worktree_info(buf)
		end
	end
	return root
end

local function run_git(args)
	local result = vim.system(args):wait()
	if result.code ~= 0 then
		return nil
	end
	return result.stdout
end

local function compute_status(wr, base)
	local args = git_cmd.with_args({ "-C", wr, "diff", base, "--name-status", "-z" })
	local out = run_git(args)
	if not out then
		vim.notify("DiffBase: invalid ref " .. base, vim.log.levels.WARN)
		return nil
	end
	local status = parser.parse_diff_name_status_output(wr, false, utils.gsplit_plain(out, "\000"))
	local ut = run_git(git_cmd.with_args({ "-C", wr, "ls-files", "--others", "--exclude-standard", "-z" }))
	if ut then
		for _, p in ipairs(parser.parse_ls_files_output(wr, utils.gsplit_plain(ut, "\000"))) do
			status[p] = "?"
		end
	end
	return status
end

local function patch()
	if not orig_find_status then
		orig_find_status = git.find_existing_status_code
	end
	git.find_existing_status_code = function(path, base_lookup)
		if diff_base then
			local wr, _ = git.find_existing_worktree(path)
			local cached = wr and status_cache[wr]
			local hit = cached and cached[path]
			if hit then
				return hit, true
			end
		end
		return orig_find_status(path, base_lookup)
	end
end

local function set(base)
	base = base or ""
	if base == "" or base == "HEAD" then
		diff_base = nil
		status_cache = {}
		pcall(function()
			require("gitsigns").change_base(nil, true)
		end)
		local e = require("neo-tree.events")
		vim.schedule(function()
			e.fire_event(e.GIT_EVENT, {})
		end)
		vim.notify("DiffBase: reset to HEAD")
		return
	end
	local wr = find_worktree_root()
	if not wr then
		vim.notify("DiffBase: not in a git repo", vim.log.levels.WARN)
		return
	end
	local s = compute_status(wr, base)
	if not s then
		return
	end
	diff_base = base
	status_cache[wr] = s
	pcall(function()
		require("gitsigns").change_base(base, true)
	end)
	patch()
	local events = require("neo-tree.events")
	if git.worktrees[wr] then
		git.worktrees[wr].status = s
	end
	vim.schedule(function()
		events.fire_event(events.GIT_STATUS_CHANGED, { git_root = wr, git_status = s, base = base })
	end)
	vim.notify("DiffBase: set to " .. base)
end

vim.api.nvim_create_user_command("DiffBase", function(opts)
	set(opts.args)
end, {
	nargs = "?",
	complete = function(l)
		return vim.tbl_filter(function(s)
			return s:find(l, 1, true) == 1
		end, { "HEAD", "HEAD~1", "HEAD~2", "HEAD~3", "main", "master", "origin/main", "origin/master" })
	end,
})