#include "MarkovNameGenerator.hpp"
#include <algorithm>
#include <cctype>

MarkovNameGenerator::MarkovNameGenerator()
{
	std::random_device rd;
	m_rng.seed(rd());
}

void MarkovNameGenerator::learn(const std::string &name)
{
	if (name.empty())
		return;

	std::string s;
	s.reserve(name.size() + 3);
	for (char c : name)
		s.push_back(static_cast<char>(std::tolower(
			static_cast<unsigned char>(c))));

	/* Pad with start/end tokens */
	std::string proc = "##" + s + "$";

	/* Order-2 transitions */
	for (size_t i = 0; i + 2 < proc.size(); ++i) {
		std::string prefix = proc.substr(i, 2);
		char next = proc[i + 2];
		m_order2[prefix][next]++;
	}

	/* Order-1 transitions (fallback) */
	for (size_t i = 0; i + 1 < proc.size(); ++i) {
		char prefix = proc[i];
		char next = proc[i + 1];
		m_order1[prefix][next]++;
	}

	/* Starting bigram: first two real characters after ## */
	if (s.size() >= 2)
		m_starters[s.substr(0, 2)]++;
	else if (s.size() == 1)
		m_starters[s + "$"]++;
}

void MarkovNameGenerator::learn(const std::vector<std::string> &names)
{
	for (const auto &n : names)
		learn(n);
}

std::string MarkovNameGenerator::generate(int minLen, int maxLen,
					  const std::unordered_set<std::string>
						  *exclude) const
{
	if (m_order2.empty())
		return "";

	for (int attempt = 0; attempt < 200; ++attempt) {
		std::string name;
		std::string prefix = "##";

		while (static_cast<int>(name.size()) < maxLen) {
			char next = 0;

			/* Try order-2 first */
			auto it2 = m_order2.find(prefix);
			if (it2 != m_order2.end() && !it2->second.empty()) {
				next = pickNext(it2->second);
			} else if (!prefix.empty()) {
				/* Fallback to order-1 */
				char fallback = prefix.back();
				auto it1 = m_order1.find(fallback);
				if (it1 != m_order1.end() && !it1->second.empty())
					next = pickNext(it1->second);
			}

			if (next == '$' || next == 0)
				break;
			if (next == '#')
				continue;

			name.push_back(next);
			prefix = prefix.substr(1) + next;
			if (prefix.size() < 2)
				prefix = "##";
		}

		if (static_cast<int>(name.size()) < minLen)
			continue;

		std::string result = titleCase(name);
		if (exclude && exclude->find(result) != exclude->end())
			continue;

		return result;
	}

	return "";
}

std::vector<std::string> MarkovNameGenerator::generateMany(
	int count, int minLen, int maxLen,
	const std::unordered_set<std::string> *exclude) const
{
	std::vector<std::string> results;
	results.reserve(count);
	std::unordered_set<std::string> localExclude;
	if (exclude)
		localExclude = *exclude;

	for (int i = 0; i < count; ++i) {
		std::string name = generate(minLen, maxLen, &localExclude);
		if (!name.empty()) {
			results.push_back(name);
			localExclude.insert(name);
		}
	}
	return results;
}

void MarkovNameGenerator::clear()
{
	m_order2.clear();
	m_order1.clear();
	m_starters.clear();
}

bool MarkovNameGenerator::empty() const
{
	return m_order2.empty();
}

void MarkovNameGenerator::saveToJson(nlohmann::json &j) const
{
	j["order2"] = nlohmann::json::object();
	for (const auto &pair : m_order2) {
		nlohmann::json inner = nlohmann::json::object();
		for (const auto &p2 : pair.second)
			inner[std::string(1, p2.first)] = p2.second;
		j["order2"][pair.first] = inner;
	}

	j["order1"] = nlohmann::json::object();
	for (const auto &pair : m_order1) {
		nlohmann::json inner = nlohmann::json::object();
		for (const auto &p2 : pair.second)
			inner[std::string(1, p2.first)] = p2.second;
		j["order1"][std::string(1, pair.first)] = inner;
	}

	j["starters"] = nlohmann::json::object();
	for (const auto &pair : m_starters)
		j["starters"][pair.first] = pair.second;
}

void MarkovNameGenerator::loadFromJson(const nlohmann::json &j)
{
	clear();

	if (j.contains("order2")) {
		for (auto &[k, v] : j["order2"].items()) {
			for (auto &[ck, cv] : v.items()) {
				if (!ck.empty())
					m_order2[k][ck[0]] = cv.get<int>();
			}
		}
	}

	if (j.contains("order1")) {
		for (auto &[k, v] : j["order1"].items()) {
			for (auto &[ck, cv] : v.items()) {
				if (!ck.empty() && !k.empty())
					m_order1[k[0]][ck[0]] = cv.get<int>();
			}
		}
	}

	if (j.contains("starters")) {
		for (auto &[k, v] : j["starters"].items())
			m_starters[k] = v.get<int>();
	}
}

char MarkovNameGenerator::pickNext(
	const std::unordered_map<char, int> &freqs) const
{
	int total = 0;
	for (const auto &p : freqs)
		total += p.second;
	if (total <= 0)
		return 0;

	std::uniform_int_distribution<int> dist(0, total - 1);
	int roll = dist(m_rng);

	for (const auto &p : freqs) {
		roll -= p.second;
		if (roll < 0)
			return p.first;
	}
	return freqs.begin()->first;
}

std::string MarkovNameGenerator::titleCase(const std::string &s)
{
	if (s.empty())
		return s;
	std::string r = s;
	r[0] = static_cast<char>(std::toupper(
		static_cast<unsigned char>(r[0])));
	return r;
}
