#include <iostream>
#include <iomanip>
#include <algorithm>
#include <sstream>
#include <cmath>
#include <chrono>
#include <fstream>
#include <span>
#include <cstring>
#include "UserScript.h"

void print_value(std::ostream& stream, const scripting::script_value& res) {
	if(std::holds_alternative<scripting::array>(res)) {
		stream << "[";
		auto max = std::get<scripting::array>(res).value.size();
		auto no_comma = max - 1;
		for(size_t idx = 0; idx < max; ++idx) {
			print_value(stream, std::get<scripting::array>(res).value[idx]);
			stream << (idx != no_comma ? ", " : "");
		}
		stream << "]";
	} else if(std::holds_alternative<std::string>(res)) {
		stream << std::get<std::string>(res);
	} else if(std::holds_alternative<scripting::null>(res)) {
		stream << "null";
	} else {
		stream << std::get<int32_t>(res);
	}
}

struct identity : public scripting::function_impl {
	std::optional<scripting::script_value> apply(scripting::UserScript* self,std::vector<scripting::argument> args, std::optional<scripting::script_error>& errors) final {
		if(args.size() != 1) {
			errors = scripting::script_error{.message = "identity expects a single argument"};
		} else {
			if(std::holds_alternative<scripting::script_value>(args.front())) {
				return std::get<scripting::script_value>(args.front());
			} else {
				return self->resolve(std::get<scripting::script_variable>(args.front()).name);
			}
		}
		return scripting::script_value({});
	}
};

struct print : public scripting::function_impl {
	std::ostream& stream;
	
	print(std::ostream& _stream) : stream(_stream) {}
	
	std::optional<scripting::script_value> apply(scripting::UserScript* self,std::vector<scripting::argument> args, std::optional<scripting::script_error>& errors) final {
		
		while(not args.empty()) {
			auto& arg = args.back();
			if(std::holds_alternative<scripting::script_value>(arg)) {
				print_value(stream, std::get<scripting::script_value>(arg));
			} else {
				print_value(stream, self->resolve(std::get<scripting::script_variable>(arg).name));
			}
			args.pop_back();
		}
		return scripting::script_value({});
	}
};

struct set : public scripting::function_impl {
	std::optional<scripting::script_value> apply(scripting::UserScript* self,std::vector<scripting::argument> args, std::optional<scripting::script_error>& errors) final {
		
		if(args.size() != 2) {
			errors = scripting::script_error{
				.message = "set expects 2 arguments"
			};
			return scripting::script_value{};
		}
		
		auto& var = args.back();
		if(not holds_alternative<scripting::script_variable>(var)) {
			errors = scripting::script_error{
				.message = "set expects the first argument to be a target variable"
			};
			return scripting::script_value{};
		}
		
		auto& arg = args.front();
		
		if(std::holds_alternative<scripting::script_value>(arg)) {
			self->setValue(get<scripting::script_variable>(var).name, std::get<scripting::script_value>(arg));
		} else {
			self->setValue(get<scripting::script_variable>(var).name, self->resolve(std::get<scripting::script_variable>(arg).name));
		}
		if(auto v = self->getValue(get<scripting::script_variable>(var).name); v) {
			return v.value();
		} else {
			return scripting::script_value{};
		}
	}
};

struct terminate : public scripting::function_impl {
	std::optional<scripting::script_value> apply(scripting::UserScript*,std::vector<scripting::argument>, std::optional<scripting::script_error>&) final {
		std::exit(1);
		// PLEASE DO NOT ACTUALLY EXIT YOU FUCKING IDIOT
		return scripting::script_value({});
	}
};

void process_bench(std::string target = "./tests/scripts/testfile.test") {
	auto engine = scripting::prepare_interpreter(std::string{});
	
	engine->registerFunction("identity", std::make_unique<identity>());
	engine->registerFunction("exit", std::make_unique<terminate>());
	engine->registerFunction("set", std::make_unique<set>());
	
	/***
	 * This is a half assed benchmark,
	 * Document results here to keep the thingy in check performance wise (release mode only)
	 *
	 * 2023-07-04 Archivist -> 2618ns - 308ns - 49ns (clang+libstdc++)
	 * 2023-07-07 Archivist -> 2481ns - 291ns - 46ns (clang+libc++)
	 * 2023-07-07 Archivist -> 106ns - 12ns - 2ns (clang+march=native+libc++)
	 */
	engine->registerFunction("print", std::make_unique<print>(std::cout));
	std::ifstream src_str(target);
	std::stringstream code;
	code << src_str.rdbuf();
	int steps = 0;
	
	decltype(std::chrono::high_resolution_clock::now()-std::chrono::high_resolution_clock::now()) per_exec{}, per_step{}, per_op{};
	
	for(int runs = 0; runs < 5000; runs++) {
		
		auto res = engine->prepare(code.str());
		
		auto begin = std::chrono::high_resolution_clock::now();
		while (not engine->getValue("exit_ctr").has_value()) {
			engine->stepOnce();
			steps++;
		}
		auto end = std::chrono::high_resolution_clock::now();
		per_exec += (end - begin);
		per_step += (end - begin);
		per_op += (end - begin);
	}
	per_exec /= 5000;
	per_step /= steps;
	per_op = per_op / 5000 / 53;
	
	std::cout << "time per exec = " << std::chrono::duration_cast<std::chrono::nanoseconds>(per_exec).count() << "ns\n";
	std::cout << "time per step = " << std::chrono::duration_cast<std::chrono::nanoseconds>(per_step).count() << "ns\n";
	std::cout << "time per avg op = " << std::chrono::duration_cast<std::chrono::nanoseconds>(per_op).count() << "ns\n";
}

void compile_bench(std::string target = "./tests/scripts/testfile.test") {
	auto engine = scripting::prepare_interpreter(std::string{});
	
	engine->registerFunction("identity", std::make_unique<identity>());
	engine->registerFunction("exit", std::make_unique<terminate>());
	engine->registerFunction("set", std::make_unique<set>());
	
	/***
	 * Same as above but for compilation times
	 *
	 * 2023-07-04 Archivist -> 386µs
	 */
	engine->registerFunction("print", std::make_unique<print>(std::cout));
	std::ifstream src_str("./tests/scripts/testfile.test");
	std::stringstream code;
	code << src_str.rdbuf();
	
	auto begin = std::chrono::high_resolution_clock::now();
	
	
	
	[&]() __attribute__((optimize("O0"))) {
		auto res = engine->prepare(code.str());
		res = engine->prepare(code.str());
		res = engine->prepare(code.str());
		res = engine->prepare(code.str());
		res = engine->prepare(code.str());
	}();
	
	auto end = std::chrono::high_resolution_clock::now();
	auto per_exec = (end - begin)/5;
	std::cout << "time per exec = " << std::chrono::duration_cast<std::chrono::microseconds>(per_exec).count() << "µs\n";
}

void compare(std::string target, std::string expect) {
	auto engine = scripting::prepare_interpreter(std::string{});
	
	engine->registerFunction("identity", std::make_unique<identity>());
	engine->registerFunction("exit", std::make_unique<terminate>());
	engine->registerFunction("set", std::make_unique<set>());
	
	
	std::stringstream str;
	std::string_view filename_source = target;
	std::string_view filename_output = expect;
	engine->registerFunction("print", std::make_unique<print>(str));
	
	std::ifstream src_str(std::string{filename_source});
	std::stringstream code;
	code << src_str.rdbuf();
	
	std::ifstream out_str(std::string{filename_output});
	std::stringstream output;
	output << out_str.rdbuf();
	
	auto res = engine->executeAtOnce(code.str());
	if (std::holds_alternative<scripting::script_value>(res)) {
	} else {
		auto &errors = std::get<std::vector<scripting::script_error>>(res);
		for (auto &line: errors) {
			str << line.message << "\n    at line " << line.location->line_number << ":"
			    << line.location->column_number << "\n";
			str << "    " << *line.location->line_contents << "\n";
			str << "    " << std::string(line.location->column_number - 1, ' ') << "^\n";
		}
	}
	
	int status = 0;
	
	while(not output.eof()) {
		std::string expected, found;
		std::getline(output, expected);
		std::getline(str, found);
		bool ok = (expected != found);
		status+= ok ;
		(ok ? std::cerr : std::cout)
			<< (not ok ? "\033[21;32m" : "\033[1;31m") << expected
			<< std::string(std::max<size_t>(0, 40 - expected.size()), ' ')<< "| " << found << std::endl;
	}
	if(status) std::exit(status);
}

void immediate_interactive() {
	auto engine = scripting::prepare_interpreter(std::string{});
	
	engine->registerFunction("identity", std::make_unique<identity>());
	engine->registerFunction("exit", std::make_unique<terminate>());
	engine->registerFunction("set", std::make_unique<set>());
	
	engine->registerFunction("print", std::make_unique<print>(std::cout));
	bool exit = false;
	while (not exit) {
		std::string code;
		std::getline(std::cin, code);
		auto res = engine->executeAtOnce(code);
		if (std::holds_alternative<scripting::script_value>(res)) {
		} else {
			auto &errors = std::get<std::vector<scripting::script_error>>(res);
			for (auto &line: errors) {
				std::cout << line.message << "\n    at line ";
				if(line.location) {
					std::cout << line.location->line_number << ":"
					          << line.location->column_number << "\n";
					std::cout << "    " << *line.location->line_contents << "\n";
					std::cout << "    " << std::string(line.location->column_number - 1, ' ') << "^\n";
				} else std::cout << "UNKNOWN\n";
			}
		}
	}
}

void exec(std::span<std::string_view> args) {
	std::vector<decltype(scripting::prepare_interpreter(std::string{}))> batch;
	
	auto engine = scripting::prepare_interpreter(std::string{});
	
	engine->registerFunction("identity", std::make_unique<identity>());
	engine->registerFunction("terminate", std::make_unique<terminate>());
	engine->registerFunction("set", std::make_unique<set>());
	
	engine->registerFunction("print", std::make_unique<print>(std::cout));
	bool exit = false;
	while (not exit) {
		std::string code;
		std::getline(std::cin, code);
		auto res = engine->executeAtOnce(code);
		if (std::holds_alternative<scripting::script_value>(res)) {
		} else {
			auto &errors = std::get<std::vector<scripting::script_error>>(res);
			for (auto &line: errors) {
				std::cout << line.message << "\n    at line ";
				if(line.location) {
					std::cout << line.location->line_number << ":"
					          << line.location->column_number << "\n";
					std::cout << "    " << *line.location->line_contents << "\n";
					std::cout << "    " << std::string(line.location->column_number - 1, ' ') << "^\n";
				} else std::cout << "UNKNOWN\n";
			}
		}
	}
}

#if defined(__linux__) or defined(WIN32)
constexpr bool trim_first_argument = true;
#else
constexpr bool trim_first_argument = false;
static_assert(false, "Undefined status of the first argument");
#endif

int cpp_main(std::span<std::string_view> args) {
	if constexpr (trim_first_argument) {
		args = args.subspan(1);
	}
	
	if(args.empty() || args.front() == "immediate") {
		immediate_interactive();
		std::exit(0);
	} else if(args.front() == "compare") {
		args = args.subspan(1);
		if(args.size() != 2) {
			std::cerr << "compare expects 2 files as arguments" << std::endl;
			std::terminate();
		}
	} else if(args.front() == "bench_exec") {
		args = args.subspan(1);
		if(args.size() > 1) {
			std::cerr << "bench_exec expects 0 or 1 file as arguments" << std::endl;
			std::terminate();
		}
		if(args.empty()) process_bench();
		else process_bench(std::string{args.front()});
	} else if(args.front() == "bench_compile") {
		args = args.subspan(1);
		if(args.size() > 1) {
			std::cerr << "bench_compile expects 0 or 1 file as arguments" << std::endl;
			std::terminate();
		}
		if(args.empty()) compile_bench();
		else compile_bench(std::string{args.front()});
	} else if(args.front() == "exec") {
		// exec(args.subspan(1));
	} else {
		std::cerr << "Unknown option" << std::endl;
	}
	return 0;
}

int main(int argc, char** argv) {
	std::vector<std::string_view> args;
	for(auto& arg : std::span(argv, argv+argc)) {
		args.emplace_back(arg, arg+strlen(arg));
	}
	return cpp_main(args);
}