#pragma once

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

namespace psemek::bt
{

	template <typename Tree, typename DurationFn>
	struct timeout_node;

	template <typename Time, typename Event, typename ... Args, typename DurationFn>
	struct timeout_node<tree<Time, Event, Args...>, DurationFn>
		: 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;

		timeout_node(node_ptr<tree_type> child, DurationFn && duration_fn)
			: child_(std::move(child))
			, duration_fn_(std::move(duration_fn))
		{}

		void start(Args ... args) override
		{
			failed_start_ = false;

			try
			{
				timer_ = duration_fn_(args...);
				child_->start(args...);
			}
			catch (assertion_failed_exception const &)
			{
				failed_start_ = true;
			}
		}

		status update(Time dt, Args ... args) override
		{
			if (failed_start_)
				return finished{false};

			timer_ -= dt;
			if (timer_ <= Time{0})
				return finished{false};

			return child_->update(dt, args...);
		}

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

	private:
		node_ptr<tree<Time, Event, Args...>> child_;
		DurationFn duration_fn_;
		Time timer_ = Time{};
		bool failed_start_ = false;
	};

	namespace detail
	{

		template <typename Tree, typename DurationFn>
		struct timeout_helper
		{
			static node_ptr<Tree> make(node_ptr<Tree> child, DurationFn && duration_fn)
			{
				return std::make_unique<timeout_node<Tree, DurationFn>>(std::move(child), std::move(duration_fn));
			}
		};

		template <typename Time, typename Event, typename ... Args>
		struct timeout_helper<tree<Time, Event, Args...>, Time>
		{
			static node_ptr<tree<Time, Event, Args...>> make(node_ptr<tree<Time, Event, Args...>> child, Time duration)
			{
				auto duration_fn = [duration](Args ...){ return duration; };
				return std::make_unique<timeout_node<tree<Time, Event, Args...>, decltype(duration_fn)>>(std::move(child), std::move(duration_fn));
			}
		};

	}

	template <typename Tree, typename DurationFn>
	node_ptr<Tree> timeout(node_ptr<Tree> child, DurationFn && duration_fn)
	{
		return detail::timeout_helper<Tree, DurationFn>::make(std::move(child), std::move(duration_fn));
	}

}
