#pragma once

#include <re2/re2.h>
#include "unordered_dense.h"

#include <cassert>
#include <limits>
#include <optional>
#include <regex>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace tiktoken
{

	static auto _byte_pair_merge(
		const std::string &piece,
		const ankerl::unordered_dense::map<std::string, int> &ranks,
		std::function<int(int, int)> func) -> std::vector<int>
	{
		std::vector<std::pair<int, int>> parts;
		parts.reserve(piece.size() + 1);
		for (auto idx = 0U; idx < piece.size() + 1; ++idx)
		{
			parts.emplace_back(idx, std::numeric_limits<int>::max());
		}

		auto get_rank = [&piece, &ranks](
							const std::vector<std::pair<int, int>> &parts,
							int start_idx,
							int skip) -> std::optional<int>
		{
			if (start_idx + skip + 2 < parts.size())
			{
				auto s = parts[start_idx].first;
				auto e = parts[start_idx + skip + 2].first;
				auto key = piece.substr(s, e - s);
				auto iter = ranks.find(key);
				if (iter != ranks.end())
				{
					return iter->second;
				}
			}
			return std::nullopt;
		};

		for (auto i = 0U; i < parts.size() - 2; ++i)
		{
			auto rank = get_rank(parts, i, 0);
			if (rank)
			{
				assert(*rank != std::numeric_limits<int>::max());
				parts[i].second = *rank;
			}
		}

		while (true)
		{
			if (parts.size() == 1)
				break;

			auto min_rank = std::make_pair<int, int>(std::numeric_limits<int>::max(), 0);
			for (auto i = 0U; i < parts.size() - 1; ++i)
			{
				auto rank = parts[i].second;
				if (rank < min_rank.first)
				{
					min_rank = {rank, i};
				}
			}

			if (min_rank.first != std::numeric_limits<int>::max())
			{
				auto i = min_rank.second;
				auto rank = get_rank(parts, i, 1);
				if (rank)
				{
					parts[i].second = *rank;
				}
				else
				{
					parts[i].second = std::numeric_limits<int>::max();
				}
				if (i > 0)
				{
					auto rank = get_rank(parts, i - 1, 1);
					if (rank)
					{
						parts[i - 1].second = *rank;
					}
					else
					{
						parts[i - 1].second = std::numeric_limits<int>::max();
					}
				}

				parts.erase(parts.begin() + (i + 1));
			}
			else
			{
				break;
			}
		}
		std::vector<int> out;
		out.reserve(parts.size() - 1);
		for (auto i = 0U; i < parts.size() - 1; ++i)
		{
			out.push_back(func(parts[i].first, parts[i + 1].first));
		}
		return out;
	}

	static auto byte_pair_encode(
		const std::string &piece,
		const ankerl::unordered_dense::map<std::string, int> &ranks) -> std::vector<int>
	{
		if (piece.size() == 1)
		{
			return {ranks.at(piece)};
		}

		auto func = [&piece, &ranks](int start, int stop) -> int
		{
			std::string key = piece.substr(start, stop - start);
			return ranks.at(key);
		};

		return _byte_pair_merge(piece, ranks, func);
	}

	class tiktoken
	{
	public:
		tiktoken() = default;
		tiktoken(
			ankerl::unordered_dense::map<std::string, int> encoder,
			ankerl::unordered_dense::map<std::string, int> special_encoder,
			const std::string &pattern)
		{
			regex_ = std::make_unique<re2::RE2>("(" + pattern + ")");

			std::string special_pattern;
			for (const auto &item : special_encoder)
			{
				if (!special_pattern.empty())
				{
					special_pattern += "|";
				}
				special_pattern += re2::RE2::QuoteMeta(item.first);
			}
			if (special_pattern.empty())
			{
				special_regex_ = nullptr;
			}
			else
			{
				special_regex_ = std::make_unique<re2::RE2>("(" + special_pattern + ")");
			}

			encoder_ = std::move(encoder);
			special_tokens_encoder = std::move(special_encoder);

			for (const auto &[k, v] : encoder_)
			{
				decoder_.emplace(v, k);
			}
			// 不知道有什么用 debug会报错，先注释了再说
			// assert(encoder_.size() != decoder_.size() && "Encoder and decoder must be of equal length; maybe you had duplicate token indices in your encoder?");

			for (const auto &[k, v] : special_tokens_encoder)
			{
				special_tokens_decoder.emplace(v, k);
			}
		}

		auto encode_ordinary(const std::string &text) const -> std::vector<int>
		{
			return _encode_ordinary_native(text);
		}

		auto encode(const std::string &text) const -> std::vector<int>
		{
			return _encode_native(text, special_tokens_encoder).first;
		}

		auto encode_single_piece(const std::string &text) const -> std::vector<int>
		{
			auto iter = encoder_.find(text);
			if (iter != encoder_.end())
			{
				return {iter->second};
			}
			return byte_pair_encode(text, encoder_);
		}

		auto decode(const std::vector<int> &tokens) const -> std::string
		{
			return _decode_native(tokens);
		}

	private:
		auto split_with_allowed_special_token(
			re2::StringPiece &input,
			const ankerl::unordered_dense::map<std::string, int> &allowed_special) const -> std::pair<std::optional<std::string>, re2::StringPiece>
		{
			if (special_regex_ == nullptr)
				return {std::nullopt, input};

			auto start = input.begin();
			std::string special;
			while (true)
			{
				if (!re2::RE2::FindAndConsume(&input, *special_regex_, &special))
				{
					break;
				}

				if (allowed_special.count(special) == 1)
				{
					return {std::move(special), re2::StringPiece(start, input.begin() - start - special.size())};
				}
			}

			return {std::nullopt, input};
		}

		auto _encode_ordinary_native(const std::string &text) const -> std::vector<int>
		{
			std::vector<int> ret;
			re2::StringPiece input(text);

			std::string piece;
			while (re2::RE2::FindAndConsume(&input, *regex_, &piece))
			{
				auto iter = encoder_.find(piece);
				if (iter != encoder_.end())
				{
					ret.push_back(iter->second);
					continue;
				}
				auto tokens = byte_pair_encode(piece, encoder_);
				ret.insert(ret.end(), tokens.begin(), tokens.end());
			}
			return ret;
		}

		auto _encode_native(
			const std::string &text,
			const ankerl::unordered_dense::map<std::string, int> &allowed_special) const -> std::pair<std::vector<int>, int>
		{
			std::vector<int> ret;
			int last_piece_token_len = 0;
			re2::StringPiece input(text);

			while (true)
			{
				auto [special, sub_input] = split_with_allowed_special_token(input, allowed_special);
				std::string piece;
				while (re2::RE2::FindAndConsume(&sub_input, *regex_, &piece))
				{
					auto iter = encoder_.find(piece);
					if (iter != encoder_.end())
					{
						last_piece_token_len = 1;
						ret.push_back(iter->second);
						continue;
					}
					auto tokens = byte_pair_encode(piece, encoder_);
					last_piece_token_len = tokens.size();
					ret.insert(ret.end(), tokens.begin(), tokens.end());
				}

				if (special)
				{
					int token = special_tokens_encoder.at(*special);
					ret.push_back(token);
					last_piece_token_len = 0;
				}
				else
				{
					break;
				}
			}

			return {ret, last_piece_token_len};
		}

		auto _decode_native(const std::vector<int> &tokens) const -> std::string
		{
			std::string ret;
			ret.reserve(tokens.size() * 2);
			for (auto token : tokens)
			{
				std::string token_bytes;
				auto iter = decoder_.find(token);
				if (iter != decoder_.end())
				{
					token_bytes = iter->second;
				}
				else
				{
					iter = special_tokens_decoder.find(token);
					if (iter != special_tokens_decoder.end())
					{
						token_bytes = iter->second;
					}
					else
					{
						throw std::runtime_error("unknown token: " + std::to_string(token));
					}
				}
				ret += token_bytes;
			}
			return ret;
		}

		ankerl::unordered_dense::map<std::string, int> encoder_;
		ankerl::unordered_dense::map<std::string, int> special_tokens_encoder;
		ankerl::unordered_dense::map<int, std::string> decoder_;
		ankerl::unordered_dense::map<int, std::string> special_tokens_decoder;
		std::unique_ptr<re2::RE2> regex_;
		std::unique_ptr<re2::RE2> special_regex_;
	};

} // namespace tiktoken
