module;

#include <charconv>
#include <iterator>
#include <variant>
#include <optional>
#include <unordered_map>
#include <string_view>
#include <ranges>
#include <format>
#include <iostream>

export module jlx:interpreter;
import :ast;
import :type_checker;

namespace jlx {
	using runtime_value = std::variant<std::string, int8_t, int16_t, int32_t, int64_t, uint8_t, uint16_t, uint32_t, uint64_t, float, double, bool>;

	struct interpreter_variable {
		std::string type;
		bool writable;
		runtime_value value;
	};

	struct interpreter_variable_scope {
		std::unordered_map<std::string, interpreter_variable> variables;
	};

	struct interpreter_function_scope {
		std::optional<runtime_value> return_value;
		bool returned = false;
	};

	struct interpreter_error : public std::runtime_error {
		interpreter_error(std::string_view msg, const token& t) :
		std::runtime_error(std::format("Interpreter error: {} at {}:{}:{}", msg, t.source_file, t.line, t.col).c_str()) {

		}
	};

	template<class... Ts>
	struct overload : Ts... { using Ts::operator()...; };

	struct base_interpreter {
		virtual std::optional<runtime_value> execute_statement(const statement&) = 0;
		virtual void push_function_scope() = 0;
		virtual void pop_function_scope() = 0;
		virtual void push_variable_scope() = 0;
		virtual void pop_variable_scope() = 0;
		virtual interpreter_function_scope& current_function_scope() = 0;
		virtual interpreter_variable_scope& current_variable_scope() = 0;
		virtual bool is_in_function() = 0;
		virtual void inject_value(const std::string_view&, runtime_value) = 0;

		virtual ~base_interpreter() = default;
	};

	struct interpreter_function {
		virtual std::optional<runtime_value> call(const std::vector<runtime_value>&, base_interpreter&) = 0;

		virtual ~interpreter_function() = default;
	};

	template<typename A, typename B>
	concept Operable = requires(A a, B b)
	{
		requires sizeof(A) >= sizeof(B);
		{ a + b };
		{ a - b };
		{ a * b };
		{ a / b };
		{ a % b };
		std::numeric_limits<A>::is_specialized;
		std::numeric_limits<B>::is_specialized;
		std::numeric_limits<A>::is_integer == std::numeric_limits<B>::is_integer;
		std::numeric_limits<A>::is_signed == std::numeric_limits<B>::is_signed;
	};

	template<typename A, typename B>
	concept Comparable = requires {
		requires std::is_same_v<A, B> ||
		requires (A a, B b)
		{
			requires std::numeric_limits<A>::is_signed == std::numeric_limits<B>::is_signed;
			{ a == b } -> std::convertible_to<bool>;
			{ a != b } -> std::convertible_to<bool>;
		};
	};

	template<typename A, typename B>
	concept Orderable = requires
	{
		requires std::is_same_v<A, B> || requires (A a, B b){
			requires std::numeric_limits<std::remove_cvref_t<A>>::is_signed == std::numeric_limits<std::remove_cvref_t<B>>::is_signed;
			{ a < b } -> std::convertible_to<bool>;
			{ a > b } -> std::convertible_to<bool>;
			{ a >= b } -> std::convertible_to<bool>;
			{ a <= b } -> std::convertible_to<bool>;
		};
	};

	static_assert(!Operable<std::string, uint32_t>);
	static_assert(Operable<uint32_t, uint16_t>);
	static_assert(!Operable<uint16_t, uint32_t>);
	static_assert(!Comparable<std::string, signed char>);
	static_assert(!Comparable<uint32_t, int32_t>);
	static_assert(!Orderable<std::string, int32_t>);
	static_assert(!Orderable<int8_t, uint32_t>);

	template<typename A, typename B>
	struct arithmetic_evaluator {
		std::remove_cvref_t<A> sum(A a, B b) requires Operable<A, B> {
			return a + b;
		}

		std::remove_cvref_t<A> mul(A a, B b) requires Operable<A, B> {
			return a * b;
		}

		std::remove_cvref_t<A> div(A a, B b) requires Operable<A, B> {
			return a / b;
		}

		std::remove_cvref_t<A> sub(A a, B b) requires Operable<A, B> {
			return a - b;
		}

		std::remove_cvref_t<A> mod(A a, B b) requires Operable<A, B> {
			return a % b;
		}

		std::remove_cvref_t<B> sum(A a, B b) requires Operable<B, A> && (!Operable<A, B>) {
			return a + b;
		}

		std::remove_cvref_t<B> mul(A a, B b) requires Operable<B, A> && (!Operable<A, B>) {
			return a * b;
		}

		std::remove_cvref_t<B> div(A a, B b) requires Operable<B, A> && (!Operable<A, B>) {
			return a / b;
		}

		std::remove_cvref_t<B> sub(A a, B b) requires Operable<B, A> && (!Operable<A, B>) {
			return a - b;
		}

		std::remove_cvref_t<B> mod(A a, B b) requires Operable<B, A> && (!Operable<A, B>) {
			return a % b;
		}

		std::string sum(std::string a, B b) requires (!std::is_same_v<std::remove_cvref_t<B>, std::string>) {
			return a + std::to_string(b);
		}

		std::string sum(std::string a, std::string b) {
			return a + b;
		}

		[[noreturn]] A sum(A, B) requires (!std::is_same_v<std::remove_cvref_t<A>, std::string> && !Operable<A, B> && !Operable<B, A>) {
			throw std::runtime_error("Arithmetically-incompatible types");
		}

		[[noreturn]] A mul(A, B) {
			throw std::runtime_error("Arithmetically-incompatible types");
		}

		[[noreturn]] A div(A, B) {
			throw std::runtime_error("Arithmetically-incompatible types");
		}

		[[noreturn]] A sub(A, B) {
			throw std::runtime_error("Arithmetically-incompatible types");
		}

		[[noreturn]] A mod(A, B) {
			throw std::runtime_error("Arithmetically-incompatible types");
		}
	};


	template<typename A, typename B> requires Comparable<A, B>
	bool are_equal(A a, B b) {
		return a == b;
	}

	template<typename A, typename B>
	[[noreturn]] bool are_equal(A, B) {
		throw std::runtime_error("Non-comparable types");
	}

	template<typename A, typename B>
	struct comparator {
		bool more_than(A a, B b) requires Orderable<A, B> {
			return a > b;
		}

		bool less_than(A a, B b) requires Orderable<A, B> {
			return a < b;
		}

		bool more_equal_than(A a, B b) requires Orderable<A, B> {
			return a >= b;
		}

		bool less_equal_than(A a, B b) requires Orderable<A, B> {
			return a <= b;
		}

		bool more_than(A, B) {
			throw std::runtime_error("Non-orderable types");
		}

		bool less_than(A, B) {
			throw std::runtime_error("Non-orderable types");
		}

		bool more_equal_than(A, B) {
			throw std::runtime_error("Non-orderable types");
		}

		bool less_equal_than(A, B) {
			throw std::runtime_error("Non-orderable types");
		}
	};


	template<typename T>
	struct lambda_function : public interpreter_function {
		std::function<T(const std::vector<runtime_value>&)> f;

		explicit lambda_function(std::function<T(const std::vector<runtime_value>&)> f) : f(f) {

		}

		virtual std::optional<runtime_value> call(const std::vector<runtime_value>& args, base_interpreter&) {
			return f(args);
		}
	};


	template<>
	struct lambda_function<void> : public interpreter_function {
		std::function<void(const std::vector<runtime_value>&)> f;

		explicit lambda_function(std::function<void(const std::vector<runtime_value>&)> f) : f(f) {

		}

		std::optional<runtime_value> call(const std::vector<runtime_value>& args, base_interpreter&) override {
			f(args);

			return std::nullopt;
		}
	};


	struct basic_function : public interpreter_function {
		const function_declaration* declaration;

		explicit basic_function(const function_declaration* declaration) : declaration(declaration) {

		}

		std::optional<runtime_value> call(const std::vector<runtime_value>& args, base_interpreter& i) override {
			if (args.size() != declaration->parameters.size()) {
				throw interpreter_error("Wrong number of arguments for function call", declaration->t);
			}

			i.push_function_scope();
			i.push_variable_scope();
			for (auto a_id = 1ULL; a_id < args.size(); a_id++) {
				auto param = declaration->parameters[a_id];
				auto arg = args[a_id];

				//TODO: Check argument type
				i.inject_value(param.name, arg);
			}

			i.execute_statement(*declaration->body);

			std::optional<runtime_value> return_value;
			auto& func_scope = i.current_function_scope();
			if (func_scope.returned) {
				return_value = func_scope.return_value;
			}

			i.pop_variable_scope();
			i.pop_function_scope();

			//TODO: Check type of RVL?
			return return_value;
		}
	};

	export template<statement_iterator T, std::sentinel_for<T> E>
	class interpreter : public base_interpreter {
		T current;
		E end;
		std::vector<interpreter_variable_scope> scopes;
		std::unordered_map<std::string, std::unique_ptr<interpreter_function>> functions;
		std::vector<interpreter_function_scope> function_scopes;
	public:
		constexpr std::string_view get_type_for_value(const runtime_value& v) {
			auto id = v.index();
			switch (id) {
				case 0:
					return "string";
				case 1:
					return "i8";
				case 2:
					return "i16";
				case 3:
					return "i32";
				case 4:
					return "i64";
				case 5:
					return "u8";
				case 6:
					return "u16";
				case 7:
					return "u32";
				case 8:
					return "u64";
				case 9:
					return "f32";
				case 10:
					return "f64";
				case 11:
					return "boolean";
				default:
					return "i64";
			}
		}

		interpreter(T start, E end) : current(start), end(end) {
			scopes.emplace_back();
		}

		void add_stdlib() {
			functions.try_emplace("print", std::make_unique<lambda_function<void>>([](const auto& args) {
				if (args.size() != 1) {
					throw std::runtime_error("No arguments to print()");
				}

				const auto& arg0 = args[0];

				if (arg0.index() != 0) {
					throw std::runtime_error("Invalid argument to print()");
				}

				const auto& str = std::get<0>(arg0);

				std::cout << str << std::endl;
			}));

			functions.insert_or_assign("boolean_to_string", std::make_unique<lambda_function<std::string>>([](const auto& args) {
				if (args.size() != 1) {
					throw std::runtime_error("No arguments to boolean_to_string()");
				}

				const auto& arg0 = args[0];

				if (arg0.index() != 11) {
					throw std::runtime_error("Invalid argument to boolean_to_string()");
				}

				bool b = std::get<11>(arg0);

				return b ? "true" : "false";
			}));
		}

		void run() {
			while (current != end) {
				execute_statement(**current);
				++current;
			}
		}

		static runtime_value parse_literal(const expression& e) {
			auto* lv = dynamic_cast<const literal_value*>(&e);

			if (lv == nullptr) {
				throw interpreter_error("Internal interpreter error", e.t);
			}

			auto& content = lv->t.content;

			const auto* begin = content.data();
			const auto* end = begin + content.size();
			switch (lv->t.type) {
				case Number:
					if (content.find('.') != std::string::npos) {
						double d{};

						auto v = std::from_chars(begin, end, d, std::chars_format::scientific);

						if (v.ec == std::errc()) {
							return d;
						} else {
							throw interpreter_error("Failed to parse float/double value", lv->t);
						}
					} else {
						uint64_t i{};

						auto v = std::from_chars(begin, end, i);

						if (v.ec == std::errc()) {
							return i;
						} else {
							throw interpreter_error("Failed to parse integer value", lv->t);
						}
					}
					break;
				case String:
					return content;
					break;
				case Boolean:
					if (content == "true") {
						return true;
					} else if (content == "false") {
						return false;
					} else {
						throw interpreter_error("Invalid boolean value", lv->t);
					}
					break;
				default:
					throw interpreter_error("Invalid literal value", lv->t);
			}
		}

		static runtime_value negate(runtime_value v) {
			if (v.index() == 0 || v.index() == 11) {
				throw std::runtime_error("Invalid literal value");
			}

			return std::visit(overload {
				[](std::string&) -> runtime_value { throw std::runtime_error("Invalid negation value"); },
				[](bool) -> runtime_value { throw std::runtime_error("Invalid boolean value"); },
				[]<typename X>(X a) -> runtime_value { return -a; }
			}, v);
		}

		bool logical_or(const dual_operation* dvo) {
			auto v0 = eval_expression(*dvo->first_operand);

			if (v0->index() != 11) {
				throw std::runtime_error("Invalid operand type");
			}

			bool b0 = std::get<11>(*v0);

			if (!b0) {
				auto v1 = eval_expression(*dvo->second_operand);

				if (v1->index() != 11) {
					throw std::runtime_error("Invalid operand type");
				}

				bool b1 = std::get<11>(*v1);

				return b1;
			}

			return true;
		}

		bool logical_and(const dual_operation* dvo) {
			auto v0 = eval_expression(*dvo->first_operand);

			if (!v0.has_value()) {
				throw interpreter_error("Invalid operand value", dvo->first_operand->t);
			}

			if (v0->index() != 11) {
				throw interpreter_error("Invalid operand type", dvo->first_operand->t);
			}

			bool b0 = std::get<11>(*v0);

			if (b0) {
				auto v1 = eval_expression(*dvo->second_operand);

				if (!v1.has_value()) {
					throw interpreter_error("Invalid operand value", dvo->second_operand->t);
				}

				if (v1->index() != 11) {
					throw interpreter_error("Invalid operand type", dvo->second_operand->t);
				}

				bool b1 = std::get<11>(*v1);

				return b1;
			}

			return false;
		}

		runtime_value arithmetic_op(const dual_operation* dvo) {
			auto v0 = eval_expression(*dvo->first_operand);
			auto v1 = eval_expression(*dvo->second_operand);

			if (!v0.has_value() || !v1.has_value()) {
				throw interpreter_error("Invalid operand value", dvo->t);
			}

			return std::visit([&dvo]<typename A, typename B>(A&& a, B&& b) {
				if (dvo->operator_token.content == "+") {
					return runtime_value(arithmetic_evaluator<A, B>().sum(a, b));
				} else if (dvo->operator_token.content == "-") {
					return runtime_value(arithmetic_evaluator<A, B>().sub(a, b));
				} else if (dvo->operator_token.content == "*") {
					return runtime_value(arithmetic_evaluator<A, B>().mul(a, b));
				} else if (dvo->operator_token.content == "/") {
					return runtime_value(arithmetic_evaluator<A, B>().div(a, b));
				} else if (dvo->operator_token.content == "%") {
					return runtime_value(arithmetic_evaluator<A, B>().mod(a, b));
				} else {
					throw interpreter_error("Invalid arithmetic operand", dvo->operator_token);
				}
			}, v0.value(), v1.value());
		}

		runtime_value compare_op(const dual_operation* dvo) {
			auto v0 = eval_expression(*dvo->first_operand);
			auto v1 = eval_expression(*dvo->second_operand);

			if (!v0.has_value() || !v1.has_value()) {
				throw interpreter_error("Invalid operand value", dvo->t);
			}

			return std::visit([&dvo]<typename A, typename B>(A&& a, B&& b) {
				if (dvo->operator_token.content == "==") {
					return are_equal(a, b);
				} else if (dvo->operator_token.content == "!=") {
					return !are_equal(a, b);
				} else {
					throw interpreter_error("Invalid comparator operator", dvo->operator_token);
				}
			}, v0.value(), v1.value());
		}

		runtime_value order_op(const dual_operation* dvo) {
			auto v0 = eval_expression(*dvo->first_operand);
			auto v1 = eval_expression(*dvo->second_operand);

			if (!v0.has_value() || !v1.has_value()) {
				throw interpreter_error("Invalid operand value", dvo->t);
			}

			return std::visit([&dvo]<typename A, typename B>(A&& a, B&& b) {
				if (dvo->operator_token.content == ">") {
					return comparator<A, B>().more_than(a, b);
				} else if (dvo->operator_token.content == "<") {
					return comparator<A, B>().less_than(a, b);
				} else if (dvo->operator_token.content == ">=") {
					return comparator<A, B>().more_equal_than(a, b);
				} else if (dvo->operator_token.content == "<=") {
					return comparator<A, B>().less_equal_than(a, b);
				} else {
					throw interpreter_error("Invalid comparator operator", dvo->operator_token);
				}
			}, v0.value(), v1.value());
		}

		std::optional<runtime_value> eval_expression(const expression& e) {
			switch (e.et) {
				case EtInvalid:
					throw interpreter_error("Invalid expression", e.t);
				case EtLiteralValue:
					return parse_literal(e);
				case EtSingleValueOperation: {
						auto* svo = dynamic_cast<const single_operation*>(&e);

						if (svo == nullptr) {
							throw interpreter_error("Internal interpreter error", e.t);
						}

						auto val = eval_expression(*svo->operand);

						if (!val.has_value()) {
							throw interpreter_error("Invalid operand value", svo->operand->t);
						}

						if (svo->operator_token.content == "+") {
							return val;
						} else if (svo->operator_token.content == "-") {
							return negate (val.value());
						} else if (svo->operator_token.content == "!") {
							if (val->index() != 11) {
								throw interpreter_error("Invalid operand type", svo->operand->t);
							} else {
								return !std::get<11>(val.value());
							}
						} else {
							throw interpreter_error("Invalid operator", svo->operator_token);
						}
					}
					break;
				case EtDualValueOperation: {
						auto* dvo = dynamic_cast<const dual_operation*>(&e);

						if (dvo == nullptr) {
							throw interpreter_error("Internal interpreter error", e.t);
						}

						if (dvo->operator_token.content == "||") {
							return logical_or(dvo);
						}

						if (dvo->operator_token.content == "&&") {
							return logical_and(dvo);
						}

						if (dvo->operator_token.content == "+" || dvo->operator_token.content == "-"
							|| dvo->operator_token.content == "*" || dvo->operator_token.content == "/"
							|| dvo->operator_token.content == "%") {
							return arithmetic_op(dvo);
						}

						if (dvo->operator_token.content == "==" || dvo->operator_token.content == "!=") {
							return compare_op(dvo);
						}

						if (dvo->operator_token.content == ">" || dvo->operator_token.content == "<" ||
							dvo->operator_token.content == "<=" || dvo->operator_token.content == ">=") {
							return order_op(dvo);
						}

						throw interpreter_error("Invalid operator", dvo->operator_token);
					}
					break;
				case EtFunctionCall: {
						auto* call = dynamic_cast<const function_call*>(&e);

						if (call == nullptr) {
							throw interpreter_error("Internal interpreter error", e.t);
						}

						auto it = functions.find(call->function_name);

						if (it == functions.end()) {
							throw interpreter_error(std::format("Function {} not found", call->function_name), call->t);
						}

						std::vector<runtime_value> args;

						for (auto& a : call->arguments) {
							auto arg = eval_expression(*a);

							if (!arg.has_value()) {
								throw interpreter_error("Internal interpreter error", a->t);
							}
							args.emplace_back(arg.value());
						}

						return it->second->call(args, *this);
					}
					break;
				case EtIdentifier: {
						auto* id = dynamic_cast<const identifier_expression*>(&e);

						if (id == nullptr) {
							throw interpreter_error("Internal interpreter error", e.t);
						}

						for (auto& s : std::ranges::reverse_view(scopes)) {
							if (s.variables.contains(id->name)) {
								auto& v = s.variables.at(id->name).value;
								return v;
							}
						}

						throw interpreter_error("Variable not found", e.t);
					}
					break;
			}
		}

		std::optional<runtime_value> execute_statement(const statement& s) override {
			switch (s.type) {
				case Expression: {
						auto* ex = dynamic_cast<const expression*>(&s);

						if (ex == nullptr) {
							throw interpreter_error("Internal interpreter error", s.t);
						}

						return eval_expression(*ex);
					}
					break;
				case Block: {
						auto* b = dynamic_cast<const block*>(&s);

						if (b == nullptr) {
							throw interpreter_error("Internal interpreter error", s.t);
						}

						push_variable_scope();

						for (auto& st : b->statements) {
							execute_statement(*st);

							if (is_in_function()) {
								auto& sc = current_function_scope();
								if (sc.returned) {
									break;
								}
							}
						}

						pop_variable_scope();

						return std::nullopt;
					}
					break;
				case ReturnStatement: {
						auto* r = dynamic_cast<const return_statement*>(&s);

						if (r == nullptr) {
							throw interpreter_error("Internal interpreter error", s.t);
						}

						auto val = r->expression != nullptr ? eval_expression(*r->expression) : std::nullopt;

						auto& sc = current_function_scope();
						sc.return_value = val;
						sc.returned = true;
					}
					break;
				case Root: {
						auto* r = dynamic_cast<const root_statement*>(&s);

						if (r == nullptr) {
							throw interpreter_error("Internal interpreter error", s.t);
						}

						for (auto& st : r->statements) {
							execute_statement(*st);
						}
					}
					break;
				case FunctionDeclaration: {
						auto* d = dynamic_cast<const function_declaration*>(&s);

						if (d == nullptr) {
							throw interpreter_error("Internal interpreter error", s.t);
						}

						if (functions.find(d->name) != functions.end()) {
							throw interpreter_error("Duplicate function name", s.t);
						}

						functions.emplace(d->name, std::make_unique<basic_function>(d));

						return std::nullopt;
					}
					break;
				case VariableDeclaration: {
						auto* d = dynamic_cast<const variable_declaration*>(&s);

						if (d == nullptr) {
							throw interpreter_error("Internal interpreter error", s.t);
						}

						auto& sc = current_variable_scope();

						if (sc.variables.find(d->name) != sc.variables.end()) {
							throw interpreter_error("Duplicate variable name", s.t);
						}

						auto val = d->initial_expression != nullptr ? eval_expression(*d->initial_expression) : throw std::runtime_error("Non-initialized variables are not supported yet."); //TODO: This

						sc.variables.insert_or_assign(d->name, interpreter_variable {
							d->type.has_value() ? d->type.value() : "void",
							!d->constant,
							val.has_value() ? val.value() : runtime_value(0),
						});

						return std::nullopt;
					}
					break;
				case IfStatement: {
					auto* is = dynamic_cast<const if_statement*>(&s);

					if (is == nullptr) {
						throw interpreter_error("Internal interpreter error", s.t);
					}

					auto condition = eval_expression(*is->condition);

					if (!condition.has_value()) {
						throw interpreter_error("No condition for if statement", is->condition->t);
					}

					if (condition->index() != 11) {
						throw interpreter_error("If condition is not boolean", is->condition->t);
					}

					if (std::get<11>(condition.value())) {
						execute_statement(*is->block);
					}
				} break;
			}
			throw interpreter_error("Internal interpreter error", s.t);
		}

		void push_variable_scope() override {
			scopes.emplace_back();
		}

		void pop_variable_scope() override {
			scopes.pop_back();
		}

		void push_function_scope() override {
			function_scopes.emplace_back();
		}

		void pop_function_scope() override {
			function_scopes.pop_back();
		}

		void inject_value(const std::string_view& name, runtime_value val) override {
			auto& s = scopes.back();

			auto type = get_type_for_value(val);

			if (s.variables.contains(std::string(name))) {
				throw std::runtime_error("Injection failed: Variable already exists");
			}

			s.variables.insert_or_assign(std::string(name), interpreter_variable {
				std::string(type),
				false,
				std::move(val)
			});
		}

		interpreter_function_scope& current_function_scope() override {
			if (function_scopes.empty()) {
				throw std::runtime_error("Internal interpreter error");
			}

			return function_scopes.back();
		}

		interpreter_variable_scope& current_variable_scope() override {
			return scopes.back();
		}

		bool is_in_function() override {
			return !function_scopes.empty();
		}
	};
}