#pragma once

#include <variant>
#include <vector>

#include "error.h"

namespace saw {
/**
 * Container with a simplistic approach to a branch
 */
template<typename T>
class branch;

/**
 * Tree object holding branches.
 *
 * The name comes from the fact a tree is acting as a node while the branch class is the
 * edge to a leaf or other nodes. A tree holds the branches while the branch either has
 * a leaf or another sub tree.
 */
template<typename T>
class tree final {
private:
	/**
	 * Object holding the treeed branch instances
	 */
    using branch = std::variant<T, tree<T>>;
	std::vector<branch> children_;
public:
	/**
	 * Default constructor
	 */
	tree() = default;

	/**
	 * Destructor
	 */
	~tree() = default;

	SAW_FORBID_COPY(tree);
	SAW_DEFAULT_MOVE(tree);

	/**
	 * Reserve space for siz elements
	 */
	error_or<void> reserve(std::size_t siz){
		try{
			children_.reserve(siz);
		}catch(const std::exception& e){
			(void) e;

			return make_error<err::out_of_memory>();
		}

		return void_t{};
	}

	/**
	 * Add a branch with a leaf attached to the tree
	 */
	error_or<std::size_t> add(T leaf) {
		std::size_t index = size();
		try {
			/**
			 * Technically we're adding a leaf on a branch
			 */
			children_.emplace_back(std::move(leaf));
		}catch(const std::exception& e){
			(void)e;

			return make_error<err::critical>();
		}

		return index;
	}

	/**
	 * Add a branch to the tree with a tree attached
	 */
	error_or<std::size_t> add() {
		std::size_t index = size();
		try {

			children_.emplace_back(tree{});
		}catch(const std::exception& e){
			(void)e;

			return make_error<err::critical>();
		}

		return index;
	}

	/**
	 * Returns the amount of branches contained within this tree level
	 */
	std::size_t size() const {
		return children_.size();
	}

	/**
	 * Returns the branch at i
	 */
	branch& at(std::size_t i){
		return children_.at(i);
	}

	/**
	 * Returns the branch at i
	 */
	const branch& at(std::size_t i) const {
		return children_.at(i);
	}
};

template<typename T>
class branch final {
private:
	using type = std::variant<tree<T>,T>;
	type tov_;

	/**
	 * We're friend classing the tree since it's way easier this way and the branch and tree
	 * class are intertwined heavily anyway.
	 */
	friend class tree<T>;
public:
	/**
	 * 
	 */
	branch():tov_{tree<T>{}}{}

	branch(tree<T> nd):tov_{std::move(nd)}{}

	branch(T val):tov_{std::move(val)}{}

	SAW_FORBID_COPY(branch);
	SAW_DEFAULT_MOVE(branch);

	template<typename NT>
	bool is() const {
		return std::holds_alternative<NT>(tov_);
	}

	bool is_tree() const {
		return std::holds_alternative<tree>(tov_);
	}

	bool is_value() const {
		return std::holds_alternative<T>(tov_);
	}

	template<typename NT>
	NT& get() {
		return std::get<NT>(tov_);
	}
	
	template<typename NT>
	const NT& get() const {
		return std::get<NT>(tov_);
	}

	tree<T>& get_tree(){
		return std::get<tree<T>>(tov_);
	}

	const tree<T>& get_tree() const {
		return std::get<tree<T>>(tov_);
	}

	T& get_value(){
		return std::get<T>(tov_);
	}

	const T& get_value() const {
		return std::get<T>(tov_);
	}

	template<typename NT>
	error_or<NT> extract(){
		if(!is<NT>()){
			return make_error<err::invalid_state>();
		}
		
		NT nd = std::move(std::get<NT>(tov_));
		tov_ = tree<T>{};

		return nd;
	}
	
	template<typename NT>
	error_or<NT> replace(type nd){
		auto eon = extract<NT>();
		if(eon.is_error()){
			return eon;
		}

		tov_ = std::move(nd);

		return eon;
	}

	error_or<tree<T>> extract_tree() {
		return extract<tree<T>>();
	}

	error_or<tree<T>> replace_tree(type nd){
		return replace<tree<T>>(std::move(nd));
	}

	error_or<T> extract_value() {
		return extract<T>();
	}

	error_or<T> replace_value(type nd){
		return replace<T>(std::move(nd));
	}
};
}