#include <devirt_utils.hpp>

namespace devirt
{
    namespace util
    {
        bool serialize_vmp2(
            std::vector< std::pair< std::uint32_t, std::vector< vm::instrs::code_block_t > > > &virt_rtns,
            std::vector< std::uint8_t > &vmp2file )
        {
            const auto file_header = reinterpret_cast< vmp2::v4::file_header * >( vmp2file.data() );

            if ( file_header->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 >( file_header ) +
                                                                    file_header->rtn_offset );

            for ( auto [ rtn_block, rtn_idx ] = std::pair{ first_rtn, 0ull }; rtn_idx < file_header->rtn_count;
                  ++rtn_idx, rtn_block = reinterpret_cast< vmp2::v4::rtn_t * >(
                                 reinterpret_cast< std::uintptr_t >( rtn_block ) + rtn_block->size ) )
            {
                virt_rtns.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 ] );

                    virt_rtns.back().second.push_back( _code_block );
                }
            }

            return true;
        }
    } // namespace util

    bool append( std::vector< std::uint8_t > &obj, std::vector< std::uint8_t > &bin )
    {
        if ( obj.empty() || bin.empty() )
            return false;

        // TODO: append a new section to the pe file...
        // put the entire obj file into this new section...
        // fix up the vmenter functions to jmp to their respective devirtualized routines...
        //
        // move relocation directory to the end of the module and append new relocations...
        // export all of the new devirtualized/llvm generated functions...
        return true;
    }
} // namespace devirt