#include "stdafx.h"
#include <regex>
#include <unordered_map>
#include "../Utilities/HexUtilities.h"
#include "../Utilities/StringUtilities.h"
#include "Assembler.h"
#include "Cpu.h"
#include "CpuDisUtils.h"
#include "DisassemblyInfo.h"
#include "LabelManager.h"

static const std::regex instRegex = std::regex(
	"^\\s*([a-zA-Z]{3})[\\s]*(#%|#){0,1}([([]{0,1})[\\s]*([$]{0,1})([^\\[\\],)(;:]*)[\\s]*((,[$][0-9a-f]{1,2}|,x\\)|\\),y|,x|,y|,s\\),y|,s|\\)|\\],y|\\]){0,1})\\s*(;*)(.*)",
	std::regex_constants::icase);
static const std::regex isCommentOrBlank = std::regex("^\\s*([;]+.*$|\\s*$)", std::regex_constants::icase);
static const std::regex labelRegex = std::regex("^\\s*([@_a-zA-Z][@_a-zA-Z0-9]*):(.*)", std::regex_constants::icase);
static const std::regex byteRegex = std::regex(
	"^\\s*[.]db\\s+((\\$[a-fA-F0-9]{1,2}[ ])*)(\\$[a-fA-F0-9]{1,2})+\\s*(;*)(.*)$", std::regex_constants::icase);

void Assembler::ProcessLine(string code, uint32_t& instructionAddress, vector<int16_t>& output,
                            std::unordered_map<string, uint32_t>& labels, bool firstPass,
                            std::unordered_map<string, uint32_t>& currentPassLabels)
{
	//Remove extra spaces as part of processing
	size_t offset = code.find_first_of(',', 0);
	if (offset != string::npos)
	{
		code.erase(std::remove(code.begin() + offset + 1, code.end(), ' '), code.end());
	}
	offset = code.find_first_of(')', 0);
	if (offset != string::npos)
	{
		code.erase(std::remove(code.begin() + offset + 1, code.end(), ' '), code.end());
	}

	//Determine if the line is blank, a comment, a label or code
	std::smatch match;
	if (std::regex_search(code, match, byteRegex))
	{
		vector<string> bytes = StringUtilities::Split(match.str(1) + match.str(3), ' ');
		for (string& byte : bytes)
		{
			output.push_back((uint8_t)(HexUtilities::FromHex(byte.substr(1))));
			instructionAddress++;
		}
		output.push_back(AssemblerSpecialCodes::EndOfLine);
	}
	else if (std::regex_search(code, match, labelRegex))
	{
		string label = match.str(1);
		string afterLabel = match.str(2);
		if (currentPassLabels.find(match.str(1)) != currentPassLabels.end())
		{
			output.push_back(AssemblerSpecialCodes::LabelRedefinition);
		}
		else
		{
			labels[match.str(1)] = instructionAddress;
			currentPassLabels[match.str(1)] = instructionAddress;
			ProcessLine(afterLabel, instructionAddress, output, labels, firstPass, currentPassLabels);
		}
		return;
	}
	else if (std::regex_search(code, match, isCommentOrBlank))
	{
		output.push_back(AssemblerSpecialCodes::EndOfLine);
		return;
	}
	else if (std::regex_search(code, match, instRegex) && match.size() > 1)
	{
		LineData lineData;
		AssemblerSpecialCodes result = GetLineData(match, lineData, labels, firstPass);
		if (result == AssemblerSpecialCodes::OK)
		{
			AssembleInstruction(lineData, instructionAddress, output, firstPass);
		}
		else
		{
			output.push_back(result);
		}
	}
	else
	{
		output.push_back(AssemblerSpecialCodes::ParsingError);
	}
}

AssemblerSpecialCodes Assembler::GetLineData(std::smatch match, LineData& lineData,
                                             std::unordered_map<string, uint32_t>& labels, bool firstPass)
{
	bool isBinary = match.str(2).length() > 1 && match.str(2)[1] == '%'; //Immediate + binary: "#%"

	lineData.OpCode = match.str(1);
	lineData.IsImmediate = !match.str(2).empty();
	lineData.IsHex = !match.str(4).empty();
	lineData.HasComment = !match.str(8).empty();
	lineData.OperandSuffix = match.str(6);
	lineData.HasOpeningParenthesis = match.str(3) == "(";
	lineData.HasOpeningBracket = match.str(3) == "[";

	std::transform(lineData.OperandSuffix.begin(), lineData.OperandSuffix.end(), lineData.OperandSuffix.begin(),
	               ::toupper);
	std::transform(lineData.OpCode.begin(), lineData.OpCode.end(), lineData.OpCode.begin(), ::toupper);

	bool foundSpace = false;
	for (char c : match.str(5))
	{
		if (c != ' ' && c != '\t')
		{
			if (foundSpace)
			{
				//can't have spaces in operands (except at the very end)
				return AssemblerSpecialCodes::InvalidSpaces;
			}
			else if (lineData.IsHex && !((c >= '0' && c <= '9') || (c >= 'A' && c <= 'F') || (c >= 'a' && c <= 'f')))
			{
				//invalid hex
				return AssemblerSpecialCodes::InvalidHex;
			}
			else if (isBinary && c != '0' && c != '1')
			{
				return AssemblerSpecialCodes::InvalidBinaryValue;
			}

			lineData.Operand.push_back(c);
		}
		else
		{
			foundSpace = true;
		}
	}

	if (isBinary)
	{
		//Convert the binary value to hex
		if (lineData.Operand.size() == 0)
		{
			return AssemblerSpecialCodes::MissingOperand;
		}
		else if (lineData.Operand.size() <= 8)
		{
			lineData.IsHex = true;
			int value = 0;
			for (size_t i = 0; i < lineData.Operand.size(); i++)
			{
				value <<= 1;
				value |= lineData.Operand[i] == '1' ? 1 : 0;
			}
			lineData.Operand = HexUtilities::ToHex(value, false);
		}
		else
		{
			return AssemblerSpecialCodes::OperandOutOfRange;
		}
	}

	if (!lineData.HasComment && !match.str(9).empty())
	{
		//something is trailing at the end of the line, and it's not a comment
		return AssemblerSpecialCodes::TrailingText;
	}

	if (!lineData.IsHex)
	{
		bool allNumeric = true;
		for (size_t i = 0; i < lineData.Operand.size(); i++)
		{
			if (lineData.Operand[i] == '-' && i == 0 && lineData.Operand.size() > 1)
			{
				//First char is a minus sign, and more characters follow, continue
				continue;
			}

			if ((lineData.Operand[i] < '0' || lineData.Operand[i] > '9'))
			{
				allNumeric = false;
				break;
			}
		}

		if (allNumeric && !lineData.Operand.empty())
		{
			//Operand is not empty, and it only contains decimal values
			lineData.IsDecimal = true;
		}
		else
		{
			lineData.IsDecimal = false;
		}
	}

	return GetAddrModeAndOperandSize(lineData, labels, firstPass);
}

AssemblerSpecialCodes Assembler::GetAddrModeAndOperandSize(LineData& lineData,
                                                           std::unordered_map<string, uint32_t>& labels, bool firstPass)
{
	int opSize = 0;
	bool invalid = false;
	string operand = lineData.Operand;

	if (lineData.IsHex)
	{
		if (operand.size() == 0)
		{
			return AssemblerSpecialCodes::MissingOperand;
		}
		else if (operand.size() <= 2)
		{
			opSize = 1;
		}
		else if (operand.size() <= 4)
		{
			opSize = 2;
		}
		else if (operand.size() <= 6)
		{
			opSize = 3;
		}
		else
		{
			return AssemblerSpecialCodes::OperandOutOfRange;
		}
	}
	else if (lineData.IsDecimal)
	{
		int value = std::stoi(operand.c_str());
		if (value < -0x800000)
		{
			//< -2^23 is invalid
			return AssemblerSpecialCodes::OperandOutOfRange;
		}
		else if (value < -0x8000)
		{
			opSize = 3;
		}
		else if (value < -0x80)
		{
			opSize = 2;
		}
		else if (value <= 0xFF)
		{
			opSize = 1;
		}
		else if (value <= 0xFFFF)
		{
			opSize = 2;
		}
		else if (value <= 0xFFFFFF)
		{
			opSize = 3;
		}
		else
		{
			//>= 2^23 is invalid
			return AssemblerSpecialCodes::OperandOutOfRange;
		}
	}
	else if (!operand.empty())
	{
		//Check if the operand is a known label
		auto findResult = labels.find(operand);
		if (findResult != labels.end())
		{
			lineData.Operand = HexUtilities::ToHex((uint16_t)findResult->second);
			lineData.IsHex = true;
			opSize = 2;
		}
		else if (operand.size() == 1 && (operand[0] == 'A' || operand[0] == 'a') && lineData.OperandSuffix.empty() && !
			lineData.IsHex && !lineData.IsImmediate && !lineData.HasOpeningParenthesis && !lineData.HasOpeningBracket)
		{
			//Allow optional "A" after AddrMode == Accumulator instructions
			lineData.Mode = AddrMode::Acc;
			opSize = 0;
		}
		else
		{
			int32_t addr = _labelManager->GetLabelRelativeAddress(operand);
			if (addr > 0xFFFF)
			{
				lineData.Operand = HexUtilities::ToHex24(addr);
				lineData.IsHex = true;
				opSize = 3;
			}
			else if (addr > 0xFF)
			{
				lineData.Operand = HexUtilities::ToHex((uint16_t)addr);
				lineData.IsHex = true;
				opSize = 2;
			}
			else if (addr >= 0)
			{
				lineData.Operand = HexUtilities::ToHex((uint8_t)addr);
				lineData.IsHex = true;
				opSize = 1;
			}
			else
			{
				if (firstPass)
				{
					//First pass, we couldn't find a matching label, so it might be defined later on
					//Pretend it exists for now
					_needSecondPass = true;
					lineData.Operand = "FFFF";
					lineData.IsHex = true;
					opSize = 2;
				}
				else
				{
					return AssemblerSpecialCodes::UnknownLabel;
				}
			}
		}
	}
	else
	{
		//No operand
		opSize = 0;
	}

	if (lineData.Mode == AddrMode::Imp)
	{
		if (lineData.OperandSuffix.substr(0, 2) == ",$")
		{
			//Used by MVP, MVN
			opSize = 2;
			lineData.Mode = AddrMode::BlkMov;
			uint8_t dest = HexUtilities::FromHex(lineData.OperandSuffix.substr(2));
			lineData.Operand += HexUtilities::ToHex(dest);
		}
		else if (lineData.IsImmediate)
		{
			if (lineData.HasOpeningParenthesis || lineData.HasOpeningBracket || opSize == 0)
			{
				invalid = true;
			}
			else if (opSize >= 3)
			{
				invalid = true;
			}
			lineData.Mode = opSize == 2 ? AddrMode::Imm16 : AddrMode::Imm8; //or Rel
		}
		else if (lineData.HasOpeningBracket)
		{
			if (lineData.OperandSuffix == "]")
			{
				switch (opSize)
				{
				case 1: lineData.Mode = AddrMode::DirIndLng;
					break;
				case 2: lineData.Mode = AddrMode::AbsIndLng;
					break;
				default: invalid = true;
					break;
				}
			}
			else if (lineData.OperandSuffix == "],Y")
			{
				if (opSize == 1)
				{
					lineData.Mode = AddrMode::DirIndLngIdxY;
				}
				else
				{
					invalid = true;
				}
			}
		}
		else if (lineData.HasOpeningParenthesis)
		{
			if (lineData.OperandSuffix == ")")
			{
				lineData.Mode = opSize == 1 ? AddrMode::DirInd : AddrMode::AbsInd;
			}
			else if (lineData.OperandSuffix == ",X)")
			{
				lineData.Mode = opSize == 1 ? AddrMode::DirIdxIndX : AddrMode::AbsIdxXInd;
			}
			else if (lineData.OperandSuffix == "),Y")
			{
				if (opSize == 1)
				{
					lineData.Mode = AddrMode::DirIndIdxY;
				}
				else
				{
					return AssemblerSpecialCodes::OperandOutOfRange;
				}
			}
			else if (lineData.OperandSuffix == ",S),Y")
			{
				if (opSize == 1)
				{
					lineData.Mode = AddrMode::StkRelIndIdxY;
				}
				else
				{
					return AssemblerSpecialCodes::OperandOutOfRange;
				}
			}
			else
			{
				invalid = true;
			}
		}
		else
		{
			if (lineData.OperandSuffix == ",X")
			{
				if (opSize == 3)
				{
					lineData.Mode = AddrMode::AbsLngIdxX;
				}
				else if (opSize == 2)
				{
					lineData.Mode = AddrMode::AbsIdxX;
				}
				else if (opSize == 1)
				{
					//Sometimes zero page addressing is not available, even if the operand is in the zero page
					lineData.Mode = AddrMode::DirIdxX;
				}
				else
				{
					invalid = true;
				}
			}
			else if (lineData.OperandSuffix == ",Y")
			{
				if (opSize == 2)
				{
					lineData.Mode = AddrMode::AbsIdxY;
				}
				else if (opSize == 1)
				{
					//Sometimes zero page addressing is not available, even if the operand is in the zero page
					lineData.Mode = AddrMode::DirIdxY;
				}
				else
				{
					invalid = true;
				}
			}
			else if (lineData.OperandSuffix == ",S")
			{
				if (opSize == 1)
				{
					lineData.Mode = AddrMode::StkRel;
				}
				else
				{
					return AssemblerSpecialCodes::OperandOutOfRange;
				}
			}
			else if (lineData.OperandSuffix.empty())
			{
				if (opSize == 0)
				{
					lineData.Mode = AddrMode::Imp; //or Acc
				}
				else if (opSize == 3)
				{
					lineData.Mode = AddrMode::AbsLng;
				}
				else if (opSize == 2)
				{
					lineData.Mode = AddrMode::Abs;
				}
				else if (opSize == 1)
				{
					//Sometimes zero page addressing is not available, even if the operand is in the zero page
					lineData.Mode = AddrMode::Dir;
				}
				else
				{
					invalid = true;
				}
			}
			else
			{
				invalid = true;
			}
		}
	}

	/*if(lineData.Mode == AddrMode::None) {
		invalid = true;
	}*/

	lineData.OperandSize = opSize;

	return invalid ? AssemblerSpecialCodes::ParsingError : AssemblerSpecialCodes::OK;
}

bool Assembler::IsOpModeAvailable(string& opCode, AddrMode mode)
{
	return _availableModesByOpName[opCode].find((int)mode) != _availableModesByOpName[opCode].end();
}

void Assembler::AssembleInstruction(LineData& lineData, uint32_t& instructionAddress, vector<int16_t>& output,
                                    bool firstPass)
{
	bool foundMatch = false;

	for (int i = 0; i < 256; i++)
	{
		AddrMode opMode = CpuDisUtils::OpMode[i];
		if (lineData.OpCode == CpuDisUtils::OpName[i])
		{
			bool modeMatch = opMode == lineData.Mode;
			if (!modeMatch)
			{
				if (lineData.Mode == AddrMode::Imp && (opMode == AddrMode::Acc || opMode == AddrMode::Stk))
				{
					modeMatch = true;
				}
				else if (lineData.Mode == AddrMode::Imm8 && opMode == AddrMode::Sig8)
				{
					modeMatch = true;
				}
				else if ((lineData.Mode == AddrMode::Imm8 || lineData.Mode == AddrMode::Imm16) && (opMode == AddrMode::ImmX
					|| opMode == AddrMode::ImmM))
				{
					modeMatch = true;
				}
				else if ((opMode == AddrMode::Rel || opMode == AddrMode::RelLng) && (lineData.Mode == AddrMode::AbsLng ||
					lineData.Mode == AddrMode::Abs || lineData.Mode == AddrMode::Dir))
				{
					bool lngBranch = opMode == AddrMode::RelLng;
					modeMatch = true;

					//Convert "absolute" jump to a relative jump
					int value = lineData.IsHex ? HexUtilities::FromHex(lineData.Operand) : std::stoi(lineData.Operand);
					if (lineData.Mode == AddrMode::Abs)
					{
						value |= (instructionAddress & 0xFF0000);
					}

					int32_t addressGap;
					if (lineData.Mode == AddrMode::Dir)
					{
						addressGap = (int8_t)value;
					}
					else
					{
						addressGap = value - (instructionAddress + (lngBranch ? 3 : 2));
					}

					if (addressGap > (lngBranch ? 32767 : 127) || addressGap < (lngBranch ? -32768 : -128))
					{
						//Gap too long, can't jump that far
						if (!firstPass)
						{
							//Pretend this is ok on first pass, we're just trying to find all labels
							output.push_back(AssemblerSpecialCodes::OutOfRangeJump);
							return;
						}
					}

					//Update data to match relative jump
					lineData.OperandSize = lngBranch ? 2 : 1;
					lineData.IsHex = true;
					lineData.Operand = lngBranch
						                   ? HexUtilities::ToHex((uint16_t)addressGap)
						                   : HexUtilities::ToHex((uint8_t)addressGap);
				}
			}

			if (modeMatch)
			{
				output.push_back(i);
				instructionAddress += (lineData.OperandSize + 1);

				if (lineData.OperandSize == 1)
				{
					int value = lineData.IsHex ? HexUtilities::FromHex(lineData.Operand) : std::stoi(lineData.Operand);
					output.push_back(value & 0xFF);
				}
				else if (lineData.OperandSize == 2)
				{
					int value = lineData.IsHex ? HexUtilities::FromHex(lineData.Operand) : std::stoi(lineData.Operand);
					output.push_back(value & 0xFF);
					output.push_back((value >> 8) & 0xFF);
				}
				else if (lineData.OperandSize == 3)
				{
					int value = lineData.IsHex ? HexUtilities::FromHex(lineData.Operand) : std::stoi(lineData.Operand);
					output.push_back(value & 0xFF);
					output.push_back((value >> 8) & 0xFF);
					output.push_back((value >> 16) & 0xFF);
				}

				foundMatch = true;
				break;
			}
		}
	}

	if (!foundMatch)
	{
		output.push_back(AssemblerSpecialCodes::InvalidInstruction);
	}
	else
	{
		output.push_back(AssemblerSpecialCodes::EndOfLine);
	}
}

Assembler::Assembler(shared_ptr<LabelManager> labelManager)
{
	_labelManager = labelManager;
}

Assembler::~Assembler()
{
}

uint32_t Assembler::AssembleCode(string code, uint32_t startAddress, int16_t* assembledCode)
{
	for (uint8_t i = 0; i < 255; i++)
	{
		if (_availableModesByOpName.find(CpuDisUtils::OpName[i]) == _availableModesByOpName.end())
		{
			_availableModesByOpName[CpuDisUtils::OpName[i]] = std::unordered_set<int>();
		}
		_availableModesByOpName[CpuDisUtils::OpName[i]].emplace((int)CpuDisUtils::OpMode[i]);
	}

	std::unordered_map<string, uint32_t> temporaryLabels;
	std::unordered_map<string, uint32_t> currentPassLabels;

	size_t i = 0;
	vector<int16_t> output;
	output.reserve(1000);

	uint32_t originalStartAddr = startAddress;

	vector<string> codeLines;
	codeLines.reserve(100);
	while (i < code.size())
	{
		size_t offset = code.find_first_of('\n', i);
		string line;
		if (offset != string::npos)
		{
			line = code.substr(i, offset - i);
			i = offset + 1;
		}
		else
		{
			line = code.substr(i);
			i = code.size();
		}
		codeLines.push_back(line);
	}

	//Make 2 passes - first one to find all labels, second to assemble
	_needSecondPass = false;
	for (string& line : codeLines)
	{
		ProcessLine(line, startAddress, output, temporaryLabels, true, currentPassLabels);
	}

	if (_needSecondPass)
	{
		currentPassLabels.clear();
		output.clear();
		startAddress = originalStartAddr;
		for (string& line : codeLines)
		{
			ProcessLine(line, startAddress, output, temporaryLabels, false, currentPassLabels);
		}
	}

	memcpy(assembledCode, output.data(), std::min<int>(100000, (int)output.size()) * sizeof(uint16_t));
	return (uint32_t)output.size();
}