#include "qvminspector.h"

qvminspector_t::qvminspector_t( qwidget_t *parent ) : qmain_window_t( parent ), file_header( nullptr ), vmctx( nullptr )
{
    ui.setupUi( this );
    ui.virt_instrs->setColumnWidth( 0, 180 );
    ui.virt_instrs->setColumnWidth( 1, 150 );
    ui.virt_instrs->setColumnWidth( 2, 190 );
    ui.virt_instrs->setColumnWidth( 3, 200 );

    connect( ui.action_open, &QAction::triggered, this, &qvminspector_t::on_open );
    connect( ui.action_close, &QAction::triggered, this, &qvminspector_t::on_close );
}

void qvminspector_t::on_close()
{
    exit( 0 );
}

void qvminspector_t::on_open()
{
    if ( file_header && vmctx )
    {
        free( file_header );
        delete vmctx;
    }

    file_header = nullptr, first_block = nullptr;
    code_block_addrs.clear();
    image_base = 0u, vm_entry_rva = 0u, module_base = 0u;

    file_path = QFileDialog::getOpenFileName(
        this, tr( "open vmp2 file" ), std::filesystem::current_path().string().c_str(), tr( "vmp2 file (*.vmp2)" ) );

    const auto &_file_path = file_path.toStdString();

    if ( file_path.isEmpty() )
    {
        dbg_msg( "invalid vmp2 file... no file selected..." );
        return;
    }

    if ( !std::filesystem::exists( _file_path ) )
    {
        dbg_msg( "vmp2 file does not exist..." );
        return;
    }

    const auto file_size = std::filesystem::file_size( _file_path );

    if ( !file_size )
    {
        dbg_msg( "invalid vmp2 file size..." );
        return;
    }

    qfile_t open_file( file_path );
    file_header = reinterpret_cast< vmp2::v3::file_header * >( malloc( file_size ) );

    if ( !open_file.open( QIODevice::ReadOnly ) )
    {
        dbg_msg( "failed to open vmp2 file..." );
        return;
    }

    memcpy( file_header, open_file.readAll().data(), file_size );

    if ( !init_data() )
    {
        dbg_msg( "failed to init vmp2 file data..." );
        return;
    }

    update_ui();
}

void qvminspector_t::dbg_print( qstring_t dbg_output )
{
    ui.dbg_output_window->appendPlainText( dbg_output );
}

void qvminspector_t::dbg_msg( qstring_t dbg_output )
{
    qmsg_box_t msg_box;
    msg_box.setText( dbg_output );
    msg_box.exec();
    dbg_print( dbg_output );
}

bool qvminspector_t::init_data()
{
    if ( file_header->magic != VMP_MAGIC )
    {
        dbg_msg( "invalid magic bytes for vmp2 file..." );

        return false;
    }

    dbg_print( "valid magic bytes for vmp2 file..." );

    if ( file_header->version != vmp2::version_t::v3 )
    {
        dbg_msg( "invalid vmp2 file version... "
                 "this vminspector is compiled for version 3...\n" );

        return false;
    }

    vm_entry_rva = file_header->vm_entry_rva;
    image_base = file_header->image_base;
    image_size = file_header->module_size;

    module_base = reinterpret_cast< std::uintptr_t >( file_header ) + file_header->module_offset;
    vmctx = new vm::ctx_t( module_base, image_base, image_size, vm_entry_rva );

    if ( !vmctx->init() )
    {
        dbg_msg( "failed to init vm::ctx_t... this can happen for many reasons..."
                 "ensure that the vmp2 file is not corrupted...\n" );

        return false;
    }

    first_block = reinterpret_cast< vmp2::v3::code_block_t * >( reinterpret_cast< std::uintptr_t >( file_header ) +
                                                                file_header->code_block_offset );
    return true;
}

void qvminspector_t::add_branch_children( qtree_widget_item_t *item, std::uintptr_t branch_addr )
{
    if ( std::find( code_block_addrs.begin(), code_block_addrs.end(), branch_addr ) != code_block_addrs.end() )
        return;

    code_block_addrs.push_back( branch_addr );

    // for each code block find the one that starts with the desired branch...
    for ( auto [ code_block, code_block_num ] = std::tuple{ first_block, 0u };
          code_block_num < file_header->code_block_count;
          code_block = reinterpret_cast< vmp2::v3::code_block_t * >( reinterpret_cast< std::uintptr_t >( code_block ) +
                                                                     code_block->next_block_offset ),
                                            ++code_block_num )
    {
        if ( code_block->vip_begin != branch_addr )
            continue;

        dbg_print( qstring_t( "> code block %1 (block_%2), number of vinstrs = %3..." )
                       .arg( code_block_num )
                       .arg( code_block->vip_begin, 0, 16 )
                       .arg( code_block->vinstr_count ) );

        // for each virtual instruction inside of this code block...
        for ( auto idx = 0u; idx < code_block->vinstr_count; ++idx )
        {
            const auto virt_instr = &code_block->vinstr[ idx ];
            const auto profile = vm::handler::get_profile( virt_instr->mnemonic_t );
            auto virt_instr_entry = new qtree_widget_item_t;

            // virtual instruction image base'ed rva... (column 1)...
            virt_instr_entry->setText(
                0, QString( "block_%1+%2" )
                       .arg( branch_addr, 0, 16 )
                       .arg( idx ? code_block->vip_begin >
                                           ( ( virt_instr->trace_data.vip - file_header->module_base ) + image_base )
                                       ? code_block->vip_begin -
                                             ( ( virt_instr->trace_data.vip - file_header->module_base ) + image_base )
                                       : ( ( virt_instr->trace_data.vip - file_header->module_base ) + image_base ) -
                                             code_block->vip_begin
                                 : 0u,
                             0, 16 ) );

            // virtual instruction operand bytes... (column 2)...
            qstring_t operand_bytes;
            operand_bytes.append( QString( "%1" ).arg( virt_instr->opcode, 0, 16 ) );

            // if virt instruction has an imm... grab its bytes...
            if ( virt_instr->operand.has_imm )
            {
                operand_bytes.append( " - " );
                for ( auto _idx = 0u; _idx < virt_instr->operand.imm.imm_size / 8; ++_idx )
                    operand_bytes.append( QString( "%1 " ).arg(
                        reinterpret_cast< std::uint8_t * >( &virt_instr->operand.imm.u )[ _idx ], 0, 16 ) );
            }

            virt_instr_entry->setText( 1, operand_bytes );

            // virtual instruction string, includes imm... (colume 3)...
            QString decoded_instr( QString( "%1" ).arg(
                profile ? profile->name : QString( "UNK(%1)" ).arg( virt_instr->opcode, 0, 16 ) ) );

            if ( virt_instr->operand.has_imm ) // if there is a second operand (imm) put its value...
                decoded_instr.append( QString( " %1" ).arg( virt_instr->operand.imm.u, 0, 16 ) );

            virt_instr_entry->setText( 2, decoded_instr );

            // add comments to the virtual instruction... (colume 4)...
            if ( virt_instr->mnemonic_t == vm::handler::LREGQ || virt_instr->mnemonic_t == vm::handler::SREGQ )
                virt_instr_entry->setText(
                    3,
                    QString( "; vreg%1" ).arg( virt_instr->operand.imm.u ? ( virt_instr->operand.imm.u / 8 ) : 0u ) );

            QVariant var;
            var.setValue( virt_instr );
            virt_instr_entry->setData( 3, Qt::UserRole, var );

            if ( virt_instr->mnemonic_t == vm::handler::JMP && code_block->jcc.has_jcc )
            {
                virt_instr_entry->setText( 3, QString( "; { %1, %2 }" )
                                                  .arg( code_block->jcc.block_addr[ 0 ], 0, 16 )
                                                  .arg( code_block->jcc.block_addr[ 1 ], 0, 16 ) );

                auto branch_entry1 = new qtree_widget_item_t(), branch_entry2 = new qtree_widget_item_t();
                const auto block1_addr = code_block->jcc.block_addr[ 0 ];
                const auto block2_addr = code_block->jcc.block_addr[ 1 ];

                branch_entry1->setText( 0, QString( "%1" ).arg( block1_addr, 0, 16 ) );
                branch_entry1->setText( 3, QString( "; block_%1" ).arg( block1_addr, 0, 16 ) );

                branch_entry2->setText( 0, QString( "%1" ).arg( block2_addr, 0, 16 ) );
                branch_entry2->setText( 3, QString( "; block_%1" ).arg( block2_addr, 0, 16 ) );

                add_branch_children( branch_entry1, code_block->jcc.block_addr[ 0 ] );
                add_branch_children( branch_entry2, code_block->jcc.block_addr[ 1 ] );

                virt_instr_entry->addChildren( { branch_entry1, branch_entry2 } );
                // if its a JMP with branches we want to insert the next code block
                // instructions into the child widget entries...
                item->addChild( virt_instr_entry );
                return;
            }
            else if ( virt_instr->mnemonic_t == vm::handler::JMP )
            {
                virt_instr_entry->setText( 3, QString( "; { %1 }" ).arg( code_block->jcc.block_addr[ 0 ], 0, 16 ) );

                // else if this jmp doesnt have two branches add the next code block to it...
                item->addChild( virt_instr_entry );
                add_branch_children( item, code_block->jcc.block_addr[ 0 ] );
            }
            item->addChild( virt_instr_entry );
        }
    }
}

void qvminspector_t::update_ui()
{
    // add vm handlers to the vm handler tree...
    ui.virt_handlers_tree->clear();
    for ( auto idx = 0u; idx < vmctx->vm_handlers.size(); ++idx )
    {
        auto new_handler_entry = new qtree_widget_item_t;
        new_handler_entry->setData( 0, Qt::UserRole, idx );
        new_handler_entry->setText( 0, QString( "%1" ).arg( idx ) );
        new_handler_entry->setText(
            1, QString( "%1" ).arg(
                   ABS_TO_IMG( vmctx->vm_handlers[ idx ].address, module_base, file_header->image_base ), 0, 16 ) );

        new_handler_entry->setText( 2, vmctx->vm_handlers[ idx ].profile ? vmctx->vm_handlers[ idx ].profile->name
                                                                         : "UNDEFINED" );

        new_handler_entry->setText( 3, QString( "%1" ).arg( vmctx->vm_handlers[ idx ].imm_size ) );

        if ( vmctx->vm_handlers[ idx ].profile && vmctx->vm_handlers[ idx ].imm_size )
            new_handler_entry->setText( 4, vmctx->vm_handlers[ idx ].profile->extention ==
                                                   vm::handler::extention_t::sign_extend
                                               ? "SIGN EXTENDED"
                                               : "ZERO EXTENDED" );
        else
            new_handler_entry->setText( 4, "UNDEFINED" );

        ui.virt_handlers_tree->addTopLevelItem( new_handler_entry );
    }
    ui.virt_handlers_tree->topLevelItem( 0 )->setSelected( true );

    // for each code block insert their virtual instructions
    // into the virtual instruction tree... also put meta data about the code
    // block above the virtual instructions... if the code block has a JCC (with two branches)
    // then make a child repeating this for loop...

    ui.virt_instrs->clear(); // clear old virtual instructions out since we are updating the UI...
    for ( auto [ code_block, code_block_num ] = std::tuple{ first_block, 0u };
          code_block_num < file_header->code_block_count;
          code_block = reinterpret_cast< vmp2::v3::code_block_t * >( reinterpret_cast< std::uintptr_t >( code_block ) +
                                                                     code_block->next_block_offset ),
                                            ++code_block_num )
    {
        code_block_addrs.push_back( code_block->vip_begin );
        dbg_print( qstring_t( "> code block %1 (block_%2), number of vinstrs = %3..." )
                       .arg( code_block_num )
                       .arg( code_block->vip_begin, 0, 16 )
                       .arg( code_block->vinstr_count ) );

        // for each virtual instruction inside of this code block...
        for ( auto idx = 0u; idx < code_block->vinstr_count; ++idx )
        {
            const auto virt_instr = &code_block->vinstr[ idx ];
            const auto profile = vm::handler::get_profile( virt_instr->mnemonic_t );
            auto virt_instr_entry = new qtree_widget_item_t;

            // virtual instruction image base'ed rva... (column 1)...
            virt_instr_entry->setText(
                0, QString( "block_%1+%2" )
                       .arg( code_block->vip_begin, 0, 16 )
                       .arg( idx ? code_block->vip_begin >
                                           ( ( virt_instr->trace_data.vip - file_header->module_base ) + image_base )
                                       ? code_block->vip_begin -
                                             ( ( virt_instr->trace_data.vip - file_header->module_base ) + image_base )
                                       : ( ( virt_instr->trace_data.vip - file_header->module_base ) + image_base ) -
                                             code_block->vip_begin
                                 : 0u,
                             0, 16 ) );

            // virtual instruction operand bytes... (column 2)...
            qstring_t operand_bytes;
            operand_bytes.append( QString( "%1" ).arg( virt_instr->opcode, 0, 16 ) );

            // if virt instruction has an imm... grab its bytes...
            if ( virt_instr->operand.has_imm )
            {
                operand_bytes.append( " - " );
                for ( auto _idx = 0u; _idx < virt_instr->operand.imm.imm_size / 8; ++_idx )
                    operand_bytes.append( QString( "%1 " ).arg(
                        reinterpret_cast< std::uint8_t * >( &virt_instr->operand.imm.u )[ _idx ], 0, 16 ) );
            }

            virt_instr_entry->setText( 1, operand_bytes );

            // virtual instruction string, includes imm... (colume 3)...
            QString decoded_instr( QString( "%1" ).arg(
                profile ? profile->name : QString( "UNK(%1)" ).arg( virt_instr->opcode, 0, 16 ) ) );

            if ( virt_instr->operand.has_imm ) // if there is a second operand (imm) put its value...
                decoded_instr.append( QString( " %1" ).arg( virt_instr->operand.imm.u, 0, 16 ) );

            virt_instr_entry->setText( 2, decoded_instr );

            // add comments to the virtual instruction... (colume 4)...
            if ( virt_instr->mnemonic_t == vm::handler::LREGQ || virt_instr->mnemonic_t == vm::handler::SREGQ )
                virt_instr_entry->setText(
                    3,
                    QString( "; vreg%1" ).arg( virt_instr->operand.imm.u ? ( virt_instr->operand.imm.u / 8 ) : 0u ) );

            // add virt_instr_t pointer for this current virtual instruction...
            QVariant var;
            var.setValue( virt_instr );
            virt_instr_entry->setData( 3, Qt::UserRole, var );

            if ( virt_instr->mnemonic_t == vm::handler::JMP && code_block->jcc.has_jcc )
            {
                virt_instr_entry->setText( 3, QString( "; { %1, %2 }" )
                                                  .arg( code_block->jcc.block_addr[ 0 ], 0, 16 )
                                                  .arg( code_block->jcc.block_addr[ 1 ], 0, 16 ) );

                auto branch_entry1 = new qtree_widget_item_t(), branch_entry2 = new qtree_widget_item_t();
                const auto block1_addr = code_block->jcc.block_addr[ 0 ];
                const auto block2_addr = code_block->jcc.block_addr[ 1 ];

                branch_entry1->setText( 0, QString( "%1" ).arg( block1_addr, 0, 16 ) );
                branch_entry1->setText( 3, QString( "; block_%1" ).arg( block1_addr, 0, 16 ) );

                branch_entry2->setText( 0, QString( "%1" ).arg( block2_addr, 0, 16 ) );
                branch_entry2->setText( 3, QString( "; block_%1" ).arg( block2_addr, 0, 16 ) );

                add_branch_children( branch_entry1, code_block->jcc.block_addr[ 0 ] );
                add_branch_children( branch_entry2, code_block->jcc.block_addr[ 1 ] );

                virt_instr_entry->addChildren( { branch_entry1, branch_entry2 } );
                // if its a JMP with branches we want to insert the next code block
                // instructions into the child widget entries...
                ui.virt_instrs->addTopLevelItem( virt_instr_entry );
                goto finish; // bad code...
            }
            else if ( virt_instr->mnemonic_t == vm::handler::JMP )
                virt_instr_entry->setText( 3, QString( "; { %1 }" ).arg( code_block->jcc.block_addr[ 0 ], 0, 16 ) );

            ui.virt_instrs->addTopLevelItem( virt_instr_entry );
        }
    }

finish: // bad code...
    ui.virt_instrs->topLevelItem( 0 )->setSelected( true );
}