#include "emulator.hpp"

#include <boost/timer/timer.hpp>

template <>
struct SignedCounterpart<uint32_t> {
	typedef int32_t type;
};

template <>
struct SignedCounterpart<uint16_t> {
	typedef int16_t type;
};

template <>
struct SignedCounterpart<uint8_t> {
	typedef int8_t type;
};

Emulator::ValueType Emulator::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();
		}
	}
}

std::unique_ptr<Emulator::RefValueType> Emulator::get_first_operand(_DInst inst) {
	_Operand dest_op = inst.ops[0];
	std::unique_ptr<RefValueType> dest;
	
	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: {
			auto segment = register_getters.at(SEGMENT_GET(inst.segment))();
			auto segment_offset = ldt[boost::apply_visitor(return_visitor(), segment) >> 3].offset;
			auto register_contents = register_getters.at(dest_op.index)();
			int32_t offset = boost::apply_visitor(return_visitor(), register_contents) + inst.disp;
			std::cout << segment_offset << std::endl << offset << std::endl << inst.disp << std::endl;
			std::cout << segment_offset + offset + inst.disp << std::endl;
			auto memory_offset = segment_offset + offset + inst.disp;
			switch(dest_op.size) {
			case 8: {
				dest = std::make_unique<RefValueType>(*reinterpret_cast<uint8_t*>(&memory[memory_offset]));
				break;
			}
			case 16: {
				dest = std::make_unique<RefValueType>(*reinterpret_cast<uint16_t*>(&memory[memory_offset]));
				break;
			}
			case 32: {
				dest = std::make_unique<RefValueType>(*reinterpret_cast<uint32_t*>(&memory[memory_offset]));
				break;
			}
			default: {
				throw UnhandledInstruction();
			}
			}
			
			break;
		}
		case O_DISP: {
			auto segment = register_getters.at(SEGMENT_GET(inst.segment))();
			auto segment_offset = ldt[boost::apply_visitor(return_visitor(), segment) >> 3].offset;
			switch(dest_op.size) {
			case 8: {
				dest = std::make_unique<RefValueType>(*reinterpret_cast<uint8_t*>(&memory[segment_offset + inst.disp]));
				break;
			}
			case 16: {
				dest = std::make_unique<RefValueType>(*reinterpret_cast<uint16_t*>(&memory[segment_offset + inst.disp]));
				break;
			}
			case 32: {
				dest = std::make_unique<RefValueType>(*reinterpret_cast<uint32_t*>(&memory[segment_offset + inst.disp]));
				break;
			}
			default: {
				throw UnhandledInstruction();
			}
			}
			break;
		}
		default: {
			throw UnhandledInstruction();
		}
	}
	return dest;;
}

Emulator::ValueType Emulator::get_second_operand(_DInst inst, RefValueType dest) {
	ValueType src;
	_Operand src_op = inst.ops[1];
	
	switch(src_op.type) {
	case O_REG: {
		auto x = register_getters.at(src_op.index)();
		src = boost::apply_visitor(to_target_type(), dest, x);
		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;
	}
	case O_DISP: {
		auto segment = register_getters.at(SEGMENT_GET(inst.segment))();
		auto segment_offset = ldt[boost::apply_visitor(return_visitor(), segment) >> 3].offset;
		switch(boost::apply_visitor(dest_size(), dest)) {
		case 8: {
			if(segment_offset + inst.disp > memory.size()) {
				std::cerr << "Warning reading outside of memory region" << std::endl;
				src = static_cast<uint8_t>(0);
			} else {
				src = *reinterpret_cast<uint8_t*>(&memory[segment_offset + inst.disp]);
			}
			break;
		}
		case 16: {
			if(segment_offset + inst.disp + 1 > memory.size()) {
				std::cerr << "Warning reading outside of memory region" << std::endl;
				src = static_cast<uint16_t>(0);
			} else {
				src = *reinterpret_cast<uint16_t*>(&memory[segment_offset + inst.disp]);
			}
			break;
		}
		case 32: {
			if(segment_offset + inst.disp + 3 > memory.size()) {
				std::cerr << "Warning reading outside of memory region" << std::endl;
				src = static_cast<uint32_t>(0);
			} else {
				src = *reinterpret_cast<uint32_t*>(&memory[segment_offset + inst.disp]);
			}
			break;
		}
		default: {
			throw UnhandledInstruction();
		}
		}
		break;
	}
	case O_SMEM: {
		auto segment = register_getters.at(SEGMENT_GET(inst.segment))();
		auto segment_offset = ldt[boost::apply_visitor(return_visitor(), segment) >> 3].offset;
		auto register_contents = register_getters.at(src_op.index)();
		int32_t offset = boost::apply_visitor(return_visitor(), register_contents) + inst.disp;
		std::cout << segment_offset << std::endl << offset << std::endl << inst.disp << std::endl;
		std::cout << segment_offset + offset + inst.disp << std::endl;
		auto memory_offset = segment_offset + offset + inst.disp;
		switch(src_op.size) {
		case 8: {
			src = *reinterpret_cast<uint8_t*>(&memory[memory_offset]);
			break;
		}
		case 16: {
			src = *reinterpret_cast<uint16_t*>(&memory[memory_offset]);
			break;
		}
		case 32: {
			src = *reinterpret_cast<uint32_t*>(&memory[memory_offset]);
			break;
		}
		case 0: {
			assert(inst.opcode == I_LEA);
			switch(boost::apply_visitor(dest_size(), dest)) {
			case 8: {
				throw UnrecognizedInstruction();
			}
			case 16: {
				src = static_cast<uint16_t>(memory_offset);
				break;
			}
			case 32: {
				src = static_cast<uint32_t>(memory_offset);
				break;
			}
			default: {
				throw UnrecognizedInstruction();
			}
			}
			
			break;
		}
		default: {
			throw UnhandledInstruction();
		}
		}
		break;
	}
	default: {
		throw UnhandledInstruction();
		
	}
	}
	return src;
}

template <typename Ret>
struct op : public boost::static_visitor<Ret> {
	CpuState& cpu;
	
	op(CpuState& cpu)
	    : cpu(cpu)
	{}
};

template <typename Ret>
struct binary_op : public op<Ret> {
	using op<Ret>::op;
	
	template <typename T, typename U>
	Ret operator()(T&, U const&) const {
		throw Emulator::IncompatibleArguments();
	}
};

template <typename Ret>
struct unary_op : public op<Ret> {
	using op<Ret>::op;
};

struct and_ : public binary_op<void> {
	using binary_op::operator();
	using binary_op<void>::binary_op;
	
	template<typename T>
	void operator()(T& lh, T const& rh) const {
		lh &= rh;
	}
};

struct or_ : public binary_op<void> {
	using binary_op::operator();
	using binary_op<void>::binary_op;
	
	template<typename T>
	void operator()(T& lh, T const& rh) const {
		lh |= rh;
	}
};

struct mov : public binary_op<void> {
	using binary_op::operator();
	using binary_op<void>::binary_op;
	
	template<typename T>
	void operator()(T& lh, T const& rh) const {
		lh = rh;
	}
};

struct add : public binary_op<void> {
	using binary_op::operator();
	using binary_op<void>::binary_op;
	
	template <typename T>
	void operator()(T& lh, T const& rh) const {
		//TODO: flags
		
		lh += rh;
	}
};

struct sub : public binary_op<void> {
	using binary_op::operator();
	using binary_op<void>::binary_op;
	
	template<typename T>
	void operator()(T& lh, T const& rh) const {
		typedef typename SignedCounterpart<T>::type S;
		auto& slh = reinterpret_cast<S&>(lh);
		auto const& srh = reinterpret_cast<S const&>(rh);
		
		cpu.of() = ((srh < 0) && (slh > std::numeric_limits<S>::max() + srh)) || ((rh > 0) && (slh < std::numeric_limits<S>::min() + srh));
		cpu.af() = !(lh & 8) && (rh & 8);
		cpu.cf() = !(static_cast<T>(std::numeric_limits<S>::min()) & lh) && (static_cast<T>(std::numeric_limits<S>::min()) & rh);
		
		lh -= rh;
		
		cpu.sf() = static_cast<T>(std::numeric_limits<S>::min()) & lh;
		cpu.zf() = lh == 0;
		cpu.pf() = ((lh & 1) + (lh >> 1 & 1) + (lh >> 2 & 1) + (lh >> 3 & 1) + (lh >> 4 & 1) + (lh >> 5 & 1) + (lh >> 6 & 1) + (lh >> 7 & 1) % 2) + 1;
	}
};

struct shr : public binary_op<void> {
	using binary_op::operator();
	using binary_op<void>::binary_op;
	
	template <typename T>
	void operator()(T& lh, T const& rh) const {
		lh >>= rh;
	}
};

struct cmp : public binary_op<void> {
	using binary_op::operator();
	using binary_op<void>::binary_op;
	
	template<typename T>
	void operator()(T& lh, T const& rh) const {
		typedef typename SignedCounterpart<T>::type S;
		auto& slh = reinterpret_cast<S&>(lh);
		auto const& srh = reinterpret_cast<S const&>(rh);
		
		cpu.of() = ((srh < 0) && (slh > std::numeric_limits<S>::max() + srh)) || ((rh > 0) && (slh < std::numeric_limits<S>::min() + srh));
		cpu.af() = !(lh & 8) && (rh & 8);
		cpu.cf() = !(static_cast<T>(std::numeric_limits<S>::min()) & lh) && (static_cast<T>(std::numeric_limits<S>::min()) & rh);
		
		T result = lh - rh;
		
		cpu.sf() = static_cast<T>(std::numeric_limits<S>::min()) & result;
		cpu.zf() = result == 0;
		cpu.pf() = ((result & 1) + (result >> 1 & 1) + (result >> 2 & 1) + (result >> 3 & 1) + (result >> 4 & 1) + (result >> 5 & 1) + (result >> 6 & 1) + (result >> 7 & 1) % 2) + 1;
	}
};

struct lea : public binary_op<void> {
	using binary_op::operator();
	using binary_op<void>::binary_op;
	
	template<typename T>
	void operator()(T& lh, T const& rh) const {
		lh = rh;
	}
};

struct dec : public unary_op<void> {
	using unary_op<void>::unary_op;
	
	template<typename T>
	void operator()(T& op) const {
		op--;
	}
};

uint32_t Emulator::get_jmp_target(_DInst inst) {
	if(inst.ops[0].type == O_PC || inst.ops[0].type == O_PTR || inst.ops[0].type == O_DISP) {
		return INSTRUCTION_GET_TARGET(&inst);
	} else if (inst.ops[0].type == O_SMEM) {
		throw UnhandledInstruction();
	} else {
		throw UnrecognizedInstruction();
	}
}

void Emulator::handle_I_JMP(_DInst inst) {
	if(inst.ops[0].type == O_PC || inst.ops[0].type == O_PTR || inst.ops[0].type == O_DISP) {
		cpu.eip() = INSTRUCTION_GET_TARGET(&inst);
	} else if (inst.ops[0].type == O_SMEM) {
		throw UnhandledInstruction();
	} else {
		throw UnrecognizedInstruction();
	}
}

void Emulator::handle_I_STI(_DInst) {
	cpu.intf() = 1;
}

void Emulator::handle_I_AND(_DInst inst) {
	handle_binary<and_>(inst);
}

void Emulator::handle_I_MOV(_DInst inst) {
	handle_binary<mov>(inst);
}

void Emulator::handle_I_SUB(_DInst inst) {
	handle_binary<sub>(inst);
}

void Emulator::handle_I_SHR(_DInst inst) {
	handle_binary<shr>(inst);
}

void Emulator::handle_I_CMP(_DInst inst) {
	handle_binary<cmp>(inst);
}

void Emulator::handle_I_OR (_DInst inst) {
	handle_binary<or_>(inst);
}

void Emulator::handle_I_ADD(_DInst inst) {
	handle_binary<add>(inst);
}

void Emulator::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 Emulator::handle_I_JNZ(_DInst inst) {
	if(!cpu.zf()) {
		cpu.eip() = get_jmp_target(inst);
	}
}

void Emulator::handle_I_JZ(_DInst inst) {
	if(cpu.zf()) {
		cpu.eip() = get_jmp_target(inst);
	}
}

void Emulator::handle_I_JG(_DInst inst)
{
	if(!cpu.zf() && (cpu.sf() == cpu.of())) {
		cpu.eip() = get_jmp_target(inst);
	}
}

void Emulator::handle_I_PUSH(_DInst inst)
{
	switch(inst.ops[0].size) {
		case 8: {
            auto value = static_cast<uint8_t>(boost::apply_visitor(return_visitor(), register_getters.at(inst.ops[0].index)()));
            auto stack_offset = static_cast<uint8_t>(boost::apply_visitor(return_visitor(), register_getters.at(R_ESP)()));
	
			uint8_t& stack_var = *reinterpret_cast<uint8_t*>(&memory[stack_offset]);
			stack_var = value;
			cpu.esp() += 8;
			break;
		}
		case 16: {
            auto value = static_cast<uint16_t>(boost::apply_visitor(return_visitor(), register_getters.at(inst.ops[0].index)()));
            auto stack_offset = static_cast<uint16_t>(boost::apply_visitor(return_visitor(), register_getters.at(R_ESP)()));
			
			uint16_t& stack_var = *reinterpret_cast<uint16_t*>(&memory[stack_offset]);
			stack_var = value;
			cpu.esp() += 16;
			break;
		}
		case 32: {
            auto value = static_cast<uint32_t>(boost::apply_visitor(return_visitor(), register_getters.at(inst.ops[0].index)()));
            auto stack_offset = static_cast<uint32_t>(boost::apply_visitor(return_visitor(), register_getters.at(R_ESP)()));
			
			uint32_t& stack_var = *reinterpret_cast<uint32_t*>(&memory[stack_offset]);
			stack_var = value;
			cpu.esp() += 32;
			break;
		}
	}
}

void Emulator::handle_I_CLD(_DInst inst)
{
	cpu.df() = 0;
}

void Emulator::handle_I_SCAS(_DInst inst)
{
	if(FLAG_GET_PREFIX(inst.flags) == FLAG_REP) {
		while(cpu.ecx() != 0 && cpu.zf() != 0) {
			throw UnhandledInstruction();
		}
	} else {
		throw UnhandledInstruction();
	}
}

void Emulator::handle_I_LEA(_DInst inst)
{
	handle_binary<lea>(inst);
}

void Emulator::handle_I_DEC(_DInst inst)
{
	handle_unary<dec>(inst);
}

void Emulator::int_0x21() {
	switch(cpu.ah()) {
		//Dos Version
		case 0x30: {
			cpu.al() = 6;
			cpu.ah() = 1;
			cpu.bh() = 0;
			break;
		}
		//DOS4GW ident
		case 0xFF: {
			cpu.eax() = 0xFFFF3447;
			break;
		}
		default: {
			throw UnhandledInstruction();
		}
	}
}

void Emulator::int_0x31()
{
	switch(cpu.ax()) {
		case 0x0006: {
			auto entry = ldt.find(cpu.bx());
			if(entry != ldt.end()) {
				cpu.cf() = 0;
				cpu.cx() = static_cast<uint16_t>(entry->second.offset >> 16);
				cpu.dx() = static_cast<uint16_t>(entry->second.offset);
				cpu.ax() = 0;
			} else {
				cpu.cf() = 1;
				cpu.ax() = 0x8022;
			}
				break;
		}
		default: {
			throw UnhandledInstruction();
		}
	}
}

Emulator::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;
	
	ldt[1] = {0};
	cpu.ds() = ((uint16_t) 1 << 3) | (uint8_t) 1 << 2 | (uint8_t) 3;
	cpu.ss() = ((uint16_t) 1 << 3) | (uint8_t) 1 << 2 | (uint8_t) 3;
	
	ldt[2] = {cpu.esp()};
	cpu.es() = ((uint16_t) 2 << 3) | (uint8_t) 1 << 2 | (uint8_t) 3;

	cpu.fs() = 0; //dos4gw or dos32a specific segment register wont be used anyway

	cpu.gs() = 0;
	
#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 );
	REGISTER_HANDLER(I_SHR );
	REGISTER_HANDLER(I_CMP );
	REGISTER_HANDLER(I_JNZ );
	REGISTER_HANDLER(I_JZ  );
	REGISTER_HANDLER(I_JG  );
	REGISTER_HANDLER(I_OR  );
	REGISTER_HANDLER(I_ADD );
	REGISTER_HANDLER(I_PUSH);
	REGISTER_HANDLER(I_CLD );
	REGISTER_HANDLER(I_SCAS);
	REGISTER_HANDLER(I_LEA );
	REGISTER_HANDLER(I_DEC );
	
#undef REGISTER_HANDLER
	
#define REGISTER_GETTER(REGISTER, R_NAME) register_getters[REGISTER] = [this]{return RefValueType(cpu.R_NAME());}
	
	REGISTER_GETTER(R_AL , al );
	REGISTER_GETTER(R_AH , ah );
	REGISTER_GETTER(R_AX , ax );
	REGISTER_GETTER(R_EAX, eax);
	
	REGISTER_GETTER(R_BH,  bh );
	REGISTER_GETTER(R_BX,  bx );
	REGISTER_GETTER(R_EBX, ebx);
	
	REGISTER_GETTER(R_CL,  cl );
	REGISTER_GETTER(R_CX,  cx );
	REGISTER_GETTER(R_ECX, ecx);
	
	REGISTER_GETTER(R_DL , dl );
	REGISTER_GETTER(R_DX , dx );
	REGISTER_GETTER(R_EDX, edx);
	
	REGISTER_GETTER(R_ESP, esp);
	
	REGISTER_GETTER(R_ESI, esi);
	REGISTER_GETTER(R_EDI, edi);
	
	REGISTER_GETTER(R_DS , ds );
	REGISTER_GETTER(R_ES , es );
	REGISTER_GETTER(R_GS , gs );
	
#undef REGISTER_GETTER
	
#define REGISTER_INTERRUPT_HANDLER(INT) interrupt_handlers[INT] = [this]{return int_##INT();}
	
	REGISTER_INTERRUPT_HANDLER(0x21);
	REGISTER_INTERRUPT_HANDLER(0x31);
	
#undef REGISTER_INTERRUPT_HANDLER
}

void Emulator::set_code_segment(binparse::Offset32 offset) {
	ldt[0] = {offset};
	cpu.cs() = ((uint16_t) 0 << 3) | (uint8_t) 1 << 2 | (uint8_t) 3;
}

void Emulator::set_data_segment(binparse::Offset32 offset) {
	ldt[1] = {offset};
	cpu.ds() = ((uint16_t) 1 << 3) | (uint8_t) 1 << 2 | (uint8_t) 3;
	cpu.ss() = ((uint16_t) 1 << 3) | (uint8_t) 1 << 2 | (uint8_t) 3;
}

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);
	relocate(file, binary);
	
	_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);
	emulator.set_code_segment(file.object_table.entries.at(1).reloc_base_address);
	emulator.set_data_segment(file.object_table.entries.at(2).reloc_base_address);
	
	unsigned int decodedInstructionsCount;
	unsigned int emulated_instructions = 0;
	
	bool run = true;
	
	boost::timer::cpu_timer timer;
	timer.start();
	
	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); 
		
		_DecodedInst inst;
		distorm_format64(&ci, &decinst, &inst);
		
		std::cout << emulator.cpu << std::endl;
		std::cout << "CurrentInstruction: " << std::hex << std::setw(8) << std::setfill('0') << inst.offset << ":\t" << inst.mnemonic.p << " " << inst.operands.p << std::dec << std::endl;
		std::cout << std::endl << std::endl;
		
		emulator.cpu.eip() += decinst.size;
		
		run = emulator.handle_instruction(decinst);
		emulated_instructions++;
	}
	
	timer.stop();
	std::cout << "Serviced: " << emulated_instructions << " Instructions in: " << timer.elapsed().user << "ns" << std::endl;
}