#pragma once
#include <transform.hpp>

namespace vm
{
	namespace calc_jmp
	{
		bool get(const zydis_routine_t& vm_entry, zydis_routine_t& calc_jmp);
	}

	namespace instrs
	{
		// decrypt transformations for encrypted virtual instruction rva...
		bool get_rva_decrypt(const zydis_routine_t& vm_entry,
			std::vector<zydis_decoded_instr_t>& transform_instrs);

		std::pair<std::uint64_t, std::uint64_t> decrypt_operand(transform::map_t& transforms,
			std::uint64_t operand, std::uint64_t rolling_key);

		std::pair<std::uint64_t, std::uint64_t> encrypt_operand(transform::map_t& transforms,
			std::uint64_t operand, std::uint64_t rolling_key);
	}

	namespace handler
	{
		using instr_callback_t = bool(*)(const zydis_decoded_instr_t& instr);

		enum mnemonic_t
		{
			INVALID,
			PUSHVSP,
			SHRQ,
			MULQ,
			DIVQ,
			JMP,
			VMEXIT,

			SREGQ,
			SREGDW,
			SREGW,

			LREGQ,
			LREGDW,

			LCONSTQ,
			LCONSTBZXW,
			LCONSTBSXDW,
			LCONSTDWSXQ,
			LCONSTWSXQ,
			LCONSTDW,

			READQ,
			READDW,
			READW,

			WRITEQ,
			WRITEDW,
			WRITEW,

			ADDQ,
			ADDDW,

			SHLQ,
			SHLDW,

			NANDQ,
			NANDDW
		};

        enum extention_t
        {
            none,
            sign_extend,
            zero_extend
        };

        struct profile_t
        {
            const char* name;
			mnemonic_t mnemonic;
            u8 imm_size;
			std::vector<instr_callback_t> signature;
            extention_t extention;
        };

		struct handler_t
		{
			u8 imm_size; // size in bits...
			vm::transform::map_t transforms;
			vm::handler::profile_t* profile;
			zydis_routine_t instrs;
			std::uintptr_t address;
		};

		bool has_imm(const zydis_routine_t& vm_handler);
		std::uint8_t imm_size(const zydis_routine_t& vm_handler);
		bool get(zydis_routine_t& vm_entry, zydis_routine_t& vm_handler, std::uintptr_t handler_addr);

		// may throw an exception...
		bool get_all(std::uintptr_t module_base, std::uintptr_t image_base,
			zydis_routine_t& vm_entry, std::uintptr_t* vm_handler_table, std::vector<handler_t>& vm_handlers);

		// can be used on calc_jmp...
		bool get_operand_transforms(const zydis_routine_t& vm_handler, transform::map_t& transforms);
		vm::handler::profile_t* get_profile(handler_t& vm_handler);

		namespace table
		{
			std::uintptr_t* get(const zydis_routine_t& vm_entry);
			bool get_transform(const zydis_routine_t& vm_entry, zydis_decoded_instr_t* transform_instr);

			std::uint64_t encrypt(zydis_decoded_instr_t& transform_instr, std::uint64_t val);
			std::uint64_t decrypt(zydis_decoded_instr_t& transform_instr, std::uint64_t val);
		}

		namespace profile
		{
			extern vm::handler::profile_t sregq;
			extern vm::handler::profile_t sregdw;
			extern vm::handler::profile_t sregw;

			extern vm::handler::profile_t lregq;
			extern vm::handler::profile_t lregdw;

			extern vm::handler::profile_t lconstq;
			extern vm::handler::profile_t lconstbzxw;
			extern vm::handler::profile_t lconstbsxdw;
			extern vm::handler::profile_t lconstdwsxq;
			extern vm::handler::profile_t lconstwsxq;
			extern vm::handler::profile_t lconstdw;

			extern vm::handler::profile_t addq;
			extern vm::handler::profile_t adddw;

			extern vm::handler::profile_t shlq;
			extern vm::handler::profile_t shldw;

			extern vm::handler::profile_t nandq;
			extern vm::handler::profile_t nanddw;

			extern vm::handler::profile_t writeq;
			extern vm::handler::profile_t writedw;

			extern vm::handler::profile_t shrq;
			extern vm::handler::profile_t pushvsp;
			extern vm::handler::profile_t mulq;
			extern vm::handler::profile_t divq;
			extern vm::handler::profile_t jmp;
			extern vm::handler::profile_t readq;
			extern vm::handler::profile_t vmexit;

			inline std::vector<vm::handler::profile_t*> all =
			{
				&sregq, &sregdw, &sregw,
				&lregq, &lregdw,
				&lconstq, &lconstbzxw, &lconstbsxdw, &lconstdwsxq, &lconstwsxq, &lconstdw,
				&addq, &adddw,
				&shlq, &shldw,
				&writeq, &writedw,
				&nandq, &nanddw,

				&shrq,
				&readq,
				&mulq,
				&pushvsp,
				&divq,
				&jmp,
				&vmexit
			};
		}
	}
}