#include <vmprofiler.hpp>

namespace vm
{
    namespace calc_jmp
    {
        bool get( const zydis_routine_t &vm_entry, zydis_routine_t &calc_jmp )
        {
            auto result =
                std::find_if( vm_entry.begin(), vm_entry.end(), []( const zydis_instr_t &instr_data ) -> bool {
                    // mov/movsx/movzx rax/eax/ax/al, [rsi]
                    return ( instr_data.instr.operand_count > 1 &&
                             ( instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOV ||
                               instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOVSX ||
                               instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOVZX ) &&
                             instr_data.instr.operands[ 0 ].type == ZYDIS_OPERAND_TYPE_REGISTER &&
                             util::reg::to64( instr_data.instr.operands[ 0 ].reg.value ) == ZYDIS_REGISTER_RAX &&
                             instr_data.instr.operands[ 1 ].type == ZYDIS_OPERAND_TYPE_MEMORY &&
                             instr_data.instr.operands[ 1 ].mem.base == ZYDIS_REGISTER_RSI );
                } );

            if ( result == vm_entry.end() )
                return false;

            calc_jmp.insert( calc_jmp.end(), result, vm_entry.end() );
            return true;
        }

        std::optional< vmp2::exec_type_t > get_advancement( const zydis_routine_t &calc_jmp )
        {
            auto result =
                std::find_if( calc_jmp.begin(), calc_jmp.end(), []( const zydis_instr_t &instr_data ) -> bool {
                    // look for any instruction with RSI being the first operand...
                    return instr_data.instr.operands[ 0 ].type == ZYDIS_OPERAND_TYPE_REGISTER &&
                           instr_data.instr.operands[ 0 ].reg.value == ZYDIS_REGISTER_RSI;
                } );

            if ( result == calc_jmp.end() )
                return {};

            const auto instr = &result->instr;

            switch ( instr->mnemonic )
            {
            case ZYDIS_MNEMONIC_LEA:
            {
                if ( instr->operands[ 1 ].type == ZYDIS_OPERAND_TYPE_MEMORY )
                {
                    if ( instr->operands[ 1 ].mem.disp.value > 0 )
                        return vmp2::exec_type_t::forward;
                    else
                        return vmp2::exec_type_t::backward;
                }
                // else we dont know what we are looking at...
                return {};
            }
            case ZYDIS_MNEMONIC_ADD:
            {
                // ADD RSI, IMM...
                if ( instr->operands[ 1 ].type == ZYDIS_OPERAND_TYPE_IMMEDIATE )
                {
                    // see if IMM is negitive...
                    if ( instr->operands[ 1 ].imm.value.s > 0 )
                        return vmp2::exec_type_t::forward;
                    else // adding a negitive number is sub...
                        return vmp2::exec_type_t::backward;
                }
                // else we dont know what we are looking at...
                return {};
            }
            case ZYDIS_MNEMONIC_SUB:
            {
                // SUB RSI, IMM...
                if ( instr->operands[ 1 ].type == ZYDIS_OPERAND_TYPE_IMMEDIATE )
                {
                    // see if IMM is negitive...
                    if ( instr->operands[ 1 ].imm.value.s > 0 )
                        return vmp2::exec_type_t::backward;
                    else // subtracting a negitive number means you are adding...
                        return vmp2::exec_type_t::forward;
                }
                // else we dont know what we are looking at...
                return {};
            }
            case ZYDIS_MNEMONIC_INC:
                return vmp2::exec_type_t::forward;
            case ZYDIS_MNEMONIC_DEC:
                return vmp2::exec_type_t::backward;
            default:
                break;
            }
            return {};
        }
    } // namespace calc_jmp
} // namespace vm