#include <Image.h>
#include <Process.h>

extern "C" {
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h> 
#include <demangle.h>
#include <paths.h>
#include <libgen.h>
#include <sys/param.h>
#include <sys/module.h>
#include <sys/linker.h>
}

bool Image::firstTimeInit = true;

const std::string Image::KERNEL_NAME( getbootfile() );
const std::string Image::BOOT_DIRECTORY( dirname( KERNEL_NAME.c_str() ) );
const std::string Image::TEXT_SECTION_NAME( ".text" );

Image* Image::kernel = 0;

Image::ImageMap Image::imageMap;
Image::LoadableImageMap Image::kernelLoadableImageMap;

Image&
Image::getKernel()
{
    if ( Image::kernel == 0 )
    {
        kernel = new Image( KERNEL_NAME );
    }
    return *kernel;
}

Image::Image( const std::string& imageName ) :
    m_fd( 0 ),
    m_symCount( 0 ),
    m_dynCount( 0 ),
    m_bfdHandle( 0 ),
    m_startAddress( 0 ),
    m_textSize( 0 ),
    m_endAddress( 0 ),
    m_textSection( 0 ),
    m_symTable( 0 ),
    m_dynTable( 0 )
{
    if ( firstTimeInit )
    {
        bfd_init();
        firstTimeInit = false;
    }

    m_fd = open( imageName.c_str(), O_RDONLY );

    if ( m_fd < 0 )
    {
        if (imageName.size())
        {
            fprintf(stderr, "%s: unable to open file %s\n", 
                    g_quitOnError ? "error" : "warning", imageName.c_str());
            if (g_quitOnError)
            {

                exit(5);
            }
        }
        return;
    }

    m_bfdHandle = bfd_fdopenr( imageName.c_str(), 0, m_fd );

    if ( m_bfdHandle == 0 )
    {
        return;
    }

    //must check the format!
    if ( bfd_check_format( m_bfdHandle, bfd_object ) == 0 )
    {
        return;
    }

    if ( ( bfd_get_file_flags( m_bfdHandle ) & HAS_SYMS ) == 0 )
    {
        return;
    }

    m_textSection = bfd_get_section_by_name( m_bfdHandle, TEXT_SECTION_NAME.c_str() );
    m_startAddress = bfd_get_section_vma( m_bfdHandle, m_textSection );
    m_textSize = bfd_get_section_size( m_textSection );
    m_endAddress = m_startAddress + m_textSize;
}

Image::~Image()
{
    if ( m_symTable != 0 )
    {
        free( m_symTable );
        m_symTable = 0;
    }

    if ( m_dynTable != 0 )
    {
        free( m_dynTable );
        m_dynTable = 0;
    }

    if ( m_bfdHandle != 0 )
    {
        bfd_close( m_bfdHandle );
        m_bfdHandle = 0;
    }
}

bool
Image::isOk() const
{
    return ( m_fd > 0 ) && ( m_textSize > 0 );
}

void
Image::loadSymtab()
{
    if ( m_symTable == 0 )
    {
        m_symCount = bfd_get_symtab_upper_bound( m_bfdHandle );
        if ( m_symCount == 0 )
        {
            return;
        }

        m_symTable = (asymbol **) malloc( m_symCount );
        m_symCount = bfd_canonicalize_symtab( m_bfdHandle, m_symTable );

        if ( m_symCount < 0 )
        {
            printf( "canonicalize symtab failed %ld\n", m_symCount );
        }
    }
}

void
Image::loadDyntab()
{
    m_dynCount = bfd_get_dynamic_symtab_upper_bound( m_bfdHandle );
    if ( m_dynCount == 0 )
    {
        return;
    }

    m_dynTable = (asymbol **) malloc( m_dynCount );
    m_dynCount = bfd_canonicalize_dynamic_symtab( m_bfdHandle, m_dynTable );

    if ( m_dynCount < 0 )
    {
        printf( "canonicalize dynamic symtab failed %ld\n", m_dynCount );
    }
}

void
Image::dumpSymtab()
{
    if ( m_symTable == 0 )
    {
        loadSymtab();
    }

    for ( long it = 0; it < m_symCount; it++ ) {

        asymbol* symbol = m_symTable[ it ];

        if ( symbol -> name == 0 || *symbol -> name == '\0' )
        {
            continue;
        }

        printf( "symtab " );
        bfd_print_symbol_vandf( m_bfdHandle, stdout, symbol );
        printf( " section: %s, symbol: %s\n", symbol -> section -> name, demangle( symbol -> name ) );
    }
}

void
Image::dumpDyntab()
{
    if ( m_dynTable == 0 )
    {
        loadDyntab();
    }

    for ( long it = 0; it < m_dynCount; it++ ) {

        asymbol* symbol = m_dynTable[ it ];

        if ( symbol -> name == 0 || *symbol -> name == '\0' )
        {
            continue;
        }

        printf( "dyntab " );
        bfd_print_symbol_vandf( m_bfdHandle, stdout, symbol );
        printf( " section: %s, symbol: %s\n", symbol -> section -> name, demangle( symbol -> name ) );
    }
}

bool
Image::isContained( Location& location, uintptr_t loadOffset )
{
    const bfd_vma& address( location.getAddress() );

    return ( ( address - loadOffset ) >= m_startAddress ) &&
        ( ( address - loadOffset ) < m_endAddress );
}

void
Image::mapLocation( Location& location, uintptr_t loadOffset )
{
    uintptr_t address = location.getAddress() - loadOffset - m_startAddress;

    loadSymtab();

    location.isMapped( false );

    if ( m_symCount > 0 )
    {
        location.isMapped( bfd_find_nearest_line( m_bfdHandle, m_textSection, m_symTable,
            address, &location.m_filename, &location.m_functionname, &location.m_linenumber ) != 0 );
    }

    if ( location.m_isMapped )
    {
        return;
    }

    loadDyntab();

    if ( m_dynCount > 0 )
    {
        location.isMapped( bfd_find_nearest_line( m_bfdHandle, m_textSection, m_dynTable,
            address, &location.m_filename, &location.m_functionname,
            &location.m_linenumber ) != 0 );
    }
}

void
Image::functionStart( Location& location )
{
    long it;

    for ( it = 0; it < m_symCount; ++it )
    {
        asymbol* symbol = m_symTable[ it ];
        const char* symName = symbol -> name;

        if ( symName == 0 )
        {
            continue;
        }

        if ( strcmp( location.m_functionname, symName ) == 0 )
        {
            Location testLocation = location;
            testLocation.m_address = bfd_asymbol_value( symbol );
            mapLocation( testLocation );
            if ( testLocation.m_isMapped )
            {
                if ( ( testLocation.m_filename != 0 && location.m_filename  ) && ( strcmp( testLocation.m_filename, location.m_filename ) == 0 ) )
                {
                    location = testLocation;
                    break;
                }
            }
        }
    }

    if ( location.m_isMapped )
    {
        return;
    }

    for ( it = 0; it < m_dynCount; ++it )
    {
        asymbol* symbol = m_dynTable[ it ];
        const char* symName = symbol -> name;

        if ( symName == 0 )
        {
            continue;
        }

        if ( strcmp( location.m_functionname, symName ) == 0 )
        {
            Location testLocation = location;
            testLocation.m_address = bfd_asymbol_value( symbol );
            mapLocation( testLocation );
            if ( testLocation.m_isMapped )
            {
                if ( ( testLocation.m_filename != 0 ) && ( strcmp( testLocation.m_filename, location.m_filename ) == 0 ) )
                {
                    location = testLocation;
                    break;
                }
            }
        }
    }
}

void
Image::mapFunctionStart( FunctionLocation& functionLocation )
{
    if ( !functionLocation.m_isMapped )
    {
        return;
    }

    Image* image = functionLocation.m_isKernel ? kernel : imageMap[ functionLocation.m_modulename ];

    if ( image == 0 )
    {
        return;
    }

    functionLocation.m_isMapped = false;
    image -> functionStart( functionLocation );
    std::sort( functionLocation.m_lineLocationList.begin(), functionLocation.m_lineLocationList.end() );

    if ( !functionLocation.m_isMapped )
    {
        functionLocation.m_linenumber = functionLocation.m_lineLocationList[ 0 ];
    }
}

std::string
Image::getLoadableImageName( const Location& location, uintptr_t& loadOffset )
{
    if ( kernelLoadableImageMap.empty() )
    {
        int fileid = 0;
        for ( fileid = kldnext( fileid ); fileid > 0; fileid = kldnext( fileid ) )
        {
            kld_file_stat stat;
            stat.version = sizeof( kld_file_stat );
            if ( kldstat( fileid, &stat ) < 0 )
            {
                printf( "can't stat module\n" );
            }
            else
            {
                kernelLoadableImageMap[ bfd_vma( stat.address ) ] = BOOT_DIRECTORY + '/' + stat.name;
            }
        }
    }

    loadOffset = 0;

    std::string loadableImageName;

    LoadableImageMap::iterator it = kernelLoadableImageMap.begin();

    for ( ; it != kernelLoadableImageMap.end(); ++it )
    {
        if ( location.getAddress() < (*it).first )
        {
            break;
        }
    }

    if ( it != kernelLoadableImageMap.begin() )
    {
        --it;
        loadOffset = (*it).first;
        loadableImageName = (*it).second;
    }

    return loadableImageName;

}

void
Image::mapAllLocations( LocationList& locationList )
{
    Image& kernel( getKernel() );

    for ( LocationList::iterator it = locationList.begin(); it != locationList.end(); ++it )
    {
        Location& location( *it );
        if ( location.m_isKernel )
        {
            if ( kernel.isContained( location ) )
            {
                kernel.mapLocation( location );
                if ( location.m_isMapped )
                {
                    location.m_modulename = KERNEL_NAME.c_str();
                }
            }
            else
            {
                uintptr_t loadOffset = 0;
                const std::string& loadableImageName( Image::getLoadableImageName( location, loadOffset ) );

                Image* kernelModuleImage = imageMap[ loadableImageName ];

                if ( kernelModuleImage == 0 )
                {
                    kernelModuleImage = new Image( loadableImageName );
                    imageMap[ loadableImageName ] = kernelModuleImage;
                }

                if ( ! kernelModuleImage -> isOk() )
                {
                    continue;
                }

                if ( kernelModuleImage -> isContained( location, loadOffset ) )
                {
                    kernelModuleImage -> mapLocation( location, loadOffset );
                    if ( location.m_isMapped )
                    {
                        location.m_modulename = loadableImageName.c_str();
                    }
                }
            }
        }
        else
        {
            Process* process = Process::getProcess( location.m_pid );

            if ( process == 0 )
            {
                continue;
            }

            const std::string& processName( process -> getName() );

            Image* processImage = imageMap[ processName ];

            if ( processImage == 0 )
            {
                processImage = imageMap[ processName ] = new Image( processName );
            }

            if ( ! processImage -> isOk() )
            {
                continue;
            }

            if ( processImage -> isContained( location ) )
            {
                processImage -> mapLocation( location );
                if ( location.m_isMapped )
                {
                    location.m_modulename = processName.c_str();
                }
            }
            else
            {
                uintptr_t loadOffset = 0;
                const std::string& loadableImageName(
                    process -> getLoadableImageName( location, loadOffset ) );

                Image* processLoadableImage = imageMap[ loadableImageName ];

                if ( processLoadableImage == 0 )
                {
                    processLoadableImage = imageMap[ loadableImageName ] = new Image( loadableImageName );
                }

                if ( ! processLoadableImage -> isOk() )
                {
                    continue;
                }

                if ( processLoadableImage -> isContained( location, loadOffset ) )
                {
                    processLoadableImage -> mapLocation( location, loadOffset );
                    if ( location.m_isMapped )
                    {
                        location.m_modulename = loadableImageName.c_str();
                    }
                }
            }
        }
    }
}

char *
Image::demangle( const char* name )
{
    if ( name == 0 )
    {
        return 0;
    }

    char* result = cplus_demangle ( name, DMGL_ANSI | DMGL_PARAMS );

    return result == 0 ? strdup( name ) : result;
}
