#include "../vm.hpp"
#include <fstream>
#include <filesystem>
#include <algorithm>
#include "../helpers.hpp"

std::string DlangVM::GetString(const ByteCodeArray& bytecode, size_t& ip)
{
	auto idx = bytecode[ip++];
	return strings[idx];
}

std::string DlangVM::GetRealString(const ByteCodeArray& bytecode, size_t& ip)
{
	double length = bytecode[ip++];
	std::string str = "";
	for (double i = 0; i < length; ++i)
		str += static_cast<char>(bytecode[ip++]);
	return str;
}

void DlangVM::Execute(ByteCodeArray bytecode)
{
	size_t ip = 0;
	auto bytecodeSize = bytecode.size();

#if DEBUG_MODE == 1
	printf("==================== VM ====================\n");
	auto start = std::chrono::high_resolution_clock::now();
#endif

	while (ip < bytecodeSize)
	{
		auto op = static_cast<Opcodes>(bytecode[ip++]);
		switch (op)
		{
			#pragma region VM CONTROL OPCODES
			case Opcodes::EOC: break; break;
			case Opcodes::POP: Pop(); break;
			
			case Opcodes::PUSH_INT: {
				auto val = bytecode[ip++];
				PushNumber(val);
			} break;

			case Opcodes::PUSH_BOOL: {
				auto val = bytecode[ip++];
				PushBoolean(val == 1);
			} break;

			case Opcodes::PUSH_STRING: {
				auto val = GetString(bytecode, ip);
				PushString(val);
			} break;

			case Opcodes::PUSH_VAR: {
				auto val = GetString(bytecode, ip);
				PushVar(val);
			} break;

			case Opcodes::PUSH_FUNC: {
				auto func = GetString(bytecode, ip);
				stack.push(
					globals.Get(func)
				);
			} break;

			case Opcodes::DEFINE_VAR:
			case Opcodes::STORE_VAR: {
				auto name = GetString(bytecode, ip);
				auto val = Pop();

				SetVariable(name, val);
			} break;

			case Opcodes::NEW_TABLE: PushTable(std::make_shared<DlangTable>()); break;
			
			case Opcodes::SET_FIELD: {
				std::string key = GetString(bytecode, ip);
				SetField(key);
			} break;
			
			case Opcodes::GET_FIELD: {
				std::string key = GetString(bytecode, ip);
				stack.push(
					GetField(key)
				);
			} break;
			
			case Opcodes::KEYS: stack.push(GetTableKeys()); break;

			case Opcodes::NEW_LIST: PushList(std::make_shared<DlangList>()); break;
			case Opcodes::PUSH_LIST: PushListIndex(); break;
			case Opcodes::SET_INDEX: SetListIndex(); break;
			case Opcodes::GET_INDEX: stack.push(GetListIndex()); break;
				
			case Opcodes::SIZE: PushNumber(GetSize()); break;
#pragma endregion

			#pragma region VM MATH OPCODES
			case Opcodes::ADD: {
				auto b = stack.pop();
				auto a = stack.pop();

				if ((a.type != StackValueType::Int || b.type != StackValueType::Int) &&
					(a.type != StackValueType::String && b.type != StackValueType::String))
					throw std::runtime_error("ADD Operation error, not all variables are integer or strings!");

				bool fullInt = (a.type == StackValueType::Int && b.type == StackValueType::Int);

				if (fullInt)
					stack.push(StackValue(StackValueType::Int,
						std::get<double>(a.data) +
						std::get<double>(b.data)
					));
				else
				{
					std::string content;
					if (a.type == StackValueType::Int)
						content += DoubleToString(std::get<double>(a.data));
					else
						content += std::get<std::string>(a.data);

					if (b.type == StackValueType::Int)
						content += DoubleToString(std::get<double>(b.data));
					else
						content += std::get<std::string>(b.data);

					stack.push(StackValue(StackValueType::String, content));
				}

			} break;

			case Opcodes::SUB: {
				auto b = stack.pop();
				auto a = stack.pop();

				if (a.type != StackValueType::Int || b.type != StackValueType::Int)
					throw std::runtime_error("SUB Operation error, not all variables are integer!");

				stack.push(StackValue(StackValueType::Int,
					std::get<double>(a.data) -
					std::get<double>(b.data)
				));
			} break;

			case Opcodes::MUL: {
				auto b = stack.pop();
				auto a = stack.pop();

				if (a.type != StackValueType::Int || b.type != StackValueType::Int)
					throw std::runtime_error("ADD Operation error, not all variables are integer!");

				stack.push(StackValue(StackValueType::Int,
					std::get<double>(a.data) *
					std::get<double>(b.data)
				));
			} break;

			case Opcodes::DIV: {
				auto b = stack.pop();
				auto a = stack.pop();

				if (a.type != StackValueType::Int || b.type != StackValueType::Int)
					throw std::runtime_error("ADD Operation error, not all variables are integer!");

				stack.push(StackValue(StackValueType::Int,
					std::get<double>(a.data) /
					std::get<double>(b.data)
				));
			} break;

			case Opcodes::GREATER_THAN: {
				auto b = stack.pop();
				auto a = stack.pop();

				if (a.type != StackValueType::Int || b.type != StackValueType::Int)
					throw std::runtime_error("GREATHER_THAN Operation error, not all variables are integer!");

				stack.push(StackValue(StackValueType::Bool,
					std::get<double>(a.data) >
					std::get<double>(b.data)
				));
			} break;

			case Opcodes::LESS_THAN: {
				auto b = stack.pop();
				auto a = stack.pop();

				if (a.type != StackValueType::Int || b.type != StackValueType::Int)
					throw std::runtime_error("LESS_THAN Operation error, not all variables are integer!");

				stack.push(StackValue(StackValueType::Bool,
					std::get<double>(a.data) <
					std::get<double>(b.data)
				));
			} break;

			case Opcodes::CMP: {
				auto b = stack.pop();
				auto a = stack.pop();

				if (a.type != StackValueType::Int || b.type != StackValueType::Int)
					throw std::runtime_error("LESS_THAN Operation error, not all variables are integer!");

				stack.push(StackValue(StackValueType::Bool,
					std::get<double>(a.data) ==
					std::get<double>(b.data)
				));
			} break;

			case Opcodes::GTE: {
				auto b = stack.pop();
				auto a = stack.pop();

				StackValue result;
				result.type = StackValueType::Bool;

				if (a.type == StackValueType::Int && b.type == StackValueType::Int) {
					result.data = (std::get<double>(a.data) >= std::get<double>(b.data));
				}

				stack.push(result);
			} break;

			case Opcodes::LTE: {
				auto b = stack.pop();
				auto a = stack.pop();

				StackValue result;
				result.type = StackValueType::Bool;

				if (a.type == StackValueType::Int && b.type == StackValueType::Int) {
					result.data = (std::get<double>(a.data) <= std::get<double>(b.data));
				}

				stack.push(result);
			} break;

			case Opcodes::AND: {
				auto b = stack.pop();
				auto a = stack.pop();

				StackValue result;
				result.type = StackValueType::Bool;

				if (a.type == StackValueType::Bool && b.type == StackValueType::Bool) {
					bool lVal = std::get<bool>(a.data);
					bool rVal = std::get<bool>(b.data);
					result.data = (lVal && rVal);
				}
				else {
					throw std::runtime_error("Logical AND (&&) requires boolean operands!");
				}

				stack.push(result);
			} break;

			case Opcodes::OR: {
				auto b = stack.pop();
				auto a = stack.pop();

				StackValue result;
				result.type = StackValueType::Bool;

				if (a.type == StackValueType::Bool && b.type == StackValueType::Bool) {
					bool lVal = std::get<bool>(a.data);
					bool rVal = std::get<bool>(b.data);
					result.data = (lVal || rVal);
				}
				else {
					throw std::runtime_error("Logical OR (||) requires boolean operands!");
				}

				stack.push(result);
			} break;

			case Opcodes::NTE: {
				auto b = stack.pop();
				auto a = stack.pop();

				StackValue result;
				result.type = StackValueType::Bool;

				if (a.type == StackValueType::Int && b.type == StackValueType::Int) {
					result.data = (std::get<double>(a.data) != std::get<double>(b.data));
				}

				stack.push(result);
			} break;
			#pragma endregion

			#pragma region LOGIC OPCODES
			case Opcodes::CALL: {
				auto argCount = bytecode[ip++];

				auto func = stack.pop();
				if (func.type != StackValueType::Function)
					throw std::runtime_error("Ohno, missing a function on the stack for call!");

				auto function = std::get<std::shared_ptr<DlangFunction>>(func.data);

				if (function->isNative)
					function->nativeFunction(this);
				else {
					CallFrame newFrame;
					newFrame.returnIp = ip;
					newFrame.previousBytecode = bytecode;

					for (double i = (argCount - 1); i >= 0; --i)
						newFrame.locals.Set(function->parameters[i], stack.pop());

					_callStack.push(newFrame);

					bytecode = function->bytecodeChunk;
					ip = 0;
				}
			} break;

			case Opcodes::RET: {
				if (_callStack.top() <= 0)
					return; // End because return was called in global file

				StackValue returnValue;
				bool assigned = false;
				if (stack.top() > 0)
				{
					returnValue = stack.pop();
					assigned = true;
				}

				CallFrame oldFrame = _callStack.pop();

				ip = oldFrame.returnIp;
				bytecode = oldFrame.previousBytecode;

				if (assigned)
					stack.push(returnValue);

			} break;

			case Opcodes::FUNC_DEF: {
				auto name = GetString(bytecode, ip);
				auto paramCount = bytecode[ip++];

				std::vector<std::string> parameters;
				for (double i = 0; i < paramCount; ++i)
					parameters.push_back(GetString(bytecode, ip));

				double bodyStart = bytecode[ip++];
				double bodyEnd = bytecode[ip++];

				std::vector<double> funcChunk(
					bytecode.begin() + static_cast<size_t>(bodyStart),
					bytecode.begin() + static_cast<size_t>(bodyEnd)
				);

				auto funcObj = std::make_shared<DlangFunction>();
				funcObj->isNative = false;
				funcObj->parameters = parameters;
				funcObj->bytecodeChunk = std::move(funcChunk);

				StackValue v;
				v.type = StackValueType::Function;
				v.data = funcObj;

				if (!name.empty())
					globals.Set(name, v);
				else
					stack.push(v);
			} break;

			case Opcodes::JZ: {
				double targetAddress = bytecode[ip++];
				auto val = stack.pop();
				bool isFalse = false;

				std::visit([&isFalse](auto&& arg) {
					using T = std::decay_t<decltype(arg)>;
					if constexpr (std::is_same_v<T, bool>) {
						isFalse = !arg;
					}
					else if constexpr (std::is_same_v<T, double>) {
						isFalse = (arg == 0);
					}
					else throw std::runtime_error("Invalid type for is zero check!");
					}, val.data);

				if (isFalse)
					ip = targetAddress;
			} break;

			case Opcodes::JUMP: {
				double targetAddress = bytecode[ip++];
				ip = targetAddress;
			} break;
			#pragma endregion

			default:
				printf("Invalid opcode 0x%llX!\n", op);
				break;
		}
	}

#if DEBUG_MODE == 1
	auto end = std::chrono::high_resolution_clock::now();
	auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
	auto durationMicroSeconds = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();

	printf("=================== END ====================\n");
	printf("VM took %lld ms (%lld microseconds)\n", duration, durationMicroSeconds);
#endif
}

StackValue DlangVM::Pop()
{
	return stack.pop();
}

void DlangVM::Pop(double amount)
{
	amount = static_cast<int>(amount);
	stack.pop(amount);
}

void DlangVM::PushNumber(double number)
{
	stack.push(
		StackValue(
			StackValueType::Int, number
		)
	);
}

void DlangVM::PushBoolean(bool value)
{
	stack.push(
		StackValue(
			StackValueType::Bool, value
		)
	);
}

void DlangVM::PushString(const std::string& str)
{
	stack.push(
		StackValue(
			StackValueType::String, str
		)
	);
}

void DlangVM::PushVar(const std::string& var)
{
	stack.push(
		GetVariable(var)
	);
}

void DlangVM::PushCFunction(CFunction func)
{
	StackValue s;
	s.type = StackValueType::Function;
	s.data = std::move(DlangFunction::CreateNative(func));
	stack.push(s);
}

bool DlangVM::CallFunc(double argCount)
{
	auto func = stack.pop();
	if (func.type != StackValueType::Function)
		return false;

	auto function = std::get<std::shared_ptr<DlangFunction>>(func.data);
	if (function->isNative)
		function->nativeFunction(this);
	else {

		CallFrame newFrame;

		std::vector<StackValue> providedArgs(argCount);
		for (int i = argCount - 1; i >= 0; --i) {
			providedArgs[i] = stack.pop();
		}

		for (size_t i = 0; i < function->parameters.size(); ++i) {
			std::string paramName = function->parameters[i];

			if (i < argCount) {
				newFrame.locals.Set(paramName, providedArgs[i]);
			}
			else {
				newFrame.locals.Set(paramName, StackValue());
			}
		}

		_callStack.push(newFrame);

		Execute(function->bytecodeChunk);
	}

	return true;
}

void DlangVM::PushTable(std::shared_ptr<DlangTable> table)
{
	stack.push(
		StackValue(
			StackValueType::Table,
			table
		)
	);
}

void DlangVM::SetField(const std::string& key)
{
	auto val = Pop();
	auto stackTable = Pop();

	if (stackTable.type == StackValueType::Table)
	{
		auto table = std::get<std::shared_ptr<DlangTable>>(stackTable.data);
		table->fields[key] = val;
		PushTable(table);
	}
}

StackValue DlangVM::GetField(const std::string& key, bool popTable)
{
	auto stackTable = Pop();

	if (stackTable.type == StackValueType::Table)
	{
		auto table = std::get<std::shared_ptr<DlangTable>>(stackTable.data);
		
		if (!popTable)
			PushTable(table);

		auto it = table->fields.find(key);
		if (it == table->fields.end())
		{
			printf("[VM Error]: Field '%s' is not found!\n", key.c_str());
			return StackValue();
		}

		return it->second;
	}

	printf("Cant get field on non table type %s!\n", StackValueTypeToString(stackTable.type).c_str());
	return StackValue();
}

StackValue DlangVM::GetTableKeys()
{
	auto stackTable = Pop();

	if (stackTable.type == StackValueType::Table)
	{
		auto table = std::get<std::shared_ptr<DlangTable>>(stackTable.data);
		auto keyList = std::make_shared<DlangList>();

		for (const auto& [key, _] : table->fields)
			keyList->values.push_back(
				StackValue(
					StackValueType::String,
					key
				)
			);

		return StackValue(StackValueType::List, std::move(keyList));
	}
	else {
		printf("Failed to keys from non table type %s!\n", StackValueTypeToString(stackTable.type).c_str());
		return StackValue();
	}
}

void DlangVM::PushList(std::shared_ptr<DlangList> list)
{
	stack.push(
		StackValue(
			StackValueType::List,
			list
		)
	);
}

double DlangVM::PushListIndex()
{
	auto val = Pop();
	auto stackList = Pop();

	if (stackList.type == StackValueType::List)
	{
		auto list = std::get<std::shared_ptr<DlangList>>(stackList.data);
		list->Push(val);
		PushList(list);
		return list->values.size() - 1;
	}
	else {
		printf("Failed to push list index to type '%s'!\n", StackValueTypeToString(stackList.type).c_str());
		return 0;
	}
}

void DlangVM::SetListIndex()
{
	auto val = Pop();
	auto stackIndex = Pop();
	auto stackList = Pop();

	if (stackIndex.type != StackValueType::Int)
	{
		printf("Failed to set non number type index %s!\n", StackValueTypeToString(stackIndex.type).c_str());
		stack.push(StackValue());
		return;
	}

	if (stackList.type == StackValueType::List)
	{
		auto list = std::get<std::shared_ptr<DlangList>>(stackList.data);
		auto index = static_cast<int>(std::get<double>(stackIndex.data));
		list->values[index] = val;
		stack.push(stackList);
	}
	else if (stackList.type == StackValueType::String)
	{
		auto& str = std::get<std::string>(stackList.data);
		auto index = static_cast<int>(std::get<double>(stackIndex.data));

		if (index > 0 && index < str.size())
		{
			if (std::holds_alternative<std::string>(val.data)) {
				str[index] = std::get<std::string>(val.data)[0];
			}
			else if (std::holds_alternative<double>(val.data)) {
				str[index] = static_cast<char>(std::get<double>(val.data));
			}
		}

		stack.push(stackList);
	}
	else {
		printf("Failed to set index on non list/string type %s!\n", StackValueTypeToString(stackList.type).c_str());
		stack.push(StackValue());
	}
}

StackValue DlangVM::GetListIndex()
{
	auto stackIndex = Pop();
	auto stackObj = Pop();

	if (stackIndex.type != StackValueType::Int && stackIndex.type != StackValueType::String)
	{
		printf("Failed to set non number/string type index %s!\n", StackValueTypeToString(stackIndex.type).c_str());
		return StackValue();
	}

	if (stackObj.type == StackValueType::List && stackIndex.type == StackValueType::Int)
	{
		auto list = std::get<std::shared_ptr<DlangList>>(stackObj.data);
		auto index = std::get<double>(stackIndex.data);
		return list->values[index];
	} else if (stackObj.type == StackValueType::String && stackIndex.type == StackValueType::Int)
	{
		auto str = std::get<std::string>(stackObj.data);
		auto index = std::get<double>(stackIndex.data);
		return StackValue(
			StackValueType::String,
			str.at(index) + ""
		);
	}
	else if (stackObj.type == StackValueType::Table && stackIndex.type == StackValueType::String)
	{
		auto table = std::get<std::shared_ptr<DlangTable>>(stackObj.data);
		auto key = std::get<std::string>(stackIndex.data);

		auto it = table->fields.find(key);
		if (it != table->fields.end())
			return it->second;
		

		return StackValue();
	}
	else {
		printf("Invalid type to index %s!\n", StackValueTypeToString(stackObj.type).c_str());
		return StackValue();
	}
}

double DlangVM::GetSize()
{
	auto stackObj = Pop();

	if (stackObj.type == StackValueType::Table)
	{
		auto table = std::get<std::shared_ptr<DlangTable>>(stackObj.data);
		return table->fields.size();
	}
	else if (stackObj.type == StackValueType::List)
	{
		auto list = std::get<std::shared_ptr<DlangList>>(stackObj.data);
		return list->values.size();
	}
	else if (stackObj.type == StackValueType::String)
	{
		auto str = std::get<std::string>(stackObj.data);
		return str.size();
	}
	else {
		printf("Cant get size of type %s!\n", StackValueTypeToString(stackObj.type).c_str());
		return 0;
	}
}

DlangVM::DlangVM()
{
	auto ver = getLatestVersion("downloads.dlang.eu", "/latest.txt");
	if (ver != VERSION && VERSION != "{DEVBUILD}")
		printf("[WARNING]: YOU ARE RUNNING AN OUTDATED VERSION OF DLANG!\n");

	if (!pluginManager)
		pluginManager = new ModuleManager();
}

void DlangVM::RunString(const std::string& code, std::vector<double>* extract)
{
	Lexer lexer;
	lexer.Parse(code);
	lexer.ParseOrganizedTokens();
	auto parsedTokens = lexer.GetOrganizedTokens();

#if DEBUG_MODE == 1
	for (const auto& token : parsedTokens)
		printf("[%s]:\t\t%s\n", TokenTypeToString(token.type), token.value.c_str());

	printf("=================================================\n");
#endif

	Parser parser(parsedTokens);
	auto program = parser.ParseProgram();

#if DEBUG_MODE == 1
	for (const auto& statement : program.statements)
		statement->Print();

	printf("=================================================\n");
#endif

	Compiler compiler(&program);
	compiler.Compile();

#if DEBUG_MODE == 1
	compiler.Print();
#endif

	if (extract)
		*extract = compiler.GetBytecode();
	else
		RunBytecode(compiler.GetBytecode());
}

void DlangVM::RunFile(const std::string& fileName, std::vector<double>* extract)
{
	if (!std::filesystem::exists(fileName))
		throw std::runtime_error("File '" + fileName + "' does not exist!");

	std::ifstream file(fileName, std::ios::in | std::ios::binary | std::ios::ate);
	if (!file.is_open())
		throw std::runtime_error("Failed to open file '" + fileName + "'!");

	std::streamsize size = file.tellg();
	file.seekg(0, std::ios::beg);

	std::string content;
	content.resize(size);

	if (!file.read(&content[0], size))
		throw std::runtime_error("Failed to read file '" + fileName + "', possibly corrupt?");

	RunString(content, extract);
}

void DlangVM::RunBinaryFile(const std::string& fileName)
{
	if (!std::filesystem::exists(fileName))
		throw std::runtime_error("File '" + fileName + "' does not exist!");

	std::ifstream file(fileName, std::ios::binary | std::ios::ate);
	if (!file.is_open())
		throw std::runtime_error("Failed to open file '" + fileName + "'!");

	std::streamsize size = file.tellg();
	file.seekg(0, std::ios::beg);

	std::vector<double> bytecode(size / sizeof(double));
	if (file.read(reinterpret_cast<char*>(bytecode.data()), size)) {
		RunBytecode(bytecode);
	}
	else throw std::runtime_error("Failed to read binary dlang file!");
}

std::vector<double> DlangVM::ParseHeader(const std::vector<double>& bytecode)
{
	if (bytecode.size() < 3)
		throw std::runtime_error("Bytecode too short!");

	size_t ip = 0;

	if (bytecode[ip++] != 0xFF ||
		bytecode[ip++] != 0x8C ||
		bytecode[ip++] != 0xFF)
		throw std::runtime_error("Failed to load dlang header from bytecode!");

	size_t stringTableSize = static_cast<size_t>(bytecode[ip++]);

	for (size_t i = 0; i < stringTableSize; ++i)
		strings.push(GetRealString(bytecode, ip));

	std::vector<double> instructions(bytecode.begin() + ip, bytecode.end());

	return instructions;
}

void DlangVM::RunBytecode(const std::vector<double>& bytecode)
{
	auto instructions = ParseHeader(bytecode);
	Execute(instructions);
}

StackValue DlangVM::GetVariable(const std::string& name)
{
	if (_callStack.top() > 0)
	{
		auto& currentFrame = _callStack.peek();
		auto local = currentFrame.locals.Get(name);
		if (local.type != StackValueType::NIL)
			return local;
	}

	auto glob = globals.Get(name);
	if (glob.type != StackValueType::NIL)
		return glob;

	return StackValue();
}

void DlangVM::SetVariable(const std::string& name, StackValue value)
{
	bool updated = false;

	if (_callStack.top() > 0)
	{
		auto& currentFrame = _callStack.peek();
		auto local = currentFrame.locals.Get(name);
		if (local.type != StackValueType::NIL)
		{
			currentFrame.locals.Set(name, value);
			updated = true;
		}
	}

	if (!updated)
	{
		auto glob = globals.Get(name);
		if (glob.type != StackValueType::NIL)
		{
			globals.Set(name, value);
			updated = true;
		}
	}

	if (!updated)
	{
		if (_callStack.top() > 0)
		{
			_callStack.peek().locals.Set(name, value);
		}
		else
		{
			globals.Set(name, value);
		}
	}
}