#pragma once

#include <variant>
#include <vector>

#include "error.h"

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

/**
 * Branch object holding sub trees
 */
template<typename T>
class branch final {
private:
	std::vector<tree<T>> children_;
public:
	branch() = default;

	~branch() = default;

	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{};
	}

	error_or<std::size_t> add(T leaf) {
		std::size_t index = size();
		try {
			children_.push_back(std::move(leaf));
		}catch(const std::exception& e){
			(void)e;

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

		return index;
	}

	error_or<std::size_t> add() {
		std::size_t index = size();
		try {
			children_.push_back(branch{});
		}catch(const std::exception& e){
			(void)e;

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

		return index;
	}

	std::size_t size() const {
		return children_.size();
	}

	tree<T>& at(std::size_t i){
		return children_.at(i);
	}

	const tree<T>& at(std::size_t i) const {
		return children_.at(i);
	}
};

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

	friend class branch<T>;
public:
	tree():tov_{branch<T>{}}{}

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

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

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

	bool is_branch() const {
		return std::holds_alternative<branch>(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_);
	}

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

	const branch<T>& get_branch() const {
		return std::get<branch<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_ = branch<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<branch<T>> extract_branch() {
		return extract<branch<T>>();
	}

	error_or<branch<T>> replace_branch(type nd){
		return replace<branch<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));
	}
};
}