#include <Process.h>
#include <Sample.h>
#include <ProcessState.h>
#include "sharedLib.h"
#include <algorithm>

extern "C" {
#include <fcntl.h>
#include <kvm.h>
#include <sys/param.h>
#include <sys/sysctl.h>
#include <sys/user.h>
#include <sys/proc.h>
#include <sys/types.h>
#include <sys/stat.h>
}

Process::ProcessMap Process::processMap;

void
Process::fillProcessMap()
{
    int processCount = 0;

    kvm_t* kd = kvm_open( 0, 0, 0, O_RDONLY, 0 );

    if ( kd == 0 )
    {
        printf( "could not get list of procs\n" );
        return;
    }

    kinfo_proc* procInfoArray = kvm_getprocs( kd, KERN_PROC_PROC, 0, &processCount );

    if ( procInfoArray == 0 )
    {
        kvm_close( kd );
        printf( "could not get list of procs\n" );
        return;
    }

    for ( int it = 0; it < processCount; it++ )
    {
        bool isFullPathValid = false;

        char ** procArgv = 0;
        procArgv = kvm_getargv( kd, procInfoArray + it, 1024 );
        if ( procArgv != 0 )
        {
            struct stat fileStat;            
            isFullPathValid = stat( *procArgv, &fileStat ) == 0;
        }
        char commandBuffer[ 1024 ];
        const char* execName = isFullPathValid ? *procArgv : ( procInfoArray + it ) -> ki_comm;
        if ( !isFullPathValid )
        {
            sprintf( commandBuffer, "%s %s\n", "which", execName );
            FILE* command = popen( commandBuffer, "r" );
            if ( command != 0 )
            {
                fgets( commandBuffer, sizeof( commandBuffer ), command );
                int freadResult = ferror( command );
                int closeResult = pclose( command );

                if ( ( freadResult == 0 ) && ( closeResult == 0 ) )
                {
                    char * carriageReturn = strchr( commandBuffer, '\n' );
                    if ( carriageReturn != 0 )
                    {
                        *carriageReturn = 0;
                    }
                    execName = commandBuffer;
                }
            }
        }
        pid_t pid = ( procInfoArray + it ) -> ki_pid;

        if ( processMap[ pid ] == 0 )
        {
            processMap[ pid ] = new Process( ProcessExec( pid, execName ) );
        }
    }

    kvm_close( kd );
}

Process::Process( const Sample& sample ) :
    m_pid( sample.getProcessID() ),
    m_sampleCount( 0 )
{
}

Process::Process( const ProcessExec& processExec ) :
    m_pid( processExec.getProcessID() ),
    m_sampleCount( 0 ),
    m_name( processExec.getProcessName() )    
{
}

std::string
Process::getLoadableImageName( const Location& location, uintptr_t& loadOffset )
{
    if ( m_loadableImageMap.empty() && !m_name.empty() )
    {

        try
        {
            sharedLibInfo sli(m_pid);

            unsigned n = sli.numLibs();
            for (unsigned i = 0; i < n; ++i)
            {
                const sharedLib& sl = sli.getLib(i);
                m_loadableImageMap[sl.getBase()] = sl.getName();
            }
        }
        catch (...) // can't get shared lib info...
        {
        }
    }

    loadOffset = 0;

    std::string loadableImageName;

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

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

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

    return loadableImageName;
}

Process&
Process::getProcess( const Sample& sample )
{
    pid_t pid = sample.getProcessID();
    Process* process = processMap[ pid ];

    if ( process == 0 )
    {
        process = processMap[ pid ] = new Process( sample );
    }
    return *process;
}

Process&
Process::getProcess( const ProcessExec& processExec )
{
    pid_t pid = processExec.getProcessID();
    Process* process = processMap[ pid ];

    if ( process == 0 )
    {
        process = processMap[ pid ] = new Process( processExec );
    }

    if ( ( *process ).m_name.empty() && !processExec.getProcessName().empty() )
    {
        process -> m_name = processExec.getProcessName();
    }

    return *process;
}

void
Process::addSample( const Sample& sample )
{
    if ( sample.isKernel() )
    {
        m_kernelSamples[ sample.getAddress() ]++;
    }
    else
    {
        m_processSamples[ sample.getAddress() ]++;
    }
    m_sampleCount++;
}

void
Process::collectLocations( LocationList& locationList )
{
    for( SampleMap::iterator it = m_processSamples.begin(); it != m_processSamples.end(); ++it )
    {
        locationList.push_back( Location( false, m_pid, (*it).first, (*it).second ) );
    }
    m_processSamples.clear();
    for( SampleMap::iterator it = m_kernelSamples.begin(); it != m_kernelSamples.end(); ++it )
    {
        locationList.push_back( Location( true, m_pid, (*it).first, (*it).second ) );
    }
    m_kernelSamples.clear();
}

void
Process::mapAllFunctions( LocationList& locationList )
{
    for ( LocationList::iterator it = locationList.begin(); it != locationList.end(); it++ )
    {
        Location& location( *it );
        const std::string& functionName( location.getFunctionName() );
        const std::string& fileName( location.getFileName() );
        const std::string& key( fileName + functionName );
        FunctionLocationMap& functionLocationMap( getProcess( location.getPid() ) -> m_functionLocationMap );

        FunctionLocationMap::iterator findPos = functionLocationMap.find( key );

        if ( findPos == functionLocationMap.end() )
        {
            functionLocationMap.insert( findPos,
                FunctionLocationMap::value_type( key, FunctionLocation( location ) ) );
        }
        else
        {
            (*findPos).second += location;
        }
    }
}

void
Process::collectActiveProcesses( ActiveProcessList& activeProcessList )
{
    for( ProcessMap::const_iterator it = processMap.begin(); it != processMap.end(); ++it )
    {
        Process* process( (*it).second );

        if ( process -> m_sampleCount > 0 )
        {
            activeProcessList.push_back( process );
        }
    }
    std::sort( activeProcessList.begin(), activeProcessList.end(), ProcessSampleCompare() );
}

void
Process::collectAllLocations( LocationList& locationList )
{
    for( ProcessMap::iterator it = processMap.begin(); it != processMap.end(); ++it )
    {
        (*it).second -> collectLocations( locationList );
    }
    std::sort( locationList.begin(), locationList.end() );
}

void
Process::getFunctionList( FunctionList& functionList )
{
    for( FunctionLocationMap::iterator it = m_functionLocationMap.begin();
        it != m_functionLocationMap.end(); ++it )
    {
        FunctionLocation functionLocation( (*it).second );
        Image::mapFunctionStart( functionLocation );
        functionList.push_back( functionLocation );
    }
    m_functionLocationMap.clear();
    std::sort( functionList.begin(), functionList.end() );
}
