#include "qvminspector.h"

QVMInspector::QVMInspector( QWidget *parent ) : QMainWindow( parent ), FileHeader( 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 );
    ui.virtual_machine_enters->setColumnWidth( 0, 180 );

    connect( ui.action_open, &QAction::triggered, this, &QVMInspector::OnOpen );
    connect( ui.action_close, &QAction::triggered, this, &QVMInspector::OnClose );
}

void QVMInspector::OnClose()
{
    exit( 0 );
}

void QVMInspector::OnOpen()
{
    if ( FileHeader )
        free( FileHeader );

    FileHeader = nullptr;
    ImgBase = 0u, ModuleBase = 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() )
    {
        DbgMsg( "invalid vmp2 file... no file selected..." );
        return;
    }

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

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

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

    QFile open_file( file_path );
    FileHeader = reinterpret_cast< vmp2::v4::FileHeader * >( malloc( file_size ) );

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

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

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

void QVMInspector::DbgPrint( QString dbg_output )
{
    ui.dbg_output_window->appendPlainText( dbg_output );
}

void QVMInspector::DbgMsg( QString dbg_output )
{
    QMessageBox msg_box;
    msg_box.setText( dbg_output );
    msg_box.exec();
    DbgPrint( dbg_output );
}

bool QVMInspector::InitData()
{
    if ( FileHeader->magic != VMP_MAGIC )
    {
        DbgMsg( "invalid magic bytes for vmp2 file..." );

        return false;
    }

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

    if ( FileHeader->version != vmp2::version_t::v4 )
    {
        DbgMsg( "invalid vmp2 file version... "
                "this vminspector is compiled for version 4...\n" );

        return false;
    }

    ImgBase = FileHeader->ImgBase;
    ImgSize = FileHeader->module_size;
    ModuleBase = reinterpret_cast< std::uintptr_t >( FileHeader ) + FileHeader->module_offset;

    if ( !SerializeVmp2( VirtRtns ) )
    {
        DbgMsg( "failed to serialize vmp2 file format...\n" );
        return false;
    }

    UpdateUI();
    return true;
}

void QVMInspector::UpdateUI()
{
    ui.virtual_machine_enters->clear();
    for ( auto &[ rtn_rva, rtn_blks ] : VirtRtns )
    {
        auto new_item = new QTreeWidgetItem();
        new_item->setText( 0, QString( "rtn_%1" ).arg( rtn_rva + FileHeader->ImgBase, 0, 16 ) );
        new_item->setText( 1, QString( "%1" ).arg( rtn_rva + FileHeader->ImgBase, 0, 16 ) );
        new_item->setText( 2, QString( "%1" ).arg( rtn_blks.size() ) );

        std::for_each( rtn_blks.begin(), rtn_blks.end(), [ & ]( vm::instrs::code_block_t &code_blk ) {
            auto new_child = new QTreeWidgetItem();
            new_child->setText( 0, QString( "blk_%1" ).arg( code_blk.vip_begin, 0, 16 ) );
            new_child->setText( 1, QString( "%1" ).arg( code_blk.vip_begin, 0, 16 ) );
            new_child->setText( 2, QString( "%1" ).arg( code_blk.vinstrs.size() ) );
            new_item->addChild( new_child );
        } );

        ui.virtual_machine_enters->addTopLevelItem( new_item );
    }
}

bool QVMInspector::SerializeVmp2( std::vector< rtn_data_t > &VirtRtns )
{
    if ( FileHeader->version != vmp2::version_t::v4 )
    {
        std::printf( "[!] invalid vmp2 file version... this build uses v3...\n" );
        return false;
    }

    auto first_rtn = reinterpret_cast< vmp2::v4::rtn_t * >( reinterpret_cast< std::uintptr_t >( FileHeader ) +
                                                            FileHeader->rtn_offset );

    for ( auto [ rtn_block, rtn_idx ] = std::pair{ first_rtn, 0ull }; rtn_idx < FileHeader->rtn_count;
          ++rtn_idx, rtn_block = reinterpret_cast< vmp2::v4::rtn_t * >(
                         reinterpret_cast< std::uintptr_t >( rtn_block ) + rtn_block->size ) )
    {
        VirtRtns.push_back( { rtn_block->vm_enter_offset, {} } );
        for ( auto [ code_block, block_idx ] = std::pair{ &rtn_block->code_blocks[ 0 ], 0ull };
              block_idx < rtn_block->code_block_count;
              ++block_idx, code_block = reinterpret_cast< vmp2::v4::code_block_t * >(
                               reinterpret_cast< std::uintptr_t >( code_block ) + code_block->next_block_offset ) )
        {
            auto block_vinstrs = reinterpret_cast< vm::instrs::virt_instr_t * >(
                reinterpret_cast< std::uintptr_t >( code_block ) + sizeof vmp2::v4::code_block_t +
                ( code_block->num_block_addrs * 8 ) );

            vm::instrs::code_block_t _code_block{ code_block->vip_begin };
            _code_block.jcc.has_jcc = code_block->has_jcc;
            _code_block.jcc.type = code_block->jcc_type;

            for ( auto idx = 0u; idx < code_block->num_block_addrs; ++idx )
                _code_block.jcc.block_addr.push_back( code_block->branch_addr[ idx ] );

            for ( auto idx = 0u; idx < code_block->vinstr_count; ++idx )
                _code_block.vinstrs.push_back( block_vinstrs[ idx ] );

            VirtRtns.back().rtn_blks.push_back( _code_block );
        }
    }

    return true;
}