#define KBLIB_DEF_MACROS 1
#define KBLIB_IOMANIP 1
#include "kblib.h"

#include "randutils.hpp"
#include <pcg_random.hpp>

#include <cstring>
#include <iostream>
#include <regex>
#include <string_view>
#include <vector>
// #include <charconv>
#include <cmath>
#include <map>
#include <utility>

using namespace std::literals;

template <typename T = int>
std::vector<T> range(T begin, T end, T step = 1) {
	return kblib::buildiota<std::vector<T>>((end - begin) / step, begin, step);
}
// range(3) -> [0, 1, 2]
template <typename T = int>
std::vector<T> range(T end) {
	return kblib::buildiota<std::vector<T>>(end, 0);
}

// die(3) -> [1, 2, 3]
template <typename T = int>
std::vector<T> die(T sides) {
	return kblib::buildiota<std::vector<T>>(sides, 1);
}

// Index an array literal without naming its type
// Caveat: the return value must not be stored unless the argument is also
// stored This is indexable because temporaries live until the end of their
// full-expression, rather than sub-expression
template <typename T>
constexpr auto a(const std::initializer_list<T>& a) {
	return a.begin();
}
// use like:
// auto v = a({2, 3, 5, 7, 9, 11})[2];

static std::size_t get_if_there(const std::string& s, int& o, std::size_t pos) {
	// This C library function is literally half-const-correct
	char* newpos = const_cast<char*>(s.c_str() + pos);
	if (pos < s.size() && std::isdigit(s[pos]))
		o = std::strtol(newpos, &newpos, 10);
	return (newpos - s.c_str());

	// const char* first{s.c_str() + pos}, last{s.c_str() + s.size()};
	// auto [end, ec] = std::from_chars(first, last, o);
	// if (ec == std::errc::result_out__of_range) {
	// throw std::out_of_range(std::string(first, end) + "is out of range");
	// } //ignore invalid_argument
	// return (end - first);
}

template <typename E = std::invalid_argument>
void kbassert(bool cond, const char* err) {
	if (!cond)
		throw E(err);
}

static std::size_t get(const std::string& s, int& o, std::size_t pos = 0) {
	// This C library function is literally half-const-correct
	char* newpos = const_cast<char*>(s.c_str() + pos);
	kbassert(pos < s.size() && std::isdigit(s[pos]), "expected a number");
	o = std::strtol(newpos, &newpos, 10);
	return (newpos - s.c_str());

	// const char* first{s.c_str() + pos}, last{s.c_str() + s.size()};
	// auto [end, ec] = std::from_chars(first, last, o);
	// if (ec == std::errc::result_out__of_range) {
	// throw std::out_of_range(std::string(first, end) + "is out of range");
	// } else if (ec == std::errc::invalid_argument) {
	// throw std::invalid_argument("expected a number");
	// }
	// return (end - first);
}

enum class Mode { pretty = 1, text = 2 };
Mode operator&(Mode l, Mode r) {
	return static_cast<Mode>(static_cast<int>(l) & static_cast<int>(r));
}
Mode operator|(Mode l, Mode r) {
	return static_cast<Mode>(static_cast<int>(l) | static_cast<int>(r));
}
bool operator!(Mode m) { return m == Mode{0}; }

static void eval_and_print(Mode, const std::string&, int budget, int linewidth,
                           std::ostream& os = std::cout);
static void roll_and_print(const std::string&, int count,
                           std::ostream& os = std::cout);

static void print_help(std::string_view pname, std::ostream& os = std::cout) {
	os << "Usage: " << pname
	   << " [opts] dice...\n"
	      "Recognized options:\n"
	      "    -R:\n"
	      "        Act as a regular dice roller.\n"
	      "    -t:\n"
	      "        Print histograms as text\n"
	      "    -p:\n"
	      "        Print histograms graphically (default)\n"
	      "    -b:\n"
	      "        Print histograms first as text, and then graphically\n"
	      "    -B{budget}:\n"
	      "        Set computation budget in rounds. Default 5 million.\n"
	      "    -W{width}:\n"
	      "        Set line width. Default 80.\n"
	      "    -h:\n"
	      "        Print this message and exit\n"
	      "\n"
	      "Dice formats recognized are:\n"
	      "{number}\n"
	      "    A constant number.\n"
	      "[{num}]d{sides}[t{count}] (basic format)\n"
	      "    Roll {num} dice with {sides} sides, taking the highest {count} "
	      "dice and summing them.\n"
	      "    If omitted, {num} is assumed to be 1 and {count} is assumed to "
	      "equal {num}\n"
	      "[{num}][d{sides}]s[{difficulty}][;{mods}] (Storyteller-like)\n"
	      "    Roll {num} dice with {sides} sides, counting successes, that is "
	      "those which are >= {difficulty}, with {mods} modifications.\n"
	      "    If omitted, {sides} is assumed to be 10 and others are as above\n"
	      "    Modifications are:\n"
	      "        b: 1s are botches, i.e. negative successes.\n"
	      "        d: 10s count as two successes each.\n"
	      "        D{num}: Dice rolling {num} or greater count as two successes "
	      "each.\n"
	      "        r: 10s are rerolled until no 10s appear\n"
	      "        R{num}: Dice rolling {num} or greater are rerolled until "
	      "they no longer appear.\n"
	      "[{num}]d{sides}m{vector}[t{count}][;r{rT}] (map format)\n"
	      "    Roll {num} {sides}-sided dice, map the greatest {count} through "
	      "{vector}, and sum them. If rT is provided, dice with (pre-mapping) "
	      "values greater than it are rerolled until they no longer appear."
	      "\n"
	      "Formats may be concatenated with the + or - operators, which "
	      "corresponds to adding the numbers obtained by the sub-formats"
	      "\n"
	      "If a format string contains spaces, it must be quoted, and they are "
	      "ignored. Printing options apply only to those formats following "
	      "them, thus you may print multiple graphs in different formats.\n";
}

int main(int argc, char** argv) {
	// Get all command-line arguments as strings
	std::vector<std::string> args(argv + 1, argv + argc);
	Mode mode{Mode::pretty};
	int budget = 5'000'000;
	int linewidth = 74;

	if (args.size() == 0 || args[0] == "-h") {
		print_help(argv[0]);
		return 0;
	}
	bool spacer = false;
	bool once = false;
	for (auto&& a : args) {
		if (a.empty())
			continue;

		if (a[0] == '-') {
			switch (a[1]) {
			case 'p':
				mode = Mode::pretty;
				break;
			case 't':
				mode = Mode::text;
				break;
			case 'b':
				mode = Mode::text | Mode::pretty;
				break;
			case 'h':
				print_help(argv[0]);
				return 0;
			case 'B':
				get(a, budget, 2);
				break;
			case 'W':
				get(a, linewidth, 2);
				linewidth -= 6;
				break;
			case 'R':
				once = true;
				break;
			default:
				std::cerr << "Unrecognized option -" << a << ", continuing...\n";
			}
		} else {
			if (std::exchange(spacer, true))
				std::cout << '\n';
			if (once) {

			} else {
				eval_and_print(mode, a, budget, linewidth);
			}
		}
	}
}

struct roll {
	int num;
	int sides;
	std::vector<int> map;
	int count;
	int rT;
};

static const char* inv_mesg(char c) {
	static thread_local char emesg[] = "invalid specifier _";
	*(std::end(emesg) - 2) = c;
	return emesg;
}

static roll parse_basic(const std::string& s) {
	// std::clog<<s<<'\n';
	int num = 1;
	std::size_t pos = 0;
	// std::clog<<kblib::repeat(" "s, pos)<<"^\t";
	pos = get_if_there(s, num, pos);
	// std::clog<<num<<'\n';
	int sides = 0;
	kbassert(pos < s.size() && s[pos] == 'd', "must give number of sides");
	++pos;
	// std::clog<<kblib::repeat(" "s, pos)<<"^\t";
	pos = get(s, sides, pos);
	// std::clog<<sides<<'\n';
	int count = num;
	// std::clog<<kblib::repeat(" "s, pos)<<"^\n";
	if (pos < s.size()) {
		kbassert(s[pos] == 't', inv_mesg(s[pos]));
		++pos;
		// std::clog<<kblib::repeat(" "s, pos)<<"^\t";
		pos = get(s, count, pos);
		// std::clog<<count<<'\n';
	}
	return {num, sides, die(sides), count, sides};
}

static roll parse_st(const std::string& s) {
	int num = 1;
	std::size_t pos = 0;
	pos = get_if_there(s, num, pos);
	int sides = 10;
	if (s[pos] == 'd') {
		++pos;
		pos = get(s, sides, pos);
	}
	int diff = sides == 10 ? 7 : 0;
	kbassert(s[pos] == 's', inv_mesg(s[pos]));
	++pos;
	if (diff) {
		pos = get_if_there(s, diff, pos);
	} else {
		try {
			pos = get(s, diff, pos);
		} catch (std::invalid_argument& e) {
			throw std::invalid_argument("if sides is specified, diff must be too");
		}
	}
	--diff; // adjust to 0-based
	bool b = false;
	int dT = sides;
	int rT = sides;
	if (pos < s.size()) {
		kbassert(s[pos] == ';', inv_mesg(s[pos]));
		for (++pos; pos < s.size(); ++pos) {
			if (s[pos] == 'b')
				b = true;
			else if (s[pos] == 'd')
				dT = sides - 1;
			else if (s[pos] == 'r')
				rT = sides - 1;
			else if (s[pos] == 'D') {
				++pos;
				pos = get(s, dT, pos);
				kbassert(dT < sides,
				         "double threshold must be less than the number of sides");
			} else if (s[pos] == 'R') {
				++pos;
				pos = get(s, rT, pos);
				kbassert(rT < sides,
				         "reroll threshold must be less than the number of sides");
			} else {
				char emesg[] = "invalid modifier _";
				*(std::end(emesg) - 2) = s[pos];
				kbassert(false, emesg);
			}
		}
	}
	std::vector<int> map(sides);
	for (int i = 0; i < sides; ++i) {
		if (b && i == 0) {
			map[i] = -1;
		} else if (i >= diff) {
			map[i] = 1;
		}
		if (i >= dT) {
			++map[i];
		}
	}
	return {num, sides, map, num, rT};
}

static roll parse_map(const std::string& s) {
	int num = 1;
	std::size_t pos = 0;
	pos = get_if_there(s, num, pos);
	int sides = 0;
	kbassert(pos < s.size() && s[pos] == 'd', "must give number of sides");
	++pos;
	pos = get(s, sides, pos);
	std::vector<int> map;
	{
		kbassert(pos < s.size() && s[pos] == 'm',
		         "mapping must follow number of sides");
		using namespace kblib::io;
		++pos;
		kbassert(pos < s.size() && isdigit(s[pos]), "invalid mapping");
		std::istringstream ss(std::string(&s[pos], s.size() - pos));
		ss >> map;
		if (ss)
			pos += ss.tellg();
		else if (ss.eof())
			pos = s.size();
		kbassert((int)map.size() == sides,
		         "mapping must have as many values as {sides}");
	}
	int count = num;
	if (pos < s.size()) {
		kbassert(s[pos] == 't', inv_mesg(s[pos]));
		++pos;
		pos = get(s, sides, pos);
	}
	int rT = sides;
	if (pos < s.size()) {
		kbassert(s[pos] == ';', inv_mesg(s[pos]));
		for (; pos < s.size(); ++pos) {
			if (s[pos] == 'r') {
				++pos;
				pos = get(s, rT, pos);
				kbassert(rT < sides,
				         "reroll threshold must be less than the number of sides");
			} else {
				char emesg[] = "invalid modifier _";
				*std::prev(std::end(emesg)) = s[pos];
				kbassert(false, emesg);
			}
		}
	}
	return {num, sides, map, count, rT};
}

static roll parse(const std::string& s) {
	enum { basic = 0, st = 1, map = 2 };
	auto format = [](const std::string& s) {
		if (s.find('s') != std::string::npos)
			return st;
		if (s.find('m') != std::string::npos)
			return map;
		return basic;
	};
	///@TODO: make it work when rolls are added together
	return a({parse_basic, parse_st, parse_map})[format(s)](s);
}

std::map<int, int> operator+(std::map<int, int> l,
                             const std::map<int, int>& r) {
	std::map<int, int> t;
	if (l.empty())
		return r;
	if (r.empty())
		return l;
	for (auto&& i1 : l) {
		for (auto&& i2 : r) {
			t[i1.first + i2.first] = i1.second + i2.second;
		}
	}
	return t;
}

std::map<int, int>& operator+=(std::map<int, int>& l,
                               const std::map<int, int>& r) {
	return l = l + r;
}

static std::map<int, int> make_hist_stochastic(const roll& r, int budget) {
	std::map<int, int> output;
	randutils::random_generator<pcg32_k64> R;
	auto& rg = R.engine();
	std::vector<int> rollcounter(r.num);
	for (int _c = 0; _c < budget; ++_c) {
		std::generate(rollcounter.begin(), rollcounter.end(),
		              [&] { return rg(r.sides - 1) + 1; });
		auto sum = [&](const std::vector<int>& v) {
			return std::accumulate(v.begin(), v.end(), 0, [&r](int s, int x) {
				return s + r.map[x - 1];
			});
		};
		++output[sum(kblib::get_max_n<std::vector<int>>(
		    rollcounter.begin(), rollcounter.end(), r.count))];
	}
	return output;
}

static std::map<int, int> make_hist(const roll& r, int budget) {
	if (r.num == 0)
		return std::map<int, int>{};
	if (std::pow(r.sides, r.num) > budget) {
		return make_hist_stochastic(r, budget);
	}
	std::map<int, int> output;
	std::vector<int> rollcounter(r.num, 1);
	auto allones = [](const std::vector<int>& v) {
		return std::all_of(v.begin(), v.end(), [](int x) { return x == 1; });
	};
	auto addone = [](std::vector<int>& v, int max) {
		for (auto& x : v) {
			if (++x > max) {
				x = 1;
			} else
				break;
		}
		return v;
	};
	auto sum = [&](const std::vector<int>& v) {
		return std::accumulate(v.begin(), v.end(), 0,
		                       [&r](int s, int x) { return s + r.map[x - 1]; });
	};
	do {
		++output[sum(kblib::get_max_n<std::vector<int>>(
		    rollcounter.begin(), rollcounter.end(), r.count))];
	} while (!allones(addone(rollcounter, r.sides)));
	return output;
}

static void print(Mode m, const std::map<int, int>& h, std::ostream& os,
                  int linewidth) {
	int padsize = kblib::digitsOf(h.rbegin()->first);
	if (h.begin()->first < 0) {
		padsize = std::max({padsize, kblib::digitsOf(h.begin()->first), 2});
	}
	if (!!(m & Mode::text)) {
		for (auto&& v : h) {
			os << std::setw(padsize) << v.first << ": " << v.second << '\n';
		}
	}
	if (m == (Mode::pretty | Mode::text))
		os << '\n';
	if (!!(m & Mode::pretty)) {
		auto scaling = 1 + std::max_element(
		                       h.begin(), h.end(),
		                       [](const auto& l, const auto& r) {
			                       return l.second < r.second;
		                       })->second /
		                       linewidth;
		if (scaling > 1)
			os << u8"* Note: 1 × ▊ = "s << scaling << '\n';
		for (auto&& v : h) {
			os << std::setw(padsize) << v.first << "| "
			   << kblib::repeat(u8"▊"s, v.second / scaling) << '\n';
		}
	}
	os << std::flush;
}

static void eval_and_print(Mode m, const std::string& a, int budget,
                           int linewidth, std::ostream& os) {
	auto query = parse(a);
	auto hist = make_hist(query, budget);
	print(m, hist, os, linewidth);
}
#if 0
static std::vector<std::pair<std::string, int>> roll_single(const roll& r, int count) {
	std::vector<std::pair<std::string, int>> output;
	randutils::random_generator<pcg32_k64> R;
	auto& rg = R.engine();
	std::vector<int> rollcounter(r.num);
	for (int _c = 0; _c < budget; ++_c) {
		std::generate(rollcounter.begin(), rollcounter.end(),
		              [&] { return rg(r.sides - 1) + 1; });
		auto sum = [&](const std::vector<int>& v) {
			return std::accumulate(v.begin(), v.end(), 0, [&r](int s, int x) {
				return s + r.map[x - 1];
			});
		};
		++output[sum(kblib::get_max_n<std::vector<int>>(
		    rollcounter.begin(), rollcounter.end(), r.count))];
	}
	return output;

	do {
		output.push_back({"", sum(kblib::get_max_n<std::vector<int>>(
		                           rollcounter.begin(), rollcounter.end(), r.count))});
	} while (!allones(addone(rollcounter, r.sides)));
	return output;
}

static void roll_and_print(const std::string& a, int count,
                           std::ostream& os) {
	auto query = parse(a);


}
#endif