// This source code is released into the public domain.
export module nihil.cli:command_tree;

import nihil.std;
import nihil.generator;
import nihil.util;
import :command;
import :registry;

namespace nihil {

// command_tree_node represents a node in the command tree.  Each node contains a command, which is
// either a real command which can be invoked, or a stub command which acts as a placeholder in the
// tree. For example, if two commands "add foo" and "add bar" are defined, then "add" will be
// implicitly created as a stub command.
struct command_tree_node final
{
	// Create an empty node with no parent.  This is used as the root of the tree.
	command_tree_node()
		: m_command(std::make_shared<nihil::command>(""))
	{
	}

	// Create a node that contains a command and has a parent.
	command_tree_node(command_tree_node *parent, std::string_view const this_word,
	                  std::shared_ptr<nihil::command> command)
		: m_parent(parent)
		, m_this_word(this_word)
		, m_command(std::move(command))
	{
	}

	// Not copyable.
	command_tree_node(command_tree_node const &) = delete;
	auto operator=(command_tree_node const &) -> command_tree_node & = delete;

	// Movable.
	command_tree_node(command_tree_node &&) = default;
	auto operator=(command_tree_node &&) -> command_tree_node & = default;

	~command_tree_node() = default;

	// Return a child node, or NULL if the child doesn't exist.
	[[nodiscard]] auto get_child(this command_tree_node const &self,
	                             std::string_view const child) -> command_tree_node const *
	{
		if (auto it = self.m_children.find(std::string(child)); it != self.m_children.end())
			return &it->second;

		return nullptr;
	}

	[[nodiscard]] auto
	get_child(this command_tree_node &self, std::string_view const child) -> command_tree_node *
	{
		if (auto it = self.m_children.find(std::string(child)); it != self.m_children.end())
			return &it->second;

		return nullptr;
	}

	// Return a child node if it exists, or insert a new empty node.
	[[nodiscard]] auto get_or_create_child(this command_tree_node &self, std::string_view child)
		-> command_tree_node *
	{
		// Return the existing child, if there is one.
		if (auto *ptr = self.get_child(child); ptr != nullptr)
			return ptr;

		// Create a new node with a dummy command.
		auto path = self.m_parent != nullptr
		                  ? std::format("{} {}", self.m_parent->command().path(), child)
		                  : std::string(child);
		auto node = command_tree_node(&self, child, std::make_shared<nihil::command>(path));

		auto [it, ok] = self.m_children.emplace(child, std::move(node));
		if (!ok)
			throw std::logic_error("failed to insert command tree node");
		return &it->second;
	}

	// Get node's command.
	[[nodiscard]] auto command(this command_tree_node const &self) -> command &
	{
		return *self.m_command;
	}

	auto command(this command_tree_node &self, std::shared_ptr<nihil::command> command) -> void
	{
		// If we have a stub command, allow it to be replaced.  This happens when
		// a stub node is replaced with a real node during tree building.  However,
		// if we already have an invocable command, this is an error.
		if (self.m_command->invocable())
			throw std::logic_error("duplicate command");

		self.m_command = std::move(command);
	}

	// Yield all children of this node, depth-first.
	auto
	all_children(this command_tree_node const &self) -> generator<command_tree_node const &>
	{
		for (auto &&node : self.m_children | std::views::values) {
			co_yield node;
			co_yield elements_of(node.all_children());
		}
	}

	// Yield direct children in this node.
	auto children(this command_tree_node const &self)
	{
		return self.m_children | std::views::values;
	}

private:
	command_tree_node                       *m_parent = nullptr;
	std::string                              m_this_word;
	std::shared_ptr<nihil::command>          m_command;
	std::map<std::string, command_tree_node> m_children;
};

// The command tree stores commands in a tree structure suitable for searching.
export struct command_tree
{
	// Add a node to the tree.  Returns false if the node already exists.
	auto insert(this command_tree &self, std::shared_ptr<command> command) -> void
	{
		auto *this_node = &self.m_root_node;

		// Find the node for this key.
		for (auto &&this_word : command->path() | std::views::split(' '))
			this_node = this_node->get_or_create_child(std::string_view(this_word));

		// Set the new value.
		this_node->command(std::move(command));
	}

	// Find a node in the tree.  Returns the node and the remaining unprocessed arguments.
	[[nodiscard]] auto find(this command_tree const &self, std::ranges::range auto &&args)
	requires(std::constructible_from<std::string_view,
					 std::ranges::range_value_t<decltype(args)>>)
	{
		auto *this_node = &self.m_root_node;

		auto rest =
			args | std::views::take_while([&](auto &&str) {
				auto const *child = this_node->get_child(std::string_view(str));
				if (child == nullptr)
					return false;
				this_node = child;
				return true;
			});

		// Force evaluation of the view, otherwise the reference to this_node
		// would be dangling.
		auto taken = std::ranges::distance(args) - std::ranges::distance(rest);
		return std::make_pair(this_node, args | std::views::drop(taken));
	}

private:
	command_tree_node m_root_node;
};

// Build a command tree from the registry.
[[nodiscard]] auto build_command_tree() -> command_tree
{
	auto tree = command_tree();

	for (auto &&command : get_registry()) {
		// Throws std::logic_error on duplicates.
		tree.insert(command);
	}

	return tree;
}

} // namespace nihil