#pragma once

#include <psemek/parser/parser.hpp>

#include <vector>
#include <optional>
#include <tuple>

namespace psemek::parser
{

	namespace detail
	{

		struct concat_tag{};

		template <std::size_t I, typename Buffer, typename Res>
		bool concat_helper(std::integral_constant<std::size_t, I>, Buffer &, Res &, error &)
		{
			return true;
		}

		template <std::size_t I, typename Buffer, typename Res, typename P, typename ... Ps>
		bool concat_helper(std::integral_constant<std::size_t, I>, Buffer & buf, Res & res, error & e, P const & p, Ps const & ... ps)
		{
			auto r = p.apply(buf);
			if (r.index() == 1)
			{
				e = std::get<1>(r);
				return false;
			}
			std::get<I>(res) = std::move(std::get<0>(r));
			return concat_helper(std::integral_constant<std::size_t, I + 1>{}, buf, res, e, ps...);
		}

		template <typename Buffer, typename ... Ps>
		auto concat_impl(Buffer & buf, Ps const & ... ps)
			-> result<std::tuple<decltype(detail::result_type(ps, buf))...>>
		{
			std::tuple<decltype(detail::result_type(ps, buf))...> res;
			error e;
			auto it = buf.it;
			if (!concat_helper(std::integral_constant<std::size_t, 0>{}, buf, res, e, ps...))
			{
				buf.it = it;
				return e;
			}
			return res;
		}

		struct one_of_tag{};

		template <std::size_t I, typename Buffer, typename Res>
		bool one_of_helper(std::integral_constant<std::size_t, I>, Buffer &, Res &, error &)
		{
			return false;
		}

		template <std::size_t I, typename Buffer, typename Res, typename P, typename ... Ps>
		bool one_of_helper(std::integral_constant<std::size_t, I>, Buffer & buf, Res & res, error & e, P const & p, Ps const & ... ps)
		{
			auto r = p.apply(buf);
			if (r.index() == 0)
			{
				res.template emplace<I>(std::move(std::get<0>(r)));
				return true;
			}
			e = std::get<1>(r);
			return one_of_helper(std::integral_constant<std::size_t, I + 1>{}, buf, res, e, ps...);
		}

		template <typename Buffer, typename ... Ps>
		auto one_of_impl(Buffer & buf, Ps const & ... ps)
			-> result<std::variant<decltype(detail::result_type(ps, buf))...>>
		{
			std::variant<decltype(detail::result_type(ps, buf))...> res;
			error e;
			auto it = buf.it;
			if (!one_of_helper(std::integral_constant<std::size_t, 0>{}, buf, res, e, ps...))
			{
				buf.it = it;
				return e;
			}
			return res;
		}

		template <typename R, typename Gen>
		struct recursive_helper
		{
			using P = decltype(std::declval<Gen>()(std::declval<parser_ref<recursive_helper>>()));
			P p;

			recursive_helper(Gen && gen)
				: p(std::forward<Gen>(gen)(parser_ref<recursive_helper>{this}))
			{}

			template <typename Buffer>
			result<R> apply(Buffer & buf) const
			{
				return p.apply(buf);
			}

			R parse(std::string_view text) const
			{
				return p.parse(text);
			}
		};

		template <typename R, typename Gen>
		struct recursive_impl
		{
			std::unique_ptr<recursive_helper<R, Gen>> p;

			recursive_impl(Gen && gen)
				: p(std::make_unique<recursive_helper<R, Gen>>(std::forward<Gen>(gen)))
			{}

			template <typename Buffer>
			result<R> apply(Buffer & buf) const
			{
				return p->apply(buf);
			}

			R parse(std::string_view text) const
			{
				return p->parse(text);
			}
		};

	}

	template <typename P, typename F>
	auto map(P && p, F && f)
	{
		return make_parser([p = std::forward<P>(p), f = std::forward<F>(f)](auto & buffer)
			-> result<decltype(f(detail::result_type(p, buffer)))>
		{
			auto res = p.apply(buffer);
			if (res.index() == 1)
				return std::get<1>(res);
			return f(std::get<0>(res));
		});
	}

	template <typename P, typename F>
	auto guard(P && p, F && f, std::string message = {})
	{
		if (message.empty())
			message = "guard failed";

		return make_parser([p = std::forward<P>(p), f = std::forward<F>(f), message = std::move(message)](auto & buffer)
			-> result<decltype(detail::result_type(p, buffer))>
		{
			auto it = buffer.it;
			auto res = p.apply(buffer);
			if (res.index() == 1)
				return std::get<1>(res);
			if (!f(std::get<0>(res)))
			{
				buffer.it = it;
				return error{buffer.offset(), message.data()};
			}
			return std::get<0>(res);
		});
	}

	template <typename P>
	auto maybe(P && p)
	{
		return make_parser([p = std::forward<P>(p)](auto & buffer)
			-> result<std::optional<decltype(detail::result_type(p, buffer))>>
		{
			auto res = p.apply(buffer);
			if (res.index() == 1)
				return std::nullopt;
			return std::get<0>(res);
		});
	}

	template <typename P>
	auto from_to(P && p, std::size_t min_count, std::optional<std::size_t> max_count, std::string message = {})
	{
		if (message.empty())
			message = "(unknown)";

		auto msg = std::string("expected at least ") + std::to_string(min_count) + std::string(" ") + message;

		return make_parser([p = std::forward<P>(p), min_count, max_count, msg = std::move(msg)](auto & buffer)
			-> result<std::vector<decltype(detail::result_type(p, buffer))>>
		{
			auto it = buffer.it;

			std::vector<decltype(detail::result_type(p, buffer))> res;
			while (true)
			{
				if (max_count && res.size() >= *max_count) break;

				auto pos = buffer.it;
				auto r = p.apply(buffer);
				if (r.index() == 1) break;
				if (buffer.it == pos)
					throw grammar_error("Infinite loop");
				res.push_back(std::move(std::get<0>(r)));
			}

			if (res.size() < min_count)
			{
				buffer.it = it;
				return error{buffer.offset(), msg.data()};
			}

			return res;
		});
	}

	template <typename P>
	auto many(P && p)
	{
		return from_to(std::forward<P>(p), 0, std::nullopt);
	}

	template <typename P>
	auto at_least(P && p, std::size_t count, std::string message = {})
	{
		return from_to(std::forward<P>(p), count, std::nullopt, std::move(message));
	}

	template <typename P>
	auto exactly(P && p, std::size_t count, std::string message = {})
	{
		return from_to(std::forward<P>(p), count, count, std::move(message));
	}

	template <typename ... Ps>
	auto concat(Ps && ... ps)
	{
		return make_parser([... ps = std::forward<Ps>(ps)](auto & buffer)
			-> result<std::tuple<decltype(detail::result_type(ps, buffer)) ...>>
		{
			return detail::concat_impl(buffer, ps...);
		}, detail::concat_tag{});
	}

	template <typename ... Ps>
	auto one_of(Ps && ... ps)
	{
		return make_parser([... ps = std::forward<Ps>(ps)](auto & buffer)
			-> result<std::variant<decltype(detail::result_type(ps, buffer)) ...>>
		{
			return detail::one_of_impl(buffer, ps...);
		}, detail::one_of_tag{});
	}

	template <typename P, typename F, typename A>
	auto fold(P && p, F && f, A && a)
	{
		return make_parser([p = std::forward<P>(p), f = std::forward<F>(f), a = std::forward<A>(a)](auto & buffer)
			-> result<std::remove_cvref_t<A>>
		{
			auto accum = a;
			while (true)
			{
				auto pos = buffer.it;
				auto res = p.apply(buffer);
				if (res.index() == 1)
					return accum;
				if (pos == buffer.it)
					throw grammar_error("Infinite loop");

				accum = f(accum, std::get<0>(res));
			}
		});
	}

	template <typename P, typename F>
	auto fold(P && p, F && f)
	{
		return make_parser([p = std::forward<P>(p), f = std::forward<F>(f)](auto & buffer)
			-> result<decltype(detail::result_type(p, buffer))>
		{
			auto res0 = p.apply(buffer);
			if (res0.index() == 1)
				return std::get<1>(res0);
			auto accum = std::move(std::get<0>(res0));
			while (true)
			{
				auto pos = buffer.it;
				auto res = p.apply(buffer);
				if (res.index() == 1)
					return accum;
				if (pos == buffer.it)
					throw grammar_error("Infinite loop");

				accum = f(std::move(accum), std::move(std::get<0>(res)));
			}
		});
	}

	struct skip_token{};

	template <typename P>
	auto skip_while(P && p)
	{
		return fold(std::forward<P>(p), [](auto const &, auto const &){ return skip_token{}; }, skip_token{});
	}

	template <typename R, typename Gen>
	auto recursive(Gen && gen)
	{
		return detail::recursive_impl<R, Gen>(std::forward<Gen>(gen));
	}

}
