#pragma once

#include <psemek/bt/node.hpp>

#include <optional>

namespace psemek::bt
{

	template <typename Tree>
	struct repeat_node;

	template <typename Time, typename Event, typename ... Args>
	struct repeat_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;

		repeat_node(node_ptr<tree_type> child, std::optional<std::size_t> max_count)
			: child_(std::move(child))
			, max_count_(max_count)
		{}

		void start(Args ... args) override
		{
			child_->start(args...);
			child_started_ = true;
			repeat_count_ = 0;
		}

		status update(Time dt, Args ... args) override
		{
			if (!child_started_)
			{
				child_->start(args...);
				child_started_ = true;
			}

			auto result = child_->update(dt, args...);
			if (auto f = std::get_if<finished>(&result))
			{
				if (!(f->result))
					return finished{true};
				else
				{
					child_started_ = false;
					if (++repeat_count_ == max_count_)
						return finished{false};
					return running{};
				}
			}
			return result;
		}

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

	private:
		node_ptr<tree<Time, Event, Args...>> child_;
		bool child_started_ = false;
		std::optional<std::size_t> max_count_;
		std::size_t repeat_count_ = 0;
	};

	template <typename Tree>
	node_ptr<Tree> repeat(node_ptr<Tree> child, std::optional<std::size_t> max_count = std::nullopt)
	{
		return std::make_unique<repeat_node<Tree>>(std::move(child), max_count);
	}

}
