#pragma once

#include <psemek/bt/node.hpp>

#include <optional>

namespace psemek::bt
{

	namespace detail
	{

		template <typename Tree>
		struct trampoline_state
		{
			std::optional<typename bt::node<Tree>::status> cached_status;
		};

	}

	template <typename Tree>
	struct trampoline_intercept_node;

	template <typename Tree>
	struct trampoline_rethrow_node;

	template <typename Time, typename Event, typename ... Args>
	struct trampoline_intercept_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::finished;
		using typename node_type::status;

		trampoline_intercept_node(node_ptr<tree_type> child, std::shared_ptr<detail::trampoline_state<tree<Time, Event, Args...>>> state)
			: child_(std::move(child))
			, state_(std::move(state))
		{}

		void start(Args ... args) override
		{
			child_->start(args...);
		}

		status update(Time dt, Args ... args) override
		{
			auto result = child_->update(dt, args...);
			if (auto f = std::get_if<finished>(&result))
				state_->cached_status = *f;
			return result;
		}

		bool event(Event const & event, Args ... args) override
		{
			return child_->event(event, args...);
		}

	private:
		node_ptr<tree<Time, Event, Args...>> child_;
		std::shared_ptr<detail::trampoline_state<tree<Time, Event, Args...>>> state_;
	};

	template <typename Time, typename Event, typename ... Args>
	struct trampoline_rethrow_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::finished;
		using typename node_type::status;

		trampoline_rethrow_node(node_ptr<tree_type> child, std::shared_ptr<detail::trampoline_state<tree<Time, Event, Args...>>> state)
			: child_(std::move(child))
			, state_(std::move(state))
		{}

		void start(Args ... args) override
		{
			child_->start(args...);
		}

		status update(Time dt, Args ... args) override
		{
			auto result = child_->update(dt, args...);
			if (std::get_if<finished>(&result) && state_->cached_status)
			{
				result = *state_->cached_status;
				state_->cached_status = std::nullopt;
			}
			return result;
		}

		bool event(Event const & event, Args ... args) override
		{
			return child_->event(event, args...);
		}

	private:
		node_ptr<tree<Time, Event, Args...>> child_;
		std::shared_ptr<detail::trampoline_state<tree<Time, Event, Args...>>> state_;
	};

	template <typename Tree>
	struct trampoline
	{
		trampoline()
			: state_(std::make_shared<detail::trampoline_state<Tree>>())
		{}

		node_ptr<Tree> intercept(node_ptr<Tree> child)
		{
			return std::make_unique<trampoline_intercept_node<Tree>>(std::move(child), state_);
		}

		node_ptr<Tree> rethrow(node_ptr<Tree> child)
		{
			return std::make_unique<trampoline_rethrow_node<Tree>>(std::move(child), state_);
		}

	private:
		std::shared_ptr<detail::trampoline_state<Tree>> state_;
	};

}
