#pragma once

#include "cpustate.hpp"

#include "le_parse_util.hpp"

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

#include <boost/variant.hpp>

#include <fstream>
#include <iomanip>
#include <limits>

template <typename T>
struct SignedCounterpart {};

template <>
struct SignedCounterpart<uint32_t>;

template <>
struct SignedCounterpart<uint16_t>;

template <>
struct SignedCounterpart<uint8_t>;

struct LDTEntry {
	uint32_t offset;
};

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;
	
	std::map<int, LDTEntry> ldt;
	
	ValueType extract_immediate_of_size(int size, _DInst inst);
	
	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 to_target_type : public boost::static_visitor<ValueType> {
		template <typename Dest, typename Src>
		ValueType operator()(Dest const& d, Src const& s) const {
			typename SignedCounterpart<Src>::type ss(s);
			typename SignedCounterpart<Dest>::type sd(ss);
			return static_cast<Dest>(sd);
		}
	};
	
	struct IncompatibleArguments : public std::runtime_error {
		IncompatibleArguments() 
		: std::runtime_error("The two arguments to this instruction are not compatible")
		{}
	};
	
	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: {
					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();
				}
			}
		}

		{
			_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;
						}
						default: {
							throw UnhandledInstruction();
						}
					}
					break;
				}
				default: {
					throw UnhandledInstruction();
				}
			}
		}
		
		boost::apply_visitor(O(cpu), *dest, src);
	}
	
	uint32_t get_jmp_target(_DInst inst);
	
	void handle_I_JMP (_DInst inst);
	void handle_I_STI (_DInst);
	void handle_I_AND (_DInst inst);
	void handle_I_MOV (_DInst inst);
	void handle_I_SUB (_DInst inst);
	void handle_I_SHR (_DInst inst);
	void handle_I_CMP (_DInst inst);
	void handle_I_INT (_DInst inst);
	void handle_I_JNZ (_DInst inst);
	void handle_I_JZ  (_DInst inst);
	void handle_I_JG  (_DInst inst);
	void handle_I_OR  (_DInst inst);
	void handle_I_ADD (_DInst inst);
	void handle_I_PUSH(_DInst inst);
	void handle_I_CLD (_DInst inst);
	
	void int_0x21();
	void int_0x31();

public:
	Emulator(binparse::Offset32 init_eip, binparse::Offset32 init_esp, std::vector<uint8_t>& memory);
	
	void set_code_segment(binparse::Offset32 offset);
	void set_data_segment(binparse::Offset32 offset);
	
	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")
		{}
	};
	
	bool handle_instruction(_DInst inst) {
		try {
			opcode_handlers.at(inst.opcode)(inst);
		} catch(std::exception const& e) {
			std::cerr << "Encountered fatal error: " << e.what() << std::endl;
			return false;
		}
		return true;
	}
	
};

void emulate(std::string file_path);