#pragma once

#include <variant>
#include <vector>

#include "error.h"

namespace saw {
template<typename T>
class tree_or;

template<typename T>
class tree final {
private:
	std::vector<tree_or<T>> children_;
public:
	tree() = default;
	~tree() = default;

	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(tree{});
		}catch(const std::exception& e){
			(void)e;

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

		return index;
	}

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

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

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

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

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

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

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

	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>(tov_);
	}

	const tree<T>& get_tree() const {
		return std::get<tree>(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));
	}
};
}