/*
 * 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 <functional>
#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 {

/*
 * 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>>
{
	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);
	}
}

template<typename T>
struct set_simple;

template<typename T>
set_simple(T *) -> set_simple<T>;

template<>
struct set_simple<bool> {
	set_simple(bool *valuep_) : valuep(valuep_) {}

	bool *valuep;

	auto operator() (std::string_view newvalue) -> void {
		if (newvalue == "on")
			*valuep = true;
		else if (newvalue == "off")
			*valuep = false;
		else
			throw zfs_error("unexpected boolean value: {}",
					newvalue);
	}
};

template<>
struct set_simple<std::string> {
	set_simple(std::string *valuep_) : valuep(valuep_) {}

	std::string *valuep;

	auto operator() (std::string_view newvalue) -> void {
		*valuep = newvalue;
	}
};

template<>
struct set_simple<std::optional<std::string>> {
	set_simple(std::optional<std::string> *valuep_) : valuep(valuep_) {}

	std::optional<std::string> *valuep;

	auto operator() (std::string_view newvalue) -> void {
		if (newvalue != "-")
			*valuep = newvalue;
	}
};


template<>
struct set_simple<std::uint64_t> {
	set_simple(std::uint64_t *valuep_) : valuep(valuep_) {}

	std::uint64_t *valuep;

	auto operator() (std::string_view newvalue) -> void {
		// XXX
		std::string cvalue(newvalue);
		*valuep = strtoull(cvalue.c_str(), NULL, 10);
	}
};

template<>
struct set_simple<unsigned> {
	set_simple(unsigned *valuep_) : valuep(valuep_) {}

	unsigned *valuep;

	auto operator() (std::string_view newvalue) -> void {
		// XXX
		std::string cvalue(newvalue);
		*valuep = std::strtoul(cvalue.c_str(), NULL, 10);
	}
};

template<>
struct set_simple<double> {
	set_simple(double *valuep_) : valuep(valuep_) {}

	double *valuep;

	auto operator() (std::string_view newvalue) -> void {
		// XXX
		std::string cvalue(newvalue);
		*valuep = std::atof(cvalue.c_str());
	}
};

template<typename Tag>
struct set_simple<enumeration<Tag>> {
	set_simple(enumeration<Tag> *valuep_) : valuep(valuep_) {}

	enumeration<Tag> *valuep;

	auto operator() (std::string_view newvalue) -> void {
		valuep->value = newvalue;
	}
};

} // anonymous namespace

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

	static const auto handlers = std::unordered_map<std::string_view,
				std::function<void (std::string_view)>>{
		{"guid"sv,		set_simple(&ret.guid)},
		{"load_guid"sv,		set_simple(&ret.load_guid)},
		{"size"sv,		set_simple(&ret.size)},
		{"free"sv,		set_simple(&ret.free)},
		{"allocated"sv,		set_simple(&ret.allocated)},
		{"altroot"sv,		set_simple(&ret.altroot)},
		{"health"sv,		set_simple(&ret.health)},
		{"bootfs"sv,		set_simple(&ret.bootfs)},
		{"delegation"sv,	set_simple(&ret.delegation)},
		{"autoreplace"sv,	set_simple(&ret.autoreplace)},
		{"cachefile"sv,		set_simple(&ret.cachefile)},
		{"failmode"sv,		set_simple(&ret.failmode)},
		{"listsnapshots"sv,	set_simple(&ret.listsnapshots)},
		{"autoexpand"sv,	set_simple(&ret.autoexpand)},
		{"readonly"sv,		set_simple(&ret.readonly)},
		{"ashift"sv,		set_simple(&ret.ashift)},
		{"comment"sv,		set_simple(&ret.comment)},
		{"expandsize"sv,	set_simple(&ret.expandsize)},
		{"freeing"sv,		set_simple(&ret.freeing)},
		{"fragmentation"sv,	set_simple(&ret.fragmentation)},
		{"leaked"sv,		set_simple(&ret.leaked)},
		{"multihost"sv,		set_simple(&ret.multihost)},
		{"checkpoint"sv,	set_simple(&ret.checkpoint)},
		{"autotrim"sv,		set_simple(&ret.autotrim)},
		{"compatibility"sv,	set_simple(&ret.compatibility)},
		{"bcloneused"sv,	set_simple(&ret.bcloneused)},
		{"bclonesaved"sv,	set_simple(&ret.bclonesaved)},
		{"bcloneratio"sv,	set_simple(&ret.bcloneratio)},
		{"dedup_table_size"sv,	set_simple(&ret.dedup_table_size)},
		{"dedupratio"sv,	set_simple(&ret.dedupratio)},
		{"last_scrubbed_txg"sv,	set_simple(&ret.last_scrubbed_txg)},
	};

	set_values(handlers, path_zpool,
		   "zpool", "get", "-Hp", "-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{};
	ret.name = name;

	static const auto handlers = std::unordered_map<std::string_view,
				std::function<void (std::string_view)>>{
		{"type"sv,		set_simple(&ret.type)},
		{"creation"sv,		set_simple(&ret.creation)},
		{"used"sv,		set_simple(&ret.used)},
		{"available"sv,		set_simple(&ret.available)},
		{"referenced"sv,	set_simple(&ret.referenced)},
		{"compressratio"sv,	set_simple(&ret.compressratio)},
		{"mounted"sv,		set_simple(&ret.mounted)},
		{"quota"sv,		set_simple(&ret.quota)},
		{"reservation"sv,	set_simple(&ret.reservation)},
		{"recordsize"sv,	set_simple(&ret.recordsize)},
		{"sharenfs"sv,		set_simple(&ret.sharenfs)},
		{"checksum"sv,		set_simple(&ret.checksum)},
		{"compression"sv,	set_simple(&ret.compression)},
		{"atime"sv,		set_simple(&ret.atime)},
		{"devices"sv,		set_simple(&ret.devices)},
		{"exec"sv,		set_simple(&ret.exec)},
		{"setuid"sv,		set_simple(&ret.setuid)},
		{"readonly"sv,		set_simple(&ret.readonly)},
		{"jailed"sv,		set_simple(&ret.jailed)},
		{"snapdir"sv,		set_simple(&ret.snapdir)},
		{"aclmode"sv,		set_simple(&ret.aclmode)},
		{"aclinherit"sv,	set_simple(&ret.aclinherit)},
		{"createtxg"sv,		set_simple(&ret.createtxg)},
		{"canmount"sv,		set_simple(&ret.canmount)},
		{"xattr"sv,		set_simple(&ret.xattr)},
		{"copies"sv,		set_simple(&ret.copies)},
		{"version"sv,		set_simple(&ret.version)},
		{"utf8only"sv,		set_simple(&ret.utf8only)},
		{"normalization"sv,	set_simple(&ret.normalization)},
		{"vscan"sv,		set_simple(&ret.vscan)},
		{"nbmand"sv,		set_simple(&ret.nbmand)},
		{"sharesmb"sv,		set_simple(&ret.sharesmb)},
		{"refquota"sv,		set_simple(&ret.refquota)},
		{"refreservation"sv,	set_simple(&ret.refreservation)},
		{"guid"sv,		set_simple(&ret.guid)},
		{"primarycache"sv,	set_simple(&ret.primarycache)},
		{"secondarycache"sv,	set_simple(&ret.secondarycache)},
		{"usedbysnapshots"sv,	set_simple(&ret.usedbysnapshots)},
		{"usedbydataset"sv,	set_simple(&ret.usedbydataset)},
		{"usedbychildren"sv,	set_simple(&ret.usedbychildren)},
		{"usedbyrefreservation"sv,
					set_simple(&ret.usedbyrefreservation)},
		{"logbias"sv,		set_simple(&ret.logbias)},
		{"objsetid"sv,		set_simple(&ret.objsetid)},
		{"dedup"sv,		set_simple(&ret.dedup)},
		{"mlslabel"sv,		set_simple(&ret.mlslabel)},
		{"sync"sv,		set_simple(&ret.sync)},
		{"dnodesize"sv,		set_simple(&ret.dnodesize)},
		{"refcompressratio"sv,	set_simple(&ret.refcompressratio)},
	};

	set_values(handlers, path_zfs, "zfs", "get", "-Hp",
		   "-oproperty,value", "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);
}

/*
 * Create a new ZFS dataset.
 */
auto destroy_dataset(std::string const &dsname) -> void {
	using namespace std::literals;

	try {
		auto ret = spawn(exec::execl(path_zfs,
				"zfs", "destroy", "-r", "--", dsname)).wait();
		if (!ret)
			throw zfs_error("failed to destroy dataset");
	} catch (std::exception const &exc) {
		throw zfs_error("failed to destroy dataset: {}", exc.what());
	}
}

} // namespace lfjail::zfs