#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <vector>
#include <map>
#include <sstream>
#include <fstream>
namespace VM{
    enum OpCode{
		NOP,
		PRINT_CCHAR,
		PRINT_CBYTE,
		PRINT_CHAR,
		PRINT_INT,
		JMP,
		JEQ,
		SET_BYTE,
		SET_INT,
		INC,
		ADD,
		MOV
	};
	class VirtualMachine{
	public:
		void load(const std::string &program){
			m_program = program;
			m_programCounter = 0;
		}

		unsigned char fetch(){
			return m_program[m_programCounter++];
		}

		int fetchInt(){
			int i = fetch();
			i = i *256 + fetch();
			i = i *256 + fetch();
			return i * 256 + fetch();
		}

		void step(){
			OpCode op = (OpCode)fetch();
			switch(op){
			case NOP:
				break;
			case PRINT_CCHAR:
				printf("%c",fetch());
				break;
			case PRINT_CHAR:
				printf("%c",(char) m_memory[fetch()]);
				break;
			case PRINT_CBYTE:
				printf("%d",(int)fetch());
				break;
			case PRINT_INT:
				printf("%d",m_memory[fetch()]);
				break;
			case JMP:
				m_programCounter = fetchInt();
				break;
			case JEQ:
				{
					int a = m_memory[fetch()];
					int b = m_memory[fetch()];
					unsigned char target = fetchInt();
					if(a == b){
						m_programCounter = target;
					}
					break;
				}
			case SET_BYTE:
				{
					unsigned char addr = fetch();
					m_memory[addr] = fetch();
					break;
				}
			case SET_INT:
				{
					unsigned char addr = fetch();
					m_memory[addr] = fetchInt();
					break;
				}
			case INC:
				m_memory[fetch()]++;
				break;
			case ADD:
				{
					unsigned char addr1 = fetch();
					unsigned char addr2 = fetch();
					m_memory[addr1] += m_memory[addr2];
					break;
				}
			case MOV:
				{
					unsigned char addr1 = fetch();
					unsigned char addr2 = fetch();
					m_memory[addr1] = m_memory[addr2];
					break;
				}
			}
			
		
		}

		void run(){
			while(m_programCounter < m_program.size()){
				step();
			}
		}
	private:
		std::string m_program;
		size_t m_programCounter;
		int m_memory[256];
	};

	class Assembler{
	public:
		void assemble(const std::string &source){
			m_bytecode = "";
			m_pos = 0;
			m_labels.clear();
			m_labelRefs.clear();

			std::vector<std::string> lines = tokenize(source,"\n");
			for(std::vector<std::string>::iterator it = lines.begin();it!=lines.end();it++){
				const std::string &line = *it;

				std::vector<std::string> tokens = tokenize(line," \t");
				
				const std::string &op = tokens[0];

				if(op[0] == '#'){
					m_labels[op] = m_pos;
				}

				if(op == "NOP"){
					pushChar(NOP);
				}else if(op == "PRINT_CCHAR"){
					pushChar(PRINT_CCHAR);
					pushChar(toInt(tokens[1]));
				}else if(op == "PRINT_CBYTE"){
					pushChar(PRINT_CBYTE);
					pushChar(toInt(tokens[1]));
				}else if(op == "PRINT_CHAR"){
					pushChar(PRINT_CHAR);
					pushChar(toInt(tokens[1]));
				}else if(op == "PRINT_INT"){
					pushChar(PRINT_INT);
					pushChar(toInt(tokens[1]));
				}else if(op == "JMP"){
					pushChar(JMP);
					pushLabel(tokens[1]);
				}else if(op == "JEQ"){
					pushChar(JEQ);
					pushChar(toInt(tokens[1]));
					pushChar(toInt(tokens[2]));
					pushLabel(tokens[3]);
				}else if(op == "SET_BYTE"){
					pushChar(SET_BYTE);
					pushChar(toInt(tokens[1]));
					pushChar(toInt(tokens[2]));
				}else if(op == "SET_INT"){
					pushChar(SET_INT);
					pushChar(toInt(tokens[1]));
					pushInt(toInt(tokens[2]));
				}else if(op == "INC"){
					pushChar(INC);
					pushChar(toInt(tokens[1]));
				}else if(op == "ADD"){
					pushChar(ADD);
					pushChar(toInt(tokens[1]));
					pushChar(toInt(tokens[2]));
				}else if(op == "MOV"){
					pushChar(MOV);
					pushChar(toInt(tokens[1]));
					pushChar(toInt(tokens[2]));
				}

			}

			for(std::vector<LabelReference>::iterator it = m_labelRefs.begin();it!=m_labelRefs.end();it++){
				const LabelReference &ref = *it;
				std::map<std::string,int>::iterator label = m_labels.find(ref.name);
				if(label == m_labels.end()){
					printf("error: label %s not found.\n",ref.name.c_str());
				}else{
					int i = (*label).second;

					m_bytecode[ref.position] = (i & 0xff000000) >> 24;
					m_bytecode[ref.position+1] = (i & 0x00ff0000) >> 16;
					m_bytecode[ref.position+2] = (i & 0x0000ff00) >> 8;
					m_bytecode[ref.position+3] = (i & 0x000000ff) >> 0;
				}
			}

		}

		const std::string &getBytecode() const{
			return m_bytecode;
		}
	private:

		std::string m_bytecode;
		size_t m_pos;
		std::map<std::string,int> m_labels;

		struct LabelReference{
			std::string name;
			int position;
		};
		std::vector<LabelReference> m_labelRefs;

		void pushChar(char c){
			m_bytecode.push_back(c);
			m_pos++;
		}

		void pushLabel(const std::string &str){
			if(str[0] == '#'){
				LabelReference ref;
				ref.name = str;
				ref.position = m_pos;
				m_labelRefs.push_back(ref);
				pushInt(0x1a);
			}else{
				pushInt(toInt(str));
			}
		}

		void pushInt(int i){
			m_bytecode.push_back((i & 0xff000000) >> 24);
			m_bytecode.push_back((i & 0x00ff0000) >> 16);
			m_bytecode.push_back((i & 0x0000ff00) >> 8);
			m_bytecode.push_back((i & 0x000000ff) >> 0);
			m_pos+= 4;
		}

		int toInt(const std::string &str){
			int i;

			const char *fmt = "%d";
			if(str[0] == '0'){
				if(str[1] == 'x' || str[1] == 'X'){
					fmt = "%x";
				}else{
					fmt = "%o";
				}
			}

			sscanf(str.c_str(),fmt,&i);
			return i;
		}

		std::vector<std::string> tokenize(const std::string &str, const char * delimiters){
			std::vector<std::string> rst;
			std::string::size_type lastPos = 0;
			std::string::size_type pos = str.find_first_of(delimiters,0);
			
			while(pos != std::string::npos){
				if(pos-lastPos){
					rst.push_back(str.substr(lastPos,pos - lastPos));
				}
				lastPos = pos + 1;
				pos = str.find_first_of(delimiters,lastPos);
			}
			rst.push_back(str.substr(lastPos));
			return rst;
		}
	};
}

int main(int argc, const char* argv[])
{
	using namespace VM;
	VirtualMachine vm;
	Assembler a;

	/*std::ifstream t(argv[1]);
	std::stringstream buffer;
	buffer << t.rdbuf();*/

	a.assemble(
		"SET_BYTE 0 1  \n"
		"SET_BYTE 1 10 \n"
		"#FOR_I\n"
		"JEQ 0 1 #FOR_I_END\n"
		"SET_BYTE 2 1 \n"
		"SET_BYTE 3 10\n"
		"SET_BYTE 4 0\n"
		"#FOR_J\n"
		"JEQ 2 3 #FOR_J_END\n"
		"ADD 4 0\n"
		"PRINT_INT 0\n"
		"PRINT_CCHAR 0x20\n"
		"PRINT_CCHAR 0x2a\n"
		"PRINT_CCHAR 0x20\n"
		"PRINT_INT 2\n"
		"PRINT_CCHAR 0x20\n"
		"PRINT_CCHAR 0x3d\n"
		"PRINT_CCHAR 0x20\n"
		"PRINT_INT 4\n"
		"PRINT_CCHAR 0x0a\n"
		"INC 2\n"
		"JMP #FOR_J\n"
		"#FOR_J_END\n"
		"INC 0\n"
		"JMP #FOR_I\n"
		"#FOR_I_END"
	);
	vm.load(a.getBytecode());
	vm.run();
	system("pause");
	return 0;
}

