#pragma once
#include <vmp2.hpp>
#include <transform.hpp>
#include <optional>

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

        std::optional< vmp2::exec_type_t > get_advancement( const 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 instrs

    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 table

        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 };
        } // namespace profile
    }     // namespace handler
} // namespace vm