#pragma once

#include "cpustate.hpp"

#include "le_parse_util.hpp"

#include <distorm.h>
#include <mnemonics.h>

#include <boost/variant.hpp>

#include <fstream>
#include <iomanip>

class Emulator {
public:
	CpuState cpu;
	std::vector<uint8_t> memory;
	
	std::map<int, std::function<void(_DInst)>> opcode_handlers;
	
	typedef boost::variant<uint32_t&, uint16_t&, uint8_t&> RefValueType;
	typedef boost::variant<uint32_t, uint16_t, uint8_t> ValueType;
	std::map<int, std::function<RefValueType()>> register_getters;
	
	std::map<int, std::function<void()>> interrupt_handlers;
	
	ValueType extract_immediate_of_size(int size, _DInst inst) {
		switch(size) {
			case 8: {
				return inst.imm.byte;
			}
			case 16: {
				return inst.imm.word;
			}
			case 32: {
				return inst.imm.dword;
			}
			default: {
				throw UnhandledInstruction();
			}
		}
	}
	
	struct print : public boost::static_visitor<std::ostream&> {
		template <typename T>
		std::ostream& operator()(T const& t) const {
			return std::cout << t;
		}
	};
	
	struct dest_size : public boost::static_visitor<unsigned int> {
		template<typename T>
		unsigned int operator()(T const&) const {
			return sizeof(T)*8;
		}
	};
	
	struct return_visitor : public boost::static_visitor<uint32_t> {
		template<typename T>
		uint32_t operator()(T const& x) const {
			return x;
		}
	};
	
	struct IncompatibleArguments : public std::runtime_error {
		IncompatibleArguments() 
		: std::runtime_error("The two arguments to this instruction are not compatible")
		{}
	};
	
	template <typename Ret>
	struct binary_op : public boost::static_visitor<Ret> {
		template <typename T, typename U>
		Ret operator()(T&, U const&) const {
			throw IncompatibleArguments();
		}
	};
	
	struct and_ : public binary_op<void> {
		using binary_op::operator();
		
		template<typename T>
		void operator()(T& lh, T const& rh) const {
			lh &= rh;
		}
	};
	
	struct mov : public binary_op<void> {
		using binary_op::operator();
		
		template<typename T>
		void operator()(T& lh, T const& rh) const {
			lh = rh;
		}
	};
	
	struct sub : public binary_op<void>  {
		using binary_op::operator();
		
		template<typename T>
		void operator()(T& lh, T const& rh) const {
			lh -= rh;
		}
	};
	
	template <typename O>
	void handle_binary(_DInst inst) {
		std::unique_ptr<RefValueType> dest;
		ValueType src;
		
		{
			_Operand dest_op = inst.ops[0];
			
			switch(dest_op.type) {
				case O_REG: {
					dest = std::make_unique<RefValueType>(register_getters.at(dest_op.index)());
					break;
				}
				case O_MEM: {
					throw UnhandledInstruction();
					break;
				}
				case O_SMEM: {
					throw UnhandledInstruction();
					break;
				}
				case O_DISP: {
					switch(dest_op.size) {
						case 8: {
							dest = std::make_unique<RefValueType>(*reinterpret_cast<uint8_t*>(&memory[inst.disp]));
							break;
						}
						case 16: {
							dest = std::make_unique<RefValueType>(*reinterpret_cast<uint16_t*>(&memory[inst.disp]));
							break;
						}
						case 32: {
							dest = std::make_unique<RefValueType>(*reinterpret_cast<uint32_t*>(&memory[inst.disp]));
							break;
						}
						default: {
							throw UnhandledInstruction();
						}
					}
					break;
				}
				default: {
					throw UnhandledInstruction();
				}
			}
		}

		{
			_Operand src_op = inst.ops[1];
			
			switch(src_op.type) {
				case O_REG: {
					auto x = register_getters.at(src_op.index)();
					src = register_getters.at(src_op.index)();
					break;
				}
				case O_IMM: {
					//Do not rely on the imm size as we need to sign extend to dest size.
					src = extract_immediate_of_size(boost::apply_visitor(dest_size(), *dest), inst);
					break;
				}
				default: {
					throw UnhandledInstruction();
				}
			}
		}
		
		boost::apply_visitor(O(), *dest, src);
	}
	
	void handle_I_JMP(_DInst inst) {
		if(inst.ops[0].type == O_PC || inst.ops[0].type == O_PTR || inst.ops[0].type == O_DISP) {
			binparse::Offset32 target = binparse::Offset32(INSTRUCTION_GET_TARGET(&inst));
			cpu.eip() = target;
		} else if (inst.ops[0].type == O_SMEM) {
			throw UnhandledInstruction();
		} else {
			throw UnrecognizedInstruction();
		}
	}
	void handle_I_STI(_DInst) {
		cpu.intf() = 1;
	}
	
	void handle_I_AND(_DInst inst) {
		handle_binary<and_>(inst);
	}
	
	void handle_I_MOV(_DInst inst) {
		handle_binary<mov>(inst);
	}
	
	void handle_I_SUB(_DInst inst) {
		handle_binary<sub>(inst);
	}
	
	void handle_I_INT(_DInst inst) {
		_Operand int_op = inst.ops[0];
		
		ValueType imm = extract_immediate_of_size(int_op.size, inst);
		
		interrupt_handlers.at(boost::apply_visitor(return_visitor(), imm))();
	}
	
	void int_0x21() {
		if(cpu.ah() == 0x30) {
			cpu.al() = 6;
			cpu.ah() = 0;
			cpu.bh() = 0;
		} 
	}
	
public:
	Emulator(binparse::Offset32 init_eip, binparse::Offset32 init_esp, std::vector<uint8_t> memory) 
	: cpu() 
	, memory(memory)
	{
		cpu.eip() = init_eip;
		cpu.esp() = init_esp;
	
		#define REGISTER_HANDLER(OPCODE) opcode_handlers[OPCODE] = std::bind(&Emulator::handle_##OPCODE, this, std::placeholders::_1)
		
		REGISTER_HANDLER(I_JMP);
		REGISTER_HANDLER(I_STI);
		REGISTER_HANDLER(I_AND);
		REGISTER_HANDLER(I_MOV);
		REGISTER_HANDLER(I_SUB);
		REGISTER_HANDLER(I_INT);
		
		#undef REGISTER_HANDLER
		
		#define REGISTER_GETTER(REGISTER, R_NAME) register_getters[REGISTER] = [this]{return RefValueType(cpu.R_NAME());}
		
		REGISTER_GETTER(R_ESP, esp);
		REGISTER_GETTER(R_AL , al );
		REGISTER_GETTER(R_AH , ah );
		REGISTER_GETTER(R_AX , ax );
		REGISTER_GETTER(R_EAX, eax);
		REGISTER_GETTER(R_EBX, ebx);
		
		#undef REGISTER_GETTER
		
		#define REGISTER_INTERRUPT_HANDLER(INT) interrupt_handlers[INT] = [this]{return int_##INT();}
		
		REGISTER_INTERRUPT_HANDLER(0x21);
		
		#undef REGISTER_INTERRUPT_HANDLER 
	}
	
	struct UnhandledInstruction : public std::runtime_error {
		UnhandledInstruction()
		: std::runtime_error("Encountered unhandled instruction")
		{}
	};
	
	struct UnrecognizedInstruction : public std::runtime_error {
		UnrecognizedInstruction()
		: std::runtime_error("Encountered unhandled instruction")
		{}
	};
	
	void handle_instruction(_DInst inst) {
		opcode_handlers.at(inst.opcode)(inst);
	}
	
};

void emulate(std::string file_path) {
	std::ifstream file_stream(file_path, std::ios::binary);
	file_stream.unsetf(std::ios::skipws);
	auto file = le::parse_file(file_stream);
	auto binary = load_binary(file);
	
	_CodeInfo ci;
	_DecodeType dt = Decode32Bits;
	
	auto code_object = file.object_table.entries.at(file.le_header.EIP_object);
	auto initial_eip = code_object.reloc_base_address + file.le_header.EIP;
	
	auto data_object = file.object_table.entries.at(file.le_header.ESP_object);
	auto initial_esp = data_object.reloc_base_address + file.le_header.ESP;
	
	Emulator emulator(initial_eip, initial_esp, binary);
	
	unsigned int decodedInstructionsCount;
	
	bool run = true;
	while(run) {
		ci.code = binary.data() + emulator.cpu.eip();
		ci.nextOffset = emulator.cpu.eip();
		ci.codeLen = binary.size() - emulator.cpu.eip();
		ci.codeOffset = emulator.cpu.eip();
		ci.dt = dt;
		ci.features = DF_NONE;
	
		_DInst decinst;
		distorm_decompose(&ci, &decinst, 1, &decodedInstructionsCount);
		
		emulator.cpu.eip() += decinst.size;
		
		_DecodedInst inst;
		distorm_format64(&ci, &decinst, &inst);
		std::cout << "CurrentInstruction: " << std::hex << std::setw(8) << std::setfill('0') << inst.offset << ":\t" << inst.mnemonic.p << " " << inst.operands.p << std::endl;
		
		emulator.handle_instruction(decinst);
	}
}