#include <vmprofiler.hpp>

namespace vm::util
{
    namespace reg
    {
        zydis_register_t to64( zydis_register_t reg )
        {
            return ZydisRegisterGetLargestEnclosing( ZYDIS_MACHINE_MODE_LONG_64, reg );
        }

        bool compare( zydis_register_t a, zydis_register_t b )
        {
            return to64( a ) == to64( b );
        }
    } // namespace reg

    bool get_fetch_operand( const zydis_routine_t &routine, zydis_instr_t &fetch_instr )
    {
        const auto result =
            std::find_if( routine.begin(), routine.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 == routine.end() )
            return false;

        fetch_instr = *result;
        return true;
    }

    std::optional< zydis_routine_t::iterator > get_fetch_operand( zydis_routine_t &routine )
    {
        auto result = std::find_if( routine.begin(), routine.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 == routine.end() )
            return {};

        return result;
    }

    void print( const zydis_decoded_instr_t &instr )
    {
        char buffer[ 256 ];
        ZydisFormatter formatter;
        ZydisFormatterInit( &formatter, ZYDIS_FORMATTER_STYLE_INTEL );
        ZydisFormatterFormatInstruction( &formatter, &instr, buffer, sizeof( buffer ), 0u );
        std::puts( buffer );
    }

    void print( zydis_routine_t &routine )
    {
        char buffer[ 256 ];
        ZydisFormatter formatter;
        ZydisFormatterInit( &formatter, ZYDIS_FORMATTER_STYLE_INTEL );

        for ( auto [ instr, raw, addr ] : routine )
        {
            ZydisFormatterFormatInstruction( &formatter, &instr, buffer, sizeof( buffer ), addr );
            std::printf( "> 0x%p %s\n", addr, buffer );
        }
    }

    bool is_jmp( const zydis_decoded_instr_t &instr )
    {
        switch ( instr.mnemonic )
        {
        case ZYDIS_MNEMONIC_JB:
        case ZYDIS_MNEMONIC_JBE:
        case ZYDIS_MNEMONIC_JCXZ:
        case ZYDIS_MNEMONIC_JECXZ:
        case ZYDIS_MNEMONIC_JKNZD:
        case ZYDIS_MNEMONIC_JKZD:
        case ZYDIS_MNEMONIC_JL:
        case ZYDIS_MNEMONIC_JLE:
        case ZYDIS_MNEMONIC_JMP:
        case ZYDIS_MNEMONIC_JNB:
        case ZYDIS_MNEMONIC_JNBE:
        case ZYDIS_MNEMONIC_JNL:
        case ZYDIS_MNEMONIC_JNLE:
        case ZYDIS_MNEMONIC_JNO:
        case ZYDIS_MNEMONIC_JNP:
        case ZYDIS_MNEMONIC_JNS:
        case ZYDIS_MNEMONIC_JNZ:
        case ZYDIS_MNEMONIC_JO:
        case ZYDIS_MNEMONIC_JP:
        case ZYDIS_MNEMONIC_JRCXZ:
        case ZYDIS_MNEMONIC_JS:
        case ZYDIS_MNEMONIC_JZ:
            return true;
        default:
            break;
        }
        return false;
    }

    bool flatten( zydis_routine_t &routine, std::uintptr_t routine_addr, bool keep_jmps )
    {
        ZydisDecoder decoder;
        zydis_decoded_instr_t instr;
        ZydisDecoderInit( &decoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_ADDRESS_WIDTH_64 );

        while ( ZYAN_SUCCESS(
            ZydisDecoderDecodeBuffer( &decoder, reinterpret_cast< void * >( routine_addr ), 0x1000, &instr ) ) )
        {
            // detect if we have already been at this instruction... if so that means there is a loop and we are
            // going to just return...
            if ( std::find_if( routine.begin(), routine.end(), [ & ]( const zydis_instr_t &zydis_instr ) -> bool {
                     return zydis_instr.addr == routine_addr;
                 } ) != routine.end() )
                return true;

            std::vector< u8 > raw_instr;
            raw_instr.insert( raw_instr.begin(), ( u8 * )routine_addr, ( u8 * )routine_addr + instr.length );

            if ( is_jmp( instr ) )
            {
                if ( instr.operands[ 0 ].type == ZYDIS_OPERAND_TYPE_REGISTER )
                {
                    routine.push_back( { instr, raw_instr, routine_addr } );
                    return true;
                }

                if ( keep_jmps )
                    routine.push_back( { instr, raw_instr, routine_addr } );

                ZydisCalcAbsoluteAddress( &instr, &instr.operands[ 0 ], routine_addr, &routine_addr );
            }
            else if ( instr.mnemonic == ZYDIS_MNEMONIC_RET )
            {
                routine.push_back( { instr, raw_instr, routine_addr } );
                return true;
            }
            else
            {
                routine.push_back( { instr, raw_instr, routine_addr } );
                routine_addr += instr.length;
            }
        }
        return false;
    }

    void deobfuscate( zydis_routine_t &routine )
    {
        static const auto _uses = []( ZydisDecodedOperand &op, zydis_register_t reg ) -> bool {
            switch ( op.type )
            {
            case ZYDIS_OPERAND_TYPE_MEMORY:
            {
                return reg::compare( op.mem.base, reg ) || reg::compare( op.mem.index, reg );
            }
            case ZYDIS_OPERAND_TYPE_REGISTER:
            {
                return reg::compare( op.reg.value, reg );
            }
            default:
                break;
            }
            return false;
        };

        static const auto _writes = []( zydis_decoded_instr_t &inst ) -> bool {
            for ( auto idx = 0; idx < inst.operand_count; ++idx )
                if ( inst.operands[ idx ].actions & ZYDIS_OPERAND_ACTION_MASK_WRITE )
                    return true;

            return false;
        };

        static const auto _remove = []( zydis_routine_t &routine, zydis_routine_t::iterator itr, zydis_register_t reg,
                                        u32 opcode_size ) -> void {
            for ( ; itr >= routine.begin(); --itr )
            {
                const auto instruction = &itr->instr;
                auto stop = false;

                if ( instruction->mnemonic == ZYDIS_MNEMONIC_JMP )
                    continue;

                for ( auto op_idx = 0u; op_idx < instruction->operand_count; ++op_idx )
                {
                    const auto op = &instruction->operands[ op_idx ];

                    if ( !_uses( *op, reg ) )
                        continue;

                    if ( op->type == ZYDIS_OPERAND_TYPE_MEMORY )
                    {
                        stop = true;
                        continue;
                    }

                    if ( opcode_size < 32 && op->size > opcode_size )
                        continue;

                    op->actions &ZYDIS_OPERAND_ACTION_MASK_WRITE ? op->actions &= ~ZYDIS_OPERAND_ACTION_MASK_WRITE
                                                                 : stop = true;
                }

                if ( !_writes( *instruction ) )
                    routine.erase( itr );

                else if ( stop )
                    break;
            }
        };

        for ( const auto &instr_data : routine )
        {
            if ( routine.empty() || routine.size() == 1 || instr_data.instr.mnemonic == ZYDIS_MNEMONIC_JMP )
                continue;

            for ( auto itr = routine.begin() + 1; itr != routine.end(); itr++ )
            {
                if ( itr->instr.mnemonic == ZYDIS_MNEMONIC_JMP || itr->instr.mnemonic == ZYDIS_MNEMONIC_RET )
                    break;

                // find the write operations that happen...
                for ( auto idx = 0u; idx < itr->instr.operand_count; ++idx )
                {
                    const auto op = &itr->instr.operands[ idx ];
                    // if its a read, continue to next opcode...
                    if ( op->actions & ZYDIS_OPERAND_ACTION_MASK_READ )
                        continue;

                    // if its not a write then continue to next opcode...
                    if ( !( op->actions & ZYDIS_OPERAND_ACTION_MASK_WRITE ) )
                        continue;

                    // if this operand is not a register then we continue...
                    if ( op->type != ZYDIS_OPERAND_TYPE_REGISTER )
                        continue;

                    // else we see if we can remove dead writes to this register...
                    _remove( routine, itr - 1, op->reg.value, op->size );
                }
            }
        }
    }
} // namespace vm::util