#include "Assembleur.h"

#include <algorithm>
#include <map>
#include <sstream>
#include <stdexcept>

enum TypeArithmetique {
	Add = 0,
	Sub,
	And,
	Or,
	Xor,
	Sl,
	Sr,
	Mul,
};

enum TypeMemoire {
	Str = 0,
	Ld,
};

enum TypeSautControle {
	Jump = 0,
	Jequ,
	Jneq,
	Jsup,
	Jinf,
	Call,
	Ret = 7,
};

static std::map<std::string, Instruction> INSTRUCTION_KEYS = {
	{"add", {Arithmetique, Add}},
	{"sub", {Arithmetique, Sub}},
	{"and", {Arithmetique, And}},
	{"or", {Arithmetique, Or}},
	{"xor", {Arithmetique, Xor}},
	{"sl", {Arithmetique, Sl}},
	{"sr", {Arithmetique, Sr}},
	{"mul", {Arithmetique, Mul}},
	{"str", {Memoire, Str}},
	{"ld", {Memoire, Ld}},
	{"jmp", {SautControle, Jump}},
	{"jequ", {SautControle, Jequ}},
	{"jneq", {SautControle, Jneq}},
	{"jsup", {SautControle, Jsup}},
	{"jinf", {SautControle, Jinf}},
	{"call", {SautControle, Call}},
	{"ret", {SautControle, Ret}},
};

constexpr int LINE_LENGTH = 32;
constexpr int INSTRUCTION_BITS_COUNT = 2;
constexpr int OPERATION_BITS_COUNT = 3;
constexpr int IMMEDIATE_BITS_COUNT = 1;
constexpr int INSTRUCTION_BLOCK_SIZE = INSTRUCTION_BITS_COUNT + OPERATION_BITS_COUNT + IMMEDIATE_BITS_COUNT;
constexpr int REGISTRY_BITS_COUNT = 3;

static constexpr int GetOffset(int a_Start, int a_BlockSize) {
	return LINE_LENGTH - a_Start - a_BlockSize;
}

std::uint32_t Assembleur::ParseLabel(const std::string& a_Label) {
	auto it = m_Labels.find(a_Label);

	if (it == m_Labels.end()) {
		throw std::invalid_argument("Label " + a_Label + " not found !");
	}

	// the address starts at 0, not 1
	return it->second - 1;
}

void Assembleur::AddLabel(const std::string& a_Label, std::uint32_t a_Line) {
	m_Labels.insert({a_Label, a_Line});
}

// 2 bits type instruction | 3 bits sous-type opération
std::uint32_t Assembleur::IToInt(Instruction a_Instruction) {
	return static_cast<std::uint32_t>(a_Instruction.m_Instruction) << GetOffset(0, INSTRUCTION_BITS_COUNT) |
		   static_cast<std::uint32_t>(a_Instruction.m_SubInstruction) << GetOffset(INSTRUCTION_BITS_COUNT, OPERATION_BITS_COUNT);
}

// 5 bits instruction | 1 bit immédiat (non utilisé) | 3 bits R1 | 3 bits R2 | 3 bits R3
std::uint32_t Assembleur::ParseOperation(Instruction a_Instruction, std::uint32_t a_R1, std::uint32_t a_R2, std::uint32_t a_R3) {
	return IToInt(a_Instruction) | 
		a_R1 << GetOffset(INSTRUCTION_BLOCK_SIZE, REGISTRY_BITS_COUNT) | 
		a_R2 << GetOffset(INSTRUCTION_BLOCK_SIZE + REGISTRY_BITS_COUNT, REGISTRY_BITS_COUNT) | 
		a_R3 << GetOffset(INSTRUCTION_BLOCK_SIZE + 2 * REGISTRY_BITS_COUNT, REGISTRY_BITS_COUNT);
}

// 5 bits instruction | 1 bit immédiat | 3 bits R1 | 3 bits R2 | 20 bits constante
std::uint32_t Assembleur::ParseOperationImmediate(
	Instruction a_Instruction, std::uint32_t a_R1, std::uint32_t a_R2, std::uint32_t a_C1) {
	return IToInt(a_Instruction) | 1 << GetOffset(INSTRUCTION_BITS_COUNT + OPERATION_BITS_COUNT, IMMEDIATE_BITS_COUNT) | 
		a_R1 << GetOffset(INSTRUCTION_BLOCK_SIZE, REGISTRY_BITS_COUNT) | 
		a_R2 << GetOffset(INSTRUCTION_BLOCK_SIZE + REGISTRY_BITS_COUNT, REGISTRY_BITS_COUNT) | 
		a_C1;
}

// 5 bits instruction | 1 bit immédiat (non utilisé) | 26 bits adresse du label
std::uint32_t Assembleur::ParseJump(Instruction a_Instruction, const std::string& a_Label) {
	return IToInt(a_Instruction) | ParseLabel(a_Label) & 0x3FFFFFF;
}

// 5 bits instruction | 1 bit immédiat (non utilisé) | 3 bits R1 | 3 bits R2 | 20 bits adresse du label
std::uint32_t Assembleur::ParseJump(Instruction a_Instruction, std::uint8_t a_R1, std::uint8_t a_R2, const std::string& a_Label) {
	return IToInt(a_Instruction) | 
		a_R1 << GetOffset(INSTRUCTION_BLOCK_SIZE, REGISTRY_BITS_COUNT) | 
		a_R2 << GetOffset(INSTRUCTION_BLOCK_SIZE + REGISTRY_BITS_COUNT, REGISTRY_BITS_COUNT) |
		ParseLabel(a_Label) & 0xFFFFF;
}

// 5 bits instruction | 1 bit immédiat (non utilisé)
std::uint32_t Assembleur::ParseJump(Instruction a_Instruction) {
	return IToInt(a_Instruction);
}

// 5 bits instruction | 1 bit immédiat (non utilisé) | 3 bits R1 | 3 bits R2
std::uint32_t Assembleur::ParseIO(Instruction a_Instruction, std::uint32_t a_R1, std::uint32_t a_R2) {
	return IToInt(a_Instruction) | 
		a_R1 << GetOffset(INSTRUCTION_BLOCK_SIZE, REGISTRY_BITS_COUNT) | 
		a_R2 << GetOffset(INSTRUCTION_BLOCK_SIZE + REGISTRY_BITS_COUNT, REGISTRY_BITS_COUNT);
}





static std::uint32_t ParseRegistry(const std::string& a_Str) {
	if (a_Str.at(0) != 'R')
		throw std::invalid_argument("Registry " + a_Str + " not found !");
	std::uint32_t registry = std::stoi(a_Str.substr(1));
	if (registry > 7)
		throw std::invalid_argument("You can only have up to 8 registries !");
	return registry;
}

static bool IsConstant(const std::string& a_Str) {
	return a_Str.at(0) == '#';
}

static std::uint32_t ParseConstant(const std::string& a_Str) {
	if (a_Str.at(0) != '#')
		throw std::invalid_argument("Registry " + a_Str + " not found !");
	return std::stoi(a_Str.substr(1));
}

std::uint32_t Assembleur::ParseInstruction(const std::string& a_Str, std::uint32_t a_Line, std::uint32_t a_RealLine) {
	std::stringstream ss{a_Str};
	std::string ins;
	ss >> ins;

	// to lower case
	std::transform(ins.begin(), ins.end(), ins.begin(), [](unsigned char c) { return std::tolower(c); });

	auto it = INSTRUCTION_KEYS.find(ins);
	if (it == INSTRUCTION_KEYS.end()) {
		throw std::invalid_argument("[Line " + std::to_string(a_RealLine) + "] " + "Instruction \"" + ins + "\" not found !");
	}

	Instruction instruction = it->second;

	try {
		switch (instruction.m_Instruction) {
			case Arithmetique: {
				std::string R1, R2, R3;
				ss >> R1 >> R2 >> R3;
				if (IsConstant(R3)) {
					return ParseOperationImmediate(instruction, ParseRegistry(R1), ParseRegistry(R2), ParseConstant(R3));
				} else {
					return ParseOperation(instruction, ParseRegistry(R1), ParseRegistry(R2), ParseRegistry(R3));
				}
			}

			case Memoire: {
				std::string R1, R2;
				ss >> R1 >> R2;
				return ParseIO(instruction, ParseRegistry(R1), ParseRegistry(R2));
			}

			case SautControle: {
				switch (instruction.m_SubInstruction) {
					case Ret:
						return ParseJump(instruction);

					case Call:
					case Jump: {
						std::string label;
						ss >> label;
						return ParseJump(instruction, label);
					}

					case Jequ:
					case Jneq:
					case Jsup:
					case Jinf: {
						std::string R1, R2, label;
						ss >> R1 >> R2 >> label;
						return ParseJump(instruction, ParseRegistry(R1), ParseRegistry(R2), label);
					}
				}
			}
		}

	} catch (std::exception& e) {
		throw std::invalid_argument(" [Line " + std::to_string(a_RealLine) + "] " + e.what() + "\n" + a_Str);
	}

	return 0;
}