#pragma once

#include <psemek/util/range.hpp>
#include <psemek/util/resource_container.hpp>

#include <cstdint>
#include <vector>

namespace psemek::util
{

	template <typename EdgesProvider, typename Index = std::uint32_t>
	struct dfs
	{
		dfs(Index count, EdgesProvider edges_provider)
			: visited_(count, visited_flag::not_visited)
			, parent_(count, null)
			, edges_provider_(std::move(edges_provider))
		{}

		struct cycle_iterator
		{
			using iterator_category = std::forward_iterator_tag;
			using value_type = Index;
			using reference_type = Index;
			using pointer = Index *;
			using difference_type = int;

			dfs * parent;
			Index i;

			Index operator *() const { return i; }

			cycle_iterator & operator ++()
			{
				i = parent->parent_[i];
				return *this;
			}

			cycle_iterator operator ++(int)
			{
				auto copy = *this;
				i = parent->parent_[i];
				return copy;
			}

			friend bool operator == (cycle_iterator const & i1, cycle_iterator const & i2)
			{
				return i1.i == i2.i;
			}

			friend bool operator != (cycle_iterator const & i1, cycle_iterator const & i2)
			{
				return i1.i != i2.i;
			}
		};

		template <typename CycleCallback>
		void run(CycleCallback && cycle_callback)
		{
			for (Index i = 0; i < visited_.size(); ++i)
				if (visited_[i] == visited_flag::not_visited)
					step(i, cycle_callback);
		}

	private:
		static constexpr Index null = static_cast<Index>(-1);

		enum class visited_flag : std::uint8_t
		{
			not_visited,
			visiting,
			visited,
		};

		std::vector<visited_flag> visited_;
		std::vector<Index> parent_;

		EdgesProvider edges_provider_;

		template <typename CycleCallback>
		void step(Index i, CycleCallback & cycle_callback)
		{
			visited_[i] = visited_flag::visiting;

			edges_provider_(i, [&](Index j)
			{
				switch (visited_[j])
				{
				case visited_flag::not_visited:
					parent_[j] = i;
					step(j, cycle_callback);
					break;
				case visited_flag::visiting:
					cycle_callback(util::range{cycle_iterator{this, i}, cycle_iterator{this, null}});
					break;
				case visited_flag::visited:
					break;
				}
			});

			visited_[i] = visited_flag::visited;
		}
	};

}
