#pragma once

#include <psemek/bt/node.hpp>
#include <psemek/bt/assert.hpp>

#include <optional>

namespace psemek::bt
{

	template <typename Tree>
	struct conditional_node;

	template <typename Time, typename Event, typename ... Args>
	struct conditional_node<tree<Time, Event, Args...>>
		: node<tree<Time, Event, Args...>>
	{
		using tree_type = tree<Time, Event, Args...>;
		using node_type = node<tree_type>;
		using typename node_type::running;
		using typename node_type::finished;
		using typename node_type::status;

		conditional_node(node_ptr<tree_type> condition, node_ptr<tree_type> true_branch, node_ptr<tree_type> false_branch)
			: condition_(std::move(condition))
			, true_branch_(std::move(true_branch))
			, false_branch_(std::move(false_branch))
		{}

		void start(Args ... args) override
		{
			condition_result_ = std::nullopt;
			condition_->start(args...);
		}

		status update(Time dt, Args ... args) override
		{
			try
			{
				if (condition_result_)
				{
					return current_node()->update(dt, args...);
				}
				else
				{
					auto result = condition_->update(dt, args...);

					if (auto finished = std::get_if<typename node_type::finished>(&result))
					{
						condition_result_ = finished->result;
						current_node()->start(args...);
						return current_node()->update(Time{}, args...);
					}

					return result;
				}
			}
			catch (assertion_failed_exception const &)
			{
				return finished{false};
			}
		}

		bool event(Event const & event, Args ... args) override
		{
			if (condition_result_)
			{
				return current_node()->event(event, args...);
			}
			else
				return condition_->event(event, args...);
		}

	private:
		node_ptr<tree_type> condition_;
		node_ptr<tree_type> true_branch_;
		node_ptr<tree_type> false_branch_;
		std::optional<bool> condition_result_;

		node<tree_type> * current_node()
		{
			return *condition_result_ ? true_branch_.get() : false_branch_.get();
		}
	};

	template <typename Tree>
	node_ptr<Tree> conditional(node_ptr<Tree> condition, node_ptr<Tree> true_branch, node_ptr<Tree> false_branch)
	{
		return std::make_unique<conditional_node<Tree>>(std::move(condition), std::move(true_branch), std::move(false_branch));
	}

}
