#include <Windows.h>
#include <filesystem>
#include <fstream>
#include <iostream>

#include <cli-parser.hpp>
#include <vmp2.hpp>
#include <vmprofiler.hpp>
#include <xtils.hpp>
#include <ZydisExportConfig.h>

int __cdecl main( int argc, const char *argv[] )
{
    argparse::argument_parser_t parser( "vmprofiler-cli", "virtual instruction pseudo code generator" );

    parser.add_argument()
        .names( { "--bin", "--vmpbin" } )
        .description( "unpacked binary protected with VMProtect 2" )
        .required( true );

    parser.add_argument()
        .names( { "--vmentry", "--entry" } )
        .description( "rva to push prior to a vm_entry" )
        .required( true );

    parser.add_argument().name( "--showhandlers" ).description( "show all vm handlers..." );

    parser.enable_help();
    auto err = parser.parse( argc, argv );

    if ( err )
    {
        std::cout << err << std::endl;
        return -1;
    }

    if ( parser.exists( "help" ) )
    {
        parser.print_help();
        return 0;
    }

    const auto module_base = reinterpret_cast< std::uintptr_t >(
        LoadLibraryExA( parser.get< std::string >( "bin" ).c_str(), NULL, DONT_RESOLVE_DLL_REFERENCES ) );

    const auto vm_entry_ptr =
        module_base + std::strtoull( parser.get< std::string >( "vmentry" ).c_str(), nullptr, 16 );

    const auto image_base = xtils::um_t::get_instance()->image_base( parser.get< std::string >( "bin" ).c_str() );

    zydis_routine_t vm_entry, calc_jmp;
    std::printf( "> vm entry start = 0x%p\n", vm_entry_ptr );

    if ( !vm::util::flatten( vm_entry, vm_entry_ptr ) )
    {
        std::printf( "> failed to flatten vm entry...\n" );
        return -1;
    }

    vm::util::deobfuscate( vm_entry );
    std::printf( "> flattened vm entry...\n" );
    std::printf( "> deobfuscated vm entry...\n" );
    std::printf( "==================================================================================\n" );
    vm::util::print( vm_entry );

    if ( !vm::calc_jmp::get( vm_entry, calc_jmp ) )
    {
        std::printf( "> failed to extract calc_jmp from vm_entry...\n" );
        return -1;
    }

    vm::util::deobfuscate( calc_jmp );
    std::printf( "> calc_jmp extracted from vm_entry... calc_jmp:\n" );
    std::printf( "==================================================================================\n" );
    vm::util::print( calc_jmp );
    std::printf( "==================================================================================\n" );

    const auto vm_handler_table = vm::handler::table::get( vm_entry );

    if ( !vm_handler_table )
    {
        std::printf( "> failed to locate vm handler table...\n" );
        return -1;
    }

    auto advancement = vm::calc_jmp::get_advancement( calc_jmp );
    if ( advancement.has_value() )
        std::printf( "> virtual instruction pointer advancement: %s\n",
                     advancement == vmp2::exec_type_t::forward ? "forward" : "backward" );
    else
        std::printf( "> virtual instruction pointer advancement was unable to be parsed!\n" );

    std::printf( "> located vm handler table... at = 0x%p, rva = 0x%p\n", vm_handler_table,
                 ( reinterpret_cast< std::uintptr_t >( vm_handler_table ) - module_base ) + image_base );

    zydis_decoded_instr_t vm_handler_transform;
    if ( !vm::handler::table::get_transform( vm_entry, &vm_handler_transform ) )
    {
        std::printf( "[!] failed to locate vm handler table entry transformation...\n" );
        return -1;
    }

    std::printf( "> vm handler table entries decrypted with = " );
    vm::util::print( vm_handler_transform );

    vm_handler_transform.mnemonic = vm::transform::inverse[ vm_handler_transform.mnemonic ];
    std::printf( "> vm handler table entries encrypted with = " );
    vm::util::print( vm_handler_transform );

    std::printf( "==================================================================================\n" );
    std::vector< zydis_decoded_instr_t > vinstr_rva_decrypt_instrs;
    if ( !vm::instrs::get_rva_decrypt( vm_entry, vinstr_rva_decrypt_instrs ) )
    {
        std::printf( "[!] failed to get virtual instruction rva decrypt transformations...\n" );
        return -1;
    }

    std::printf( "> virtual instruction rva decryption instructions:\n" );
    for ( auto &transform : vinstr_rva_decrypt_instrs )
    {
        std::printf( "\t" );
        vm::util::print( transform );
    }
    std::printf( "==================================================================================\n" );

    std::vector< vm::handler::handler_t > vm_handlers;
    if ( !vm::handler::get_all( module_base, image_base, vm_entry, vm_handler_table, vm_handlers ) )
    {
        std::printf( "> failed to get all vm handler meta data...\n" );
        return -1;
    }

    if ( parser.exists( "showhandlers" ) )
    {
        for ( auto idx = 0u; idx < vm_handlers.size(); ++idx )
        {
            auto vm_handler = vm_handlers[ idx ];
            std::printf( "==========[vm handler %s, idx = %d, imm size = %d]========\n",
                         vm_handler.profile ? vm_handler.profile->name : "UNKNOWN", idx, vm_handler.imm_size );

            std::printf( "================[vm handler instructions]==============\n" );
            vm::util::print( vm_handler.instrs );

            if ( vm_handler.imm_size )
            {
                std::printf( "=================[vm handler transforms]===============\n" );
                for ( const auto &[ transform_type, transform ] : vm_handler.transforms )
                    vm::util::print( transform );
            }

            std::printf( "=======================================================\n\n" );
        }
    }
    std::printf( "> finished...\n" );
}