#pragma once
#include <vector>

/// Port of tree.py to C++
template <class T>
struct Tree {
	// FIXME make T not require default constructability
	struct Node {
	   protected:
		friend Tree;
		Node* parent = nullptr;
		std::vector<Node*> children {};
		T item {};

	   public:
		~Node() {
			for(auto& child : children)
				delete child;
		}

		T& data() { return item; }

		const T& data() const { return item; }

		bool is_leaf() const { return children.size() == 0; }

		Node* parent_node() { return parent; }

		Node* create_leaf(const T& item) {
			auto* node = new Node;
			node->parent = this;
			node->item = item;
			children.push_back(node);
			return node;
		}

		template <class Fn>
		void walk(Fn&& fn) {
			fn(this);
			if(!is_leaf()) {
				for(auto& child : children)
					child->walk(fn);
			}
		}

		template <class Pred>
		Node* find_child(Pred&& predicate) {
			if(predicate(this) == true)
				return this;

			for(auto& child : children)
				if(predicate(child) == true)
					return child;

			return nullptr;
		}

		std::size_t parent_count() const {
			auto* parent = this->parent;
			auto parent_count = 0z;
			while(parent) {
				parent_count++;
				parent = parent->parent;
			}
			return parent_count;
		}
	};

	Tree() { root = new Node; }

	~Tree() { delete root; }

	// Trees are not copyable but they can move
	Tree(const Tree&) = delete;
	Tree(Tree&&) = default;

	template <class Fn>
	void walk(Fn&& fn) {
		root->walk(fn);
	}

	Node* create_leaf(const T& item) { return root->create_leaf(item); }

	Node* root_node() { return root; }

   private:
	Node* root;
};