#include "tstrings.h"
#include "interpolate.h"
#include "error.h"
// #include "kblib.h"

#include "srell.hpp"

#include <iterator>
#include <cstdlib>

#include <iostream>

using kblib::a;
using namespace std::literals;

struct token {
	op type=op::null;
	std::string o_val;
};

[[nodiscard]] bool operator==(const token& a, const token& b)
{
	return std::tie(a.type, a.o_val) == std::tie(b.type, b.o_val);
}

[[nodiscard]] bool operator<(const token& a, const token& b)
{
	return std::tie(a.type, a.o_val) < std::tie(b.type, b.o_val);
}

namespace {

[[nodiscard]] std::string strip_escapes(std::string in)
{
	srell::regex double_lb("<<");
	in = srell::regex_replace(in, double_lb, "<");
	srell::regex double_rb(">>");
	in = srell::regex_replace(in, double_rb, ">");
	srell::regex backslash(R"__(\\(.))__");
	return srell::regex_replace(in, backslash, "$1");
}

class token_iterator {
public:
	using value_type = token;
	using pointer = const value_type*;
	using reference = const value_type&;
	using const_reference = const value_type&;
	using iterator_category = std::input_iterator_tag;
	
	token_iterator() noexcept = default;
	token_iterator(std::string_view _s) : data(_s) {
		// try {
			extract_next();
		// } catch (srell::regex_error& e) {
			// std::cerr<<e.what()<<'\n';
			// std::abort();
		// }
	}
	
	[[nodiscard]] const_reference operator*() const noexcept {
		return curr;
	}
	[[nodiscard]] pointer operator->() const noexcept {
		return &curr;
	}
	
	token_iterator& operator++() {
		do {
			extract_next();
		} while (!data.empty() && curr.type == op::null);
		return *this;
	}
	token_iterator operator++(int) {
		token_iterator tmp{*this};
		do {
			extract_next();
		} while (!data.empty() && curr.type == op::null);
		return tmp;
	}
	
	[[nodiscard]] explicit operator bool() const noexcept {
		return curr.o_val.size();
	}
	
	friend bool operator==(const token_iterator& a, const token_iterator& b) noexcept;
private:
	
	[[noreturn]] void error(ec err, const char* epos)
	{
		parse_error(err, data, epos);
	}
	
	[[noreturn]] void error(ec err)
	{
		parse_error(err, data, cpos);
	}
	
	//actual tokenizer
	void extract_next()
	{
		try {
		if (cpos == ptrdiff_t(data.size())) {
			//revert to sentinel on end of range
			if (data.size()) {
				if (parenDepth != -1) {
					error(ec::eos);
				}
				*this = {};
			}
			return;
		}
		struct match_action {
			srell::regex match;
			void(*action)(token_iterator*, srell::smatch);
		};
		const auto flag = srell::regex_constants::dotall;
		const static std::vector<match_action> textmode {
			{srell::regex{u8R"__(^((?:[^<>\\]|<<|>>|\\.)+))__"s, flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::text, strip_escapes(matches.str(1))};
			}},
			{srell::regex{u8R"__(^<>)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::early_close);
			}},
			{srell::regex{u8R"__(^<(\d+)>)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::argument, matches.str(1)};
			}},
			{srell::regex{u8R"__(^<([\w\d]+)>)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::special_argument, matches.str(1)};
			}},
			{srell::regex{u8R"__(^<(\.\.\.)>)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::ellipsis, matches.str(1)};
			}},
			{srell::regex{u8R"__(^<(#[\w\d]+)>)__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::overmarked_arg);
			}},
			{srell::regex{u8R"__(^<([^(>]*\\[^>]*)>)__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::escape_nostr);
			}},
			{srell::regex{u8R"__(^<\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::blank_fname);
			}},
			{srell::regex{u8R"__(^<raw\(([^)]*)\)>)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::text, matches.str(1)};
			}},
			{srell::regex{u8R"__(^<raw\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::invalid_raw);
			}},
			{srell::regex{u8R"__(^<([-\w\d+*/^?=.]+)\()__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->state = 1;
				t->parenDepth = 1;
				// don't actually return the < token
				t->curr = {op::push_func, matches.str(1)};
			}},
#if 0
			{srell::regex{u8R"__(^<(?=[^>]+>))__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->state = 1;
				//don't actually return the < token
				t->curr = {op::null, {}};
			}},
#endif
			{srell::regex{u8R"__(^<[^>]*>)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::bad_format);
			}},
			{srell::regex{u8R"__(^<[^>]*)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::extra_ab);
			}},
			
			{srell::regex{u8R"__(^.+)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::bad_format);
			}},
			{srell::regex{u8R"__(^$)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->extract_next();
			}},
		};
		const static std::vector<match_action> escmode0 = {{
			{srell::regex{u8R"__(^$)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::eos);
			}},
			{srell::regex{u8R"__(^\))__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::extra_rp);
			}},
#if 0
			{srell::regex{u8R"__(^\|)__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::extra_p);
			}},
			{srell::regex{u8R"__(^\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::blank_fname);
			}},
			{srell::regex{u8R"__(^raw\(([^)]*)\))__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::text, matches.str(1)};
			}},
			{srell::regex{u8R"__(^([-\w\d+*/^?=.]+)\()__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::push_func, matches.str(1)};
				++(t->parenDepth);
			}},
			{srell::regex{u8R"__(^\#[\w\d]+\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::call_arg);
			}},
			{srell::regex{u8R"__(^[^>(|]*\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::invalid_fname);
			}},
			{srell::regex{u8R"__(^(\.\.\.))__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::ellipsis, matches.str(1)};
			}},
			{srell::regex{u8R"__(^>$)__", flag},
			 [](token_iterator* t, srell::smatch) {
				 //end of stream, correctly formed
				t->state = 0;
				t->parenDepth = -1;
				t->curr = {op::text, ""s};
			}},
			{srell::regex{u8R"__(^>((?:[^<>\\]|<<|>>|\\.)+))__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->state = 0;
				t->parenDepth = -1;
				//don't actually return the > token
				t->curr = {op::text, strip_escapes(matches.str(1))};
			}},
#endif
			{srell::regex{u8R"__(^.+)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::bad_format);
			}},
		}};
		const static std::vector<match_action> escmode1 = {{
			{srell::regex{u8R"__(^$)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::eos);
			}},
			{srell::regex{u8R"__(^\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::blank_fname);
			}},
			{srell::regex{u8R"__(^raw\(([^)]*)\)\|?)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::text, matches.str(1)};
			}},
			{srell::regex{u8R"__(^raw\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::invalid_raw);
			}},
			{srell::regex{u8R"__(^([-\w\d+*/^?=.]+)\()__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::push_func, matches.str(1)};
				++(t->parenDepth);
			}},
			{srell::regex{u8R"__(^\#[\w\d]+\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::call_arg);
			}},
			{srell::regex{u8R"__(^\#(\d+)\|?)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::argument, matches.str(1)};
			}},
			{srell::regex{u8R"__(^\#([\w\d]+)\|?)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::special_argument, matches.str(1)};
			}},
			{srell::regex{u8R"__(^\#([\w\d#]+))__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::extra_h);
			}},
			{srell::regex{u8R"__(^\#[^()|<>])__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::early_close);
			}},
			{srell::regex{u8R"__(^[^>(|]*\()__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::invalid_fname);
			}},
			{srell::regex{u8R"__(^(\))\|>)__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::extra_p);
			}},
			{srell::regex{u8R"__(^(\))[^)|>]+>)__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::extra_c);
			}},
			{srell::regex{u8R"__(^(\))>)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				if (t->parenDepth != 1) {
					t->error(ec::extra_lp);
				}
				t->curr = {op::call_func, matches.str(1)};
				t->parenDepth = -1;
				t->state = 0;
			}},
			{srell::regex{u8R"__(^(\))\|?)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::call_func, matches.str(1)};
				--(t->parenDepth);
			}},
			{srell::regex{u8R"__(^(\.\.\.)\|?)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::ellipsis, matches.str(1)};
			}},
			{srell::regex{u8R"__(^<)__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::extra_ab);
			}},
			{srell::regex{u8R"__(^>)__", flag},
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::extra_lp);
			}},
			{srell::regex{u8R"__(^\|)__", flag},
			 [](token_iterator* t, srell::smatch) {
				//don't actually return the | token
				t->curr = {op::text, ""};
			}},
			{srell::regex{u8R"__(^((?:[^|()<>\\]|\\.)*)\|?)__", flag},
			 [](token_iterator* t, srell::smatch matches) {
				t->curr = {op::text, strip_escapes(matches.str(1))};
			}},
			
			{srell::regex{u8R"__(^.+)__", flag}, 
			 [](token_iterator* t, srell::smatch) {
				t->error(ec::bad_format);
			}},
		}};
		const static auto apply = [&](const std::vector<match_action>& cases) {
			std::string to_match = data.substr(cpos);
			srell::smatch m;
			for (auto& ma : cases) {
				if (srell::regex_search(to_match, m, ma.match)) {
					ma.action(this, m);
					cpos += m.str(0).size();
					return;
				}
			}
			error(ec::p_internal);
		};
		apply(state == 0 ? textmode : parenDepth == 0 ? escmode0 : escmode1);
		} catch (srell::regex_error& e) {
			std::cerr<<e.code()<<'\n';
			std::abort();
		}
	}
	std::string data;
	std::ptrdiff_t cpos{0};
	value_type curr;
	int parenDepth{-1};
	int state{0};
};

[[nodiscard]] bool operator==(const token_iterator& a, const token_iterator& b) noexcept {
	return (a.data == b.data)
		&& (a.cpos == b.cpos);
}

[[nodiscard]] bool operator!=(const token_iterator& a, const token_iterator& b) noexcept {
	return !(a == b);
}

class tokenize {
public:
	using iterator = token_iterator;
	using const_iterator = token_iterator;
	
	tokenize(std::string_view _s) noexcept : data(_s) {}
	
	[[nodiscard]] token_iterator begin() const {
		return {data};
	}
	
	[[nodiscard]] token_iterator end() const noexcept {
		return {};
	}
	
private:
	std::string_view data;
};
}

[[nodiscard]] bytecode_op token_to_bytecode(token in, const bytecode_machine& b)
{
	assert(in.type != op::null);
	auto at = [](const bytecode_machine& b, std::string_view s) -> transformer* {
		auto it = b.transformers.find(s);
		if (it == b.transformers.end()) {
			throw wordgen_error(ec::no_function, kblib::quoted(s), "\n", "");
		}
		return it->second.get();
	};
	switch (in.type) {
	case op::text:
	case op::ellipsis:
		return {in.type, std::string(in.o_val)};
		break;
	case op::push_func:
		return {in.type, at(b, in.o_val)};
		break;
	case op::call_func:
		return {in.type, {}};
		break;
	case op::argument:
		//Would use from_chars here but <charconv> doesn't work
		//So a copy has to be made because there's no other (char* first, char* last) -> int function in the standard library and stoi is the only one which doesn't take an NTBS (which in.o_val.data() is not)
		return {in.type, std::stoi(std::string(in.o_val))};
		break;
	case op::special_argument:
		return {in.type, in.o_val};
		break;
	default:
		return {op::null, {}};
	}
}

[[nodiscard]] bytecodes bytecode_machine::to_bytes(const std::string& in) const
{
	auto tokens = tokenize(in);
	//wrap entire string in a top-level echo, making eval simpler
	bytecodes ret = {
		{op::push_func, transformers.at(".").get()},
	};
	std::transform(tokens.begin(), tokens.end(), std::back_inserter(ret), [this](auto a) {return token_to_bytecode(a, *this);});
	ret.push_back({op::call_func, {}});
	return ret;
}

[[nodiscard]] std::string bytecode_machine::eval(const bytecodes& in, const std::vector<std::string>& p_args, const counters& c, int argc) const
{
	auto sa = [&](std::string_view a) {
		return special_argument(a);
	};
	
	ArgsType args{origin_tag{min_argument()}, -min_argument() + p_args.size()};
	std::copy(p_args.begin(), p_args.end(), args.begin()-min_argument());
	args[sa("d")] = std::to_string(c.depth);
	args[sa("D")] = std::to_string(c.depthLimit);
	args[sa("e")] = std::to_string(*c.expansions);
	args[sa("E")] = std::to_string(c.expansionsLimit);
	
	args[sa("c")] = std::to_string(args.high());
	args[sa("C")] = std::to_string(argc);
	
	args[sa("a")] = std::vector<std::string>{p_args.begin(), p_args.begin() + argc};
	args[sa("...")] = std::vector<std::string>{p_args.begin() + argc, p_args.end()};
	args[sa("A")] = std::vector<std::string>{p_args.begin(), p_args.end()};
	
	struct frame {
		const transformer* call = nullptr;
		argslist args;
	};
	std::vector<frame> stack;
	//This precondition is established by to_bytes
	assert(!in.empty());
	assert(in.front().type == op::push_func);
	for (auto&& op : in) {
		try {
			switch (op.type) {
			case op::push_func:
				stack.push_back({std::get<const transformer*>(op.data), {}});
				break;
			case op::call_func:
				{
					auto v = (*stack.back().call)(std::move(stack.back().args));
					stack.pop_back();
					if (!stack.empty()) {
						stack.back().args.push_back(std::move(v));
					} else {
						return stringize(v);
					}
				}
				break;
			case op::text:
				stack.back().args.push_back(std::get<std::string>(op.data));
				break;
			case op::argument:
				stack.back().args.push_back(kblib::variant_cast<arg_t>(args.at(std::get<int>(op.data))));
				break;
			case op::special_argument:
				stack.back().args.push_back(kblib::variant_cast<arg_t>(args.at(special_argument(std::get<std::string>(op.data)))));
				break;
			case op::ellipsis:
				stack.back().args.push_back(kblib::variant_cast<arg_t>(args.at(special_argument("..."))));
				break;
			default:
				throw std::invalid_argument("invalid bytecode received: type = " + std::to_string(kblib::etoi(op.type)));
			case op::null:
				;//do nothing
			}
		} catch (std::bad_variant_access&) {
			throw std::invalid_argument("invalid bytecode received: ill-formed operation\nop = " + to_string(op, *this));
		}
	}
	//Only way to get here is if we run out of bytecodes without popping the stack all the way
	throw std::invalid_argument("invalid bytecode received: not enough pops");
}

[[nodiscard]] int bytecode_machine::special_argument(std::string_view s) const
{
	for (const auto& [key, val] : arg_names) {
		if (key == s) {
			return val;
		}
	}
	throw std::runtime_error(kblib::concat("invalid argref '#", s, "'"));
}

[[nodiscard]] const std::string& bytecode_machine::reverse_lookup(const transformer& t) const noexcept(false)
{
	auto it = std::find_if(
		transformers.begin(), transformers.end(),
		[&t](const auto& u){return u.second.get() == &t;}
	);
	if (it != transformers.end()) {
		return it->first;
	} else {
		throw std::out_of_range("");
	}
}

[[nodiscard]] std::string to_string(const op_data& data, const bytecode_machine& b)
{
	return std::visit(kblib::visitor{
		[](std::monostate){return "NULL"s;},
		[](int i){return std::to_string(i);},
		[](const std::string& s){return kblib::quoted(s);},
		[&b](const transformer* t){
			try {
				return '"' + b.reverse_lookup(*t) + "(\"";
			} catch (std::out_of_range&) {
				return "INVALID"s;
			}
		},
	}, data);
}

[[nodiscard]] std::string to_string(const bytecode_op& op, const bytecode_machine& b)
{
	std::string what
		= "{type: ";
	what += std::to_string(kblib::etoi(op.type));
	what += ", data: ";
	what += to_string(op.data, b) + "}";
	return what;
}

/*enum class op : unsigned char {
	null=0,				//monostate
	text,					//string
	push_func,			//transformer*
	call_func,			//monostate
	argument,			//int
	special_argument,	//int
	ellipsis,			//string
};*/

[[nodiscard]] std::string serialize(const bytecode_op& op, const bytecode_machine& b)
{
	auto data = std::visit(kblib::visitor{
		[](std::monostate){return "null"s;},
		[](int i){return std::to_string(i);},
		[](const std::string& s){return kblib::quoted(s);},
		[&b](const transformer* t){
			try {
				return b.reverse_lookup(*t);
			} catch (std::out_of_range&) {
				return "INVALID"s;
			}
		},
	}, op.data);
	switch (op.type) {
		case op::null:					return "null	" + data;
		case op::text:					return "text	" + data;
		case op::push_func:			return "func(	" + data;
		case op::call_func:			return "call)	" + data;
		case op::argument:			return "arg 	" + data;
		case op::special_argument:	return "sparg	" + data;
		case op::ellipsis:			return "dots	" + data;
		default:							return "ERROR	" + data;
	}
}