/*
 * This source code is released into the public domain.
 */

/*
 * Generic (non-jail-related) ZFS handling.
 *
 * It would be nice to use libzfs_core for this, but it still has some issues
 * on FreeBSD, e.g. https://bugs.freebsd.org/bugzilla/show_bug.cgi?id=257222.
 */

#include <unordered_map>

#include <unistd.h>

#include "ctype.hh"
#include "spawn.hh"
#include "words.hh"
#include "zfs.hh"

using namespace std::literals;

namespace lfjail::zfs {

namespace {

// Hardcode the program locations so we don't rely on $PATH.
constexpr std::string path_zpool = "/sbin/zpool";
constexpr std::string path_zfs = "/sbin/zfs";

/*
 * Run a command and split each line of output into whitespace-separated
 * columns.
 */
auto columns(exec::executor auto &&executor) 
	-> std::generator<std::vector<std::string_view>>
{
	std::string output;

	auto result = spawn(std::move(executor),
			    capture(STDOUT_FILENO, output)).wait();

	if (!result)
		throw zfs_error("failed to run zfs command");

	for (auto &&line : words(output, '\n')) {
		co_yield std::vector(std::from_range, words(line));
	}
}

// Run a command which produces two columns of output, assumed to be
// a name and a value; call the function from the provided map on
// each line.
auto set_values(auto &&map, std::string command, auto &&...args) -> void
{
	auto lines = columns(
		exec::execl(std::move(command), std::move(args)...));

	for (auto &&words : lines) {
		if (std::ranges::size(words) != 2)
			throw zfs_error("unexpected output from zfs command");

		auto property = words[0];
		auto value = words[1];

		auto handler = map.find(property);
		if (handler != map.end())
			handler->second(value);
	}
}

} // anonymous namespace

auto get_pool(std::string_view name) -> pool {
	auto ret = pool{};

	auto handlers = std::unordered_map{
		std::pair{"name"sv, [&ret](auto &&v) { ret.name = v; }},
	};

	set_values(handlers, path_zpool,
		   "zpool", "get", "-H", "-oproperty,value", "all", name);

	return ret;
}

auto get_pools() -> std::generator<pool> {
	auto lines = columns(exec::execl(path_zpool,
					 "zpool", "get", "-H", "-o", "name"));

	for (auto &&words : lines)
		co_yield get_pool(words[0]);
}

auto get_dataset(std::string_view name) -> dataset {
	auto ret = dataset{};

	auto handlers = std::unordered_map{
		std::pair{"name"sv, [&ret](auto &&v) { ret.name = v; }},
	};

	set_values(handlers, path_zfs, "zfs", "get", "-H",
		   "-oproperty,value"sv, "all", name);

	return ret;
}

auto get_dataset_names() -> std::generator<std::string> {
	auto lines = columns(exec::execl(path_zfs,
					 "zfs", "get", "-H", "-o", "name"));

	for (auto &&words : lines)
		co_yield std::string(words[0]);
}

auto get_datasets() -> std::generator<dataset> {
	for (auto &&dataset : get_dataset_names())
		co_yield get_dataset(dataset);
}

auto dataset_exists(std::string name) -> bool {
	return std::ranges::contains(get_dataset_names(), name);
}

} // namespace lfjail::zfs