#include "parser.hpp"
#include <vector>
#include <utility>

template<typename PREV_T, typename FUN_T>
class map_wrapper;

template<typename PREV_T>
class path_wrapper;

class path_node
{
	// bool _is_int;
	std::string str;
	int i;
public:
	// path_node(int i) : _is_int(true), str(""), i(i) { }
	path_node(int i) : str(""), i(i) { }
	// path_node(const std::string& s) : _is_int(false), str(s), i(-1) { }
	path_node(const std::string& s) : str(s), i(-1) { }
	// int get_int() const { return _is_int ? i : -1; }
	int get_int() const { return i; }
	// std::string get_string() const { return !_is_int ? str : ""; }
	const std::string & get_string() const { return str; }
	bool operator==(const path_node& other) const { return get_int() == other.get_int() && get_string() == other.get_string(); }
	bool operator!=(const path_node& other) const { return !operator==(other); }
};

class path
{
	std::vector<path_node> list;
	path(const std::vector<path_node>& l) : list(l) { }
public:
	path() = default;

	path& extend(int i)
	{
		list.emplace_back(i);
		return *this;
	}

	path& extend(const std::string& str)
	{
		list.emplace_back(str);
		return *this;
	}

	int front_i() { return list.front().get_int(); }
	std::string front_s() { return list.front().get_string(); }
	int back_i() { return list.back().get_int(); }
	std::string back_s() { return list.back().get_string(); }
	path tail() { return path(std::vector<path_node>(list.begin() + 1, list.end())); }
	bool empty() { return list.empty(); }

	bool starts_with(const path& other) const
	{
		if (list.size() < other.list.size())
			return false;
		for (size_t i = 0; i < other.list.size(); i++)
			if (list[i] != other.list[i])
				return false;
		return true;
	}

	// this has to starts with prefix, otherwise ub
	path remove_prefix(const path& prefix) const
	{
		return path(std::vector<path_node>(list.begin() + prefix.list.size(), list.end()));
	}
};

class query_res;

template<typename PREV_IT, typename FUN_T>
class filter_iterator
{
	PREV_IT iterator;
	PREV_IT end;
	FUN_T fun;
public:
	filter_iterator(PREV_IT iterator, PREV_IT end, FUN_T fun)
		: iterator(iterator), end(end), fun(fun)
	{
		if (iterator != end)
		{
			const auto& temp = *iterator;
			if (!fun(temp.first, temp.second))
				operator++();
		}
	}
	const auto operator*() const 
	{
		return *iterator;
	}
	filter_iterator& operator++()
	{
		while (iterator != end)
		{
			++iterator;
			if (iterator != end)
			{
				const auto& temp = *iterator;
				if (fun(temp.first, temp.second))
					return *this;
			}
		}
		return *this;
	}
	bool operator!=(const filter_iterator<PREV_IT, FUN_T>& other) const { return iterator != other.iterator; }
};

template<typename PREV_T, typename FUN_T>
class filter_wrapper
{
public:
	const PREV_T& prev;
	FUN_T fun;

	filter_wrapper(const PREV_T& prev, FUN_T fun) : prev(prev), fun(fun) {}

	auto begin() const { return filter_iterator(prev.begin(), prev.end(), fun); }
	auto end() const { return filter_iterator(prev.end(), prev.end(), fun); }

	template <typename F>
	auto map(F f) const { return map_wrapper<filter_wrapper<PREV_T, FUN_T>, F>(*this, f); }

	template <typename F>
	auto filter(F f) const { return filter_wrapper<filter_wrapper<PREV_T, FUN_T>, F>(*this, f); }

	auto operator[](const path& p) const { return path_wrapper<filter_wrapper<PREV_T, FUN_T>>(*this, p); }
};

template<typename PREV_IT, typename FUN_T>
class map_iterator
{
	PREV_IT iterator;
	FUN_T fun;
public:
	map_iterator(PREV_IT iterator, FUN_T fun) : iterator(iterator), fun(fun) { }
	const auto operator*() const
	{
		auto temp = *iterator;
		return std::pair<path, decltype(fun(temp.first, std::move(temp.second)))>(temp.first, fun(temp.first, std::move(temp.second)));
	}
	map_iterator& operator++() { ++iterator; return *this; }
	bool operator!=(const map_iterator<PREV_IT, FUN_T>& other) const
	{
		return iterator != other.iterator;
	}
};

template<typename PREV_T, typename FUN_T>
class map_wrapper
{
public:
	const PREV_T& prev;
	FUN_T fun;

	map_wrapper(const PREV_T& prev, FUN_T fun) : prev(prev), fun(fun) { }

	auto begin() const { return map_iterator(prev.begin(), fun); }
	auto end() const { return map_iterator(prev.end(), fun); }

	template <typename F>
	auto map(F f) const { return map_wrapper<map_wrapper<PREV_T, FUN_T>, F>(*this, f); }

	template <typename F>
	auto filter(F f) const { return filter_wrapper<map_wrapper<PREV_T, FUN_T>, F>(*this, f); }

	auto operator[](const path& p) const { return path_wrapper<map_wrapper<PREV_T, FUN_T>>(*this, p); }
};

template<typename PREV_IT>
class path_iterator
{
	PREV_IT iterator;
	PREV_IT end;
	const path& _path;
public:
	bool operator!=(const path_iterator<PREV_IT>& other) const
	{
		return iterator != other.iterator;
	}
	path_iterator(PREV_IT iterator, PREV_IT end, const path& p) : iterator(iterator), end(end), _path(p)
	{
		if (iterator != end)
		{
			if (!(*iterator).first.starts_with(_path))
				operator++();
		}
	}
	const auto operator*() const
	{
		auto temp = *iterator;
		return std::pair<path, decltype(temp.second)>(temp.first.remove_prefix(_path), std::move(temp.second));
	}
	path_iterator& operator++()
	{
		while (iterator != end)
		{
			++iterator;
			if (iterator != end)
			{
				if ((*iterator).first.starts_with(_path))
					return *this;
			}
		}
		return *this;
	}
};

template<typename PREV_T>
class path_wrapper
{
	const PREV_T& prev;
	const path& _path;
public:
	path_wrapper(const PREV_T& prev, const path& p) : prev(prev), _path(p) { }

	auto begin() const { return path_iterator(prev.begin(), prev.end(), _path); }
	auto end() const { return path_iterator(prev.end(), prev.end(), _path); }

	template <typename F>
	auto map(F f) const { return map_wrapper(*this, f); }

	template <typename F>
	auto filter(F f) const { return filter_wrapper(*this, f); }

	auto operator[](const path& p) const { return path_wrapper<path_wrapper<PREV_T>>(*this, p); }
};

template<typename iterator_t>
class query_res_iterator
{
	iterator_t iterator;
public:
	query_res_iterator(iterator_t iterator) : iterator(iterator) { }
	query_res_iterator& operator++() { ++iterator; return *this; }
	bool operator!=(const query_res_iterator& other) const { return iterator != other.iterator; }
	const std::pair<path, const json_value&> operator*() const
	{
		const auto& temp = *iterator;
		return { temp.first, *(temp.second) };
	}
};

class query_res
{
public: // TODO - remove public
	std::vector<std::pair<path, const json_value*>> vec;
	json_ptr ptr_to_json_tree;
public:
	query_res(const std::vector<std::pair<path, const json_value*>>& vec, json_ptr ptr)
		: vec(vec), ptr_to_json_tree(std::move(ptr)) { }

	auto begin() const { return query_res_iterator(vec.begin()); }
	auto end() const { return query_res_iterator(vec.end()); }

	template <typename F>
	auto map(F f) const { return map_wrapper(*this, f); }

	template <typename F>
	auto filter(F f) const { return filter_wrapper(*this, f); }

	auto operator[](const path& p) const { return path_wrapper<query_res>(*this, p); }

};

static void post_order(const json_value* ptr, std::vector<std::pair<path, const json_value*>>& vector, const path& p)
{
	if (ptr->type() == json_type::integer)
		return vector.push_back({ std::move(p), ptr });

	if (ptr->type() == json_type::object)
	{
		auto& map = ptr->getMap();
		for (auto& [str, val] : map)
		{
			path new_path = p;
			new_path.extend(str);
			post_order(val.get(), vector, new_path);
		}
		return vector.push_back({ std::move(p), ptr });
	}

	for (int i = 0; i < ptr->length(); i++)
	{
		path new_path = p;
		new_path.extend(i);
		post_order(&ptr->item_at(i), vector, new_path);
	}
	return vector.push_back({ std::move(p), ptr });
}

query_res query(const std::string& str)
{
	auto temp = json_parse(str);
	std::vector<std::pair<path, const json_value*>> res;
	post_order(temp.get(), res, path());
	return query_res(std::move(res), std::move(temp));
}