#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")
		{}
	};
	
	std::unique_ptr<RefValueType> get_first_operand(_DInst inst);
	
	ValueType get_second_operand(_DInst inst, RefValueType dest);
	
	template <typename O>
	void handle_binary(_DInst inst) {
		std::unique_ptr<RefValueType> dest = get_first_operand(inst);
		ValueType src = get_second_operand(inst, *dest);
		
		boost::apply_visitor(O(cpu), *dest, src);
	}
	
	template<typename O> 
	void handle_unary(_DInst inst) {
		std::unique_ptr<RefValueType> op = get_first_operand(inst);
		
		boost::apply_visitor(O(cpu), *op);
	}
	
	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 handle_I_SCAS(_DInst inst);
	void handle_I_LEA (_DInst inst);
	void handle_I_DEC (_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 unrecognized 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);