#!/usr/libexec/flua

-- SPDX-License-Identifier: BSD-2-Clause
--
-- Copyright(c) 2025 The FreeBSD Foundation.
--
-- This software was developed by Isaac Freund <ifreund@freebsdfoundation.org>
-- under sponsorship from the FreeBSD Foundation.

local sys_wait = require("posix.sys.wait")
local unistd = require("posix.unistd")

local all_libcompats <const> = "%%_ALL_libcompats%%"

-- Run a command using the OS shell and capture the stdout
-- Strips exactly one trailing newline if present, does not strip any other whitespace.
-- Asserts that the command exits cleanly
local function capture(command)
	local p = io.popen(command)
	local output = p:read("*a")
	assert(p:close())
	-- Strip exactly one trailing newline from the output, if there is one
	return output:match("(.-)\n$") or output
end

local function append_list(list, other)
	for _, item in ipairs(other) do
		table.insert(list, item)
	end
end

-- Read from the given fd until EOF
-- Returns all the data read as a single string
local function read_all(fd)
	local ret = ""
	repeat
		local buffer = assert(unistd.read(fd, 1024))
		ret = ret .. buffer
	until buffer == ""
	return ret
end

-- Run bsddialog with the given argument list
-- Returns the exit code and stderr output of bsddialog
local function bsddialog(args)
	local r, w = assert(unistd.pipe())

	local pid = assert(unistd.fork())
	if pid == 0 then
		assert(unistd.close(r))
		assert(unistd.dup2(w, 2))
		assert(unistd.execp("bsddialog", args))
		unistd._exit()
	end
	assert(unistd.close(w))

	local output = read_all(r)
	assert(unistd.close(r))

	local _, _, exit_code = assert(sys_wait.wait(pid))
	return exit_code, output
end

-- Prompts the user for a yes/no answer to the given question using bsddialog
-- Returns true if the user answers yes and false if the user answers no.
local function prompt_yn(question)
	local exit_code = bsddialog({
		"--yesno",
		"--disable-esc",
		question,
		0, 0, -- autosize
	})
	return exit_code == 0
end

-- Creates a dialog for component selection mirroring the
-- traditional tarball component selection dialog.
local function select_components(components, options)
	local descriptions = {
		kernel_dbg = "Kernel debug info",
		base_dbg = "Base system debug info",
		src = "System source tree",
		tests = "Test suite",
		lib32 = "32-bit compatibility libraries",
		lib32_dbg = "32-bit compatibility libraries debug info",
	}
	local defaults = {
		kernel_dbg = "on",
		base_dbg = "off",
		src = "off",
		tests = "off",
		lib32 = "on",
		lib32_dbg = "off",
	}

	-- Sorting the components is necessary to ensure that the ordering is
	-- consistent in the UI.
	local sorted_components = {}
	for component, _ in pairs(components) do
		table.insert(sorted_components, component)
	end
	table.sort(sorted_components)

	local checklist_items = {}
	for _, component in ipairs(sorted_components) do
		if component ~= "base" and component ~= "kernel" and
		    not (component == "kernel_dbg" and options.no_kernel) and
		    #components[component] > 0 then
			local description = descriptions[component] or "''"
			local default = defaults[component]  or "off"
			table.insert(checklist_items, component)
			table.insert(checklist_items, description)
			table.insert(checklist_items, default)
		end
	end

	local bsddialog_args = {
		"--backtitle", "FreeBSD Installer",
		"--title", "Select System Components",
		"--nocancel",
		"--disable-esc",
		"--separate-output",
		"--checklist", "Choose optional system components to install:",
		"0", "0", "0", -- autosize
	}
	append_list(bsddialog_args, checklist_items)

	local exit_code, output = bsddialog(bsddialog_args)
	-- This should only be possible if bsddialog is killed by a signal
	-- or buggy, we disable the cancel option and esc key.
	-- If this does happen, there's not much we can do except exit with a
	-- hopefully useful stack trace.
	assert(exit_code == 0)

	local selected = {"base"}
	if not options.no_kernel then
		table.insert(selected, "kernel")
	end
	for component in output:gmatch("[^\n]+") do
		table.insert(selected, component)
	end

	return selected
end

-- Returns a list of pkgbase packages selected by the user
local function select_packages(pkg, options)
	local components = {
		kernel = {},
		kernel_dbg = {},
		base = {},
		base_dbg = {},
		src = {},
		tests = {},
	}

	for compat in all_libcompats:gmatch("%S+") do
		components["lib" .. compat] = {}
		components["lib" .. compat .. "_dbg"] = {}
	end

	local rquery = capture(pkg .. "rquery -U -r FreeBSD-base %n")
	for package in rquery:gmatch("[^\n]+") do
		if package == "FreeBSD-src" or package:match("^FreeBSD%-src%-.*") then
			table.insert(components["src"], package)
		elseif package == "FreeBSD-tests" or package:match("^FreeBSD%-tests%-.*") then
			table.insert(components["tests"], package)
		elseif package:match("^FreeBSD%-kernel%-.*") then
			-- Kernels other than FreeBSD-kernel-generic are ignored
			if package == "FreeBSD-kernel-generic" then
				table.insert(components["kernel"], package)
			elseif package == "FreeBSD-kernel-generic-dbg" then
				table.insert(components["kernel_dbg"], package)
			end
		elseif package:match(".*%-dbg$") then
			table.insert(components["base_dbg"], package)
		else
			local found = false
			for compat in all_libcompats:gmatch("%S+") do
				if package:match(".*%-dbg%-lib" .. compat .. "$") then
					table.insert(components["lib" .. compat .. "_dbg"], package)
					found = true
					break
				elseif package:match(".*%-lib" .. compat .. "$") then
					table.insert(components["lib" .. compat], package)
					found = true
					break
				end
			end
			if not found then
				table.insert(components["base"], package)
			end
		end
	end
	-- Don't assert the existence of dbg, tests, and src packages here. If using
	-- a custom local repository with BSDINSTALL_PKG_REPOS_DIR we shouldn't
	-- require it to have all packages.
	assert(#components["kernel"] == 1)
	assert(#components["base"] > 0)

	local selected = {}
	for _, component in ipairs(select_components(components, options)) do
		append_list(selected, components[component])
	end

	return selected
end

local function parse_options()
	local options = {}
	for _, a in ipairs(arg) do
		if a == "--no-kernel" then
			options.no_kernel = true
		else
			io.stderr:write("Error: unknown option " .. a .. "\n")
			os.exit(1)
		end
	end
	return options
end

-- Fetch and install pkgbase packages to BSDINSTALL_CHROOT.
-- Respect BSDINSTALL_PKG_REPOS_DIR if set, otherwise use pkg.freebsd.org.
local function pkgbase()
	local options = parse_options()

	-- TODO Support fully offline pkgbase installation by taking a new enough
	-- version of pkg.pkg as input.
	if not os.execute("pkg -N > /dev/null 2>&1") then
		print("Bootstrapping pkg on the host system")
		assert(os.execute("pkg bootstrap -y"))
	end

	local chroot = assert(os.getenv("BSDINSTALL_CHROOT"))
	assert(os.execute("mkdir -p " .. chroot))

	local repos_dir = os.getenv("BSDINSTALL_PKG_REPOS_DIR")
	if not repos_dir then
		repos_dir = chroot .. "/usr/local/etc/pkg/repos/"
		assert(os.execute("mkdir -p " .. repos_dir))
		assert(os.execute("cp /usr/share/bsdinstall/FreeBSD-base.conf " .. repos_dir))

		-- Since pkg always interprets fingerprints paths as relative to
		-- the --rootdir we must copy the key from the host.
		assert(os.execute("mkdir -p " .. chroot .. "/usr/share/keys"))
		assert(os.execute("cp -R /usr/share/keys/pkg " .. chroot .. "/usr/share/keys/"))
	end

	-- We must use --repo-conf-dir rather than -o REPOS_DIR here as the latter
	-- is interpreted relative to the --rootdir. BSDINSTALL_PKG_REPOS_DIR must
	-- be allowed to point to a path outside the chroot.
	local pkg = "pkg --rootdir " .. chroot ..
		" --repo-conf-dir " .. repos_dir .. " -o IGNORE_OSVERSION=yes "

	while not os.execute(pkg .. "update") do
		if not prompt_yn("Updating repositories failed, try again?") then
			os.exit(1)
		end
	end

	local packages = table.concat(select_packages(pkg, options), " ")

	while not os.execute(pkg .. "install -U -F -y -r FreeBSD-base " .. packages) do
		if not prompt_yn("Fetching packages failed, try again?") then
			os.exit(1)
		end
	end

	if not os.execute(pkg .. "install -U -y -r FreeBSD-base " .. packages) then
		os.exit(1)
	end
end

pkgbase()