#pragma once

#include <psemek/util/flat_list.hpp>

#include <variant>
#include <tuple>
#include <memory>
#include <vector>
#include <optional>
#include <set>

namespace psemek::util
{

	template <typename Time, typename Event, typename ... Args>
	struct behavior_tree
	{
		struct running
		{};

		struct finished
		{
			bool result;
		};

		struct suspended
		{
			Time duration;
		};

		using status = std::variant<running, finished, suspended>;

		struct any_node_base
		{
			virtual void start(Args ...) = 0;
			virtual status update(Time, Args ...) = 0;
			virtual bool event(Event const &, Args ...) = 0;

			virtual ~any_node_base() {}
		};

		template <typename Node>
		struct any_node_impl
			: any_node_base
		{
			Node node;

			any_node_impl(Node node)
				: node(std::move(node))
			{}

			void start(Args ... args) override
			{
				node.start(args...);
			}

			status update(Time dt, Args ... args) override
			{
				return node.update(dt, args...);
			}

			bool event(Event const & e, Args ... args) override
			{
				return node.event(e, args...);
			}
		};

		struct any_node
		{
			any_node() = default;

			explicit operator bool() const
			{
				return static_cast<bool>(impl_);
			}

			template <typename Node>
			any_node(Node node)
				: impl_(new any_node_impl<Node>{std::move(node)})
			{}

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

			status update(Time dt, Args ... args)
			{
				return impl_->update(dt, args...);
			}

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

		private:
			std::unique_ptr<any_node_base> impl_;
		};

		struct updater
		{
			using instance_id = std::uint32_t;

			updater()
				: suspended_(suspended_comparator{this})
			{}

			template <typename Instance>
			instance_id add(Instance instance, Args ... args)
			{
				auto id = instances_.insert(instance_data{std::move(instance), {args...}, 0, std::nullopt});
				new_.push_back(id);
				return id;
			}

			size_t event(Event const & e, instance_id id)
			{
				auto & n = instances_[id];
				if (std::apply([&](Args ... args){ return n.tree.event(e, args...); }, n.args))
				{
					if (n.suspend && n.suspend->end > time_)
					{
						suspended_.erase(n.suspend->iterator);
						n.suspend->end = time_;
						n.suspend->iterator = suspended_.insert(id).first;
					}

					return 1;
				}
				return 0;
			}

			template <typename Filter>
			size_t event(Event const & e, Filter && filter)
			{
				size_t count = 0;

				for (auto id : active_)
				{
					auto & n = instances_[id];
					if (std::apply(filter, n.args))
						count += event(e, id);
				}

				for (auto it = suspended_.begin(); it != suspended_.end();)
				{
					auto id = *it++;
					auto & n = instances_[id];
					if (std::apply(filter, n.args))
						count += event(e, id);
				}

				return count;
			}

			size_t event(Event const & e)
			{
				return event(e, [](auto const & ...){ return true; });
			}

			struct update_statistics
			{
				int active_count = 0;
				int suspended_count = 0;
			};

			update_statistics update(Time dt)
			{
				time_ += dt;

				fill_dt(dt);
				activate();
				wake_up();

				update_statistics result;
				result.active_count = active_.size();
				result.suspended_count = suspended_.size();

				do_update();

				return result;
			}

		private:
			struct suspended_comparator
			{
				updater * parent;

				bool operator()(instance_id id1, instance_id id2) const
				{
					auto tie = [&](instance_id id) -> std::tuple<Time, instance_id> { return {parent->instances_[id].suspend->end, id}; };

					return tie(id1) < tie(id2);
				}
			};

			using suspended_set = std::set<instance_id, suspended_comparator>;

			struct instance_data
			{
				any_node tree;
				std::tuple<Args...> args;
				Time elapsed;

				struct suspend_data
				{
					Time start;
					Time end;
					suspended_set::iterator iterator;
				};

				std::optional<suspend_data> suspend;
			};

			flat_list<instance_data, instance_id> instances_;
			Time time_ = 0;

			std::vector<instance_id> new_;
			std::vector<instance_id> active_;
			std::vector<instance_id> new_active_;
			suspended_set suspended_;

			void fill_dt(Time dt)
			{
				for (auto id : active_)
					instances_[id].elapsed = dt;
			}

			void activate()
			{
				for (auto id : new_)
				{
					auto & n = instances_[id];
					n.elapsed = 0;
					std::apply([&](Args ... args){ n.tree.start(args...); }, n.args);
					active_.push_back(id);
				}
				new_.clear();
			}

			void wake_up()
			{
				while (!suspended_.empty())
				{
					auto id = *suspended_.begin();
					if (instances_[id].suspend->end > time_)
						break;
					suspended_.erase(suspended_.begin());

					active_.push_back(id);
					instances_[id].elapsed = time_ - instances_[id].suspend->start;
					instances_[id].suspend = std::nullopt;
				}
			}

			void do_update()
			{
				for (auto id : active_)
				{
					auto & n = instances_[id];
					auto result = std::apply([&](Args ... args){ return n.tree.update(n.elapsed, args...); }, n.args);
					if (auto s = std::get_if<finished>(&result))
					{
						new_.push_back(id);
					}
					else if (auto s = std::get_if<suspended>(&result))
					{
						n.elapsed = 0;
						n.suspend = {time_, time_ + s->duration, suspended_.end()};
						n.suspend->iterator = suspended_.insert(id).first;
					}
					else
						new_active_.push_back(id);
				}
				std::swap(active_, new_active_);
				new_active_.clear();
			}
		};

		template <typename ActionFn>
		struct action
		{
			ActionFn actionFn;

			action(ActionFn actionFn)
				: actionFn(std::move(actionFn))
			{}

			void start(Args ...)
			{}

			status update(Time dt, Args ... args)
			{
				actionFn(dt, args...);
				return finished{true};
			}

			bool event(Event const &, Args ...)
			{
				return false;
			}
		};

		template <typename TimeFn>
		struct wait
		{
			TimeFn timeFn;

			wait(TimeFn timeFn)
				: timeFn(std::move(timeFn))
			{}

			Time remaining_time = Time{0};

			void start(Args ... args)
			{
				remaining_time = timeFn(args...);
			}

			status update(Time dt, Args ...)
			{
				remaining_time -= dt;
				if (remaining_time <= 0)
					return finished{true};
				return suspended{remaining_time};
			}

			bool event(Event const &, Args ...)
			{
				return false;
			}
		};

		template <typename CondFn>
		struct condition
		{
			CondFn condFn;

			condition(CondFn condFn)
				: condFn(std::move(condFn))
			{}

			void start(Args ...)
			{}

			status update(Time, Args ... args)
			{
				return finished{condFn(args...)};
			}

			bool event(Event const &, Args ...)
			{
				return false;
			}
		};

		template <typename EventFn, typename Child>
		struct on_event
		{
			EventFn eventFn;
			Child child;

			std::optional<status> cached_status;

			on_event(EventFn eventFn, Child child)
				: eventFn(std::move(eventFn))
				, child(std::move(child))
			{}

			void start(Args ... args)
			{
				child.start(args...);
			}

			status update(Time dt, Args ... args)
			{
				if (cached_status)
				{
					auto result = *cached_status;
					cached_status = std::nullopt;
					return result;
				}

				return child.update(dt, std::forward<Args>(args)...);
			}

			bool event(Event const & e, Args ... args)
			{
				cached_status = eventFn(e, args...);
				if (!cached_status)
					return child.event(e, args...);
				return true;
			}
		};

		template <typename Child>
		struct success
		{
			Child child;

			success(Child child)
				: child(std::move(child))
			{}

			void start(Args ... args)
			{
				child.start(args...);
			}

			status update(Time dt, Args ... args)
			{
				auto child_status = child.update(dt, std::forward<Args>(args)...);
				if (child_status.index() == 1)
					return finished{true};
				return child_status;
			}

			bool event(Event const & e, Args ... args)
			{
				return child.event(e, args...);
			}
		};

		template <typename Child>
		struct failure
		{
			Child child;

			failure(Child child)
				: child(std::move(child))
			{}

			void start(Args ... args)
			{
				child.start(args...);
			}

			status update(Time dt, Args ... args)
			{
				auto child_status = child.update(dt, std::forward<Args>(args)...);
				if (child_status.index() == 1)
					return finished{false};
				return child_status;
			}

			bool event(Event const & e, Args ... args)
			{
				return child.event(e, args...);
			}
		};

		template <typename Child>
		struct repeat
		{
			Child child;

			repeat(Child child)
				: child(std::move(child))
			{}

			void start(Args ... args)
			{
				child.start(args...);
			}

			status update(Time dt, Args ... args)
			{
				auto result = child.update(dt, std::forward<Args>(args)...);
				if (auto f = std::get_if<finished>(&result))
				{
					if (!(f->result))
						return finished{true};
					else
					{
						start(args...);
						return running{};
					}
				}
				return result;
			}

			bool event(Event const & e, Args ... args)
			{
				return child.event(e, args...);
			}
		};

		template <typename Condition, typename Child>
		struct repeat_while
		{
			Condition condition;
			Child child;
			bool current_started = false;
			int current = 0;

			repeat_while(Condition condition, Child child)
				: condition(std::move(condition))
				, child(std::move(child))
			{}

			void start(Args ...)
			{
				current = 0;
				current_started = false;
			}

			status update(Time dt, Args ... args)
			{
				if (current == 0)
				{
					if (!current_started)
					{
						condition.start(args...);
						current_started = true;
					}

					auto result = condition.update(dt, args...);
					if (auto f = std::get_if<finished>(&result))
					{
						if (!(f->result))
							return finished{true};
						else
						{
							current = 1;
							current_started = false;
							dt = 0;
						}
					}
					else
						return result;
				}

				// current == 1

				if (!current_started)
				{
					child.start(args...);
					current_started = true;
				}

				auto result = child.update(dt, args...);
				if (auto f = std::get_if<finished>(&result))
				{
					if (!(f->result))
						return finished{false};
					else
					{
						current = 0;
						current_started = false;
						return running{};
					}
				}
				else
					return result;
			}

			bool event(Event const & e, Args ... args)
			{
				return condition.event(e, args...) || child.event(e, args...);
			}
		};

		template <typename Child>
		struct retry
		{
			Child child;
			std::optional<int> max_count;
			int count;

			retry(Child child, std::optional<int> max_count = std::nullopt)
				: child(std::move(child))
				, max_count(max_count)
			{}

			void start(Args ... args)
			{
				count = 0;
				child.start(args...);
			}

			status update(Time dt, Args ... args)
			{
				auto result = child.update(dt, std::forward<Args>(args)...);
				if (auto f = std::get_if<finished>(&result))
				{
					if (f->result)
						return finished{true};
					else
					{
						++count;
						if (max_count && count == *max_count)
							return finished{false};
						child.start(args...);
						return running{};
					}
				}
				return result;
			}

			bool event(Event const & e, Args ... args)
			{
				return child.event(e, args...);
			}
		};

		template <typename ... Children>
		struct sequence
		{
			std::tuple<Children...> children;
			bool current_started = false;
			size_t current = 0;

			sequence(Children ... children)
				: children{std::move(children)...}
			{}

			void start(Args ...)
			{
				if constexpr (sizeof...(Children) == 0)
				{
					current = 1;
				}
				else
				{
					current = 0;
					current_started = false;
				}
			}

			status update(Time dt, Args ... args)
			{
				return update_impl<0>(dt, args...);
			}

			bool event(Event const & e, Args ... args)
			{
				return event_impl<0>(e, args...);
			}

		private:
			template <size_t I>
			status update_impl(Time dt, Args ... args)
			{
				if constexpr (I == sizeof...(Children))
				{
					return finished{true};
				}
				else
				{
					if (current != I)
					{
						return update_impl<I + 1>(dt, args...);
					}
					else
					{
						if (!current_started)
						{
							std::get<I>(children).start(args...);
							current_started = true;
						}
						auto result = std::get<I>(children).update(dt, args...);
						dt = 0;
						if (auto f = std::get_if<finished>(&result))
						{
							if (f->result)
							{
								current_started = false;
								++current;
								return update_impl<I + 1>(dt, args...);
							}
						}
						return result;
					}
				}
			}

			template <size_t I>
			bool event_impl(Event const & e, Args ... args)
			{
				if constexpr (I == sizeof...(Children))
				{
					return false;
				}
				else
				{
					if (current == I)
					{
						return std::get<I>(children).event(e, args...);
					}
					else
					{
						return event_impl<I + 1>(e, args...);
					}
				}
			}
		};

		template <typename ... Children>
		struct selector
		{
			std::tuple<Children...> children;
			bool current_started = false;
			size_t current = 0;

			selector(Children ... children)
				: children{std::move(children)...}
			{}

			void start(Args ...)
			{
				if constexpr (sizeof...(Children) == 0)
				{
					current = 1;
				}
				else
				{
					current = 0;
					current_started = false;
				}
			}

			status update(Time dt, Args ... args)
			{
				return update_impl<0>(dt, args...);
			}

			bool event(Event const & e, Args ... args)
			{
				return event_impl<0>(e, args...);
			}

		private:
			template <size_t I>
			status update_impl(Time dt, Args ... args)
			{
				if constexpr (I == sizeof...(Children))
				{
					return finished{false};
				}
				else
				{
					if (current != I)
					{
						return update_impl<I + 1>(dt, args...);
					}
					else
					{
						if (!current_started)
						{
							std::get<I>(children).start(args...);
							current_started = true;
						}
						auto result = std::get<I>(children).update(dt, args...);
						dt = 0;
						if (auto f = std::get_if<finished>(&result))
						{
							if (!(f->result))
							{
								current_started = false;
								++current;
								return update_impl<I + 1>(dt, args...);
							}
							else
								return finished{true};
						}
						return result;
					}
				}
			}

			template <size_t I>
			bool event_impl(Event const & e, Args ... args)
			{
				if constexpr (I == sizeof...(Children))
				{
					return false;
				}
				else
				{
					if (current == I)
					{
						return std::get<I>(children).event(e, args...);
					}
					else
					{
						return event_impl<I + 1>(e, args...);
					}
				}
			}
		};
	};

}
