#include <cstdlib>
#include <iostream>
#include <vector>
#include <string>
using namespace std;

// Command line parsing
#include <getopt.h>

// DyninstAPI includes
#include "BPatch.h"
#include "BPatch_binaryEdit.h"
#include "BPatch_flowGraph.h"
#include "BPatch_function.h"
#include "BPatch_point.h"

#include "BPatch_object.h"
#include "BPatch_module.h"

#include "Type.h"
#include "InstSpec.h"

#ifdef HAVE_SELINUX
#include <selinux/selinux.h>
#endif

using namespace Dyninst;

/* Every Dyninst mutator needs to declare one instance of BPatch */
BPatch bpatch;

int ends_with_function(char *str, const char *str2)
{
    char *dot = strstr(str, str2);

    if (NULL == dot) return 0;
    return strcmp(dot, str2) == 0;
}

BPatch_function *findFuncByName (BPatch_image * appImage, char *funcName)
{
    /* fundFunctions returns a list of all functions with the name 'funcName' in the binary */
    BPatch_Vector < BPatch_function * >funcs;
    if (NULL == appImage->findFunction (funcName, funcs) || !funcs.size ()
            || NULL == funcs[0]) {
        cerr << "Failed to find " << funcName <<
            " function in the instrumentation library" << endl;
        return NULL;
    }
    return funcs[0];
}

bool insertFuncEntry (BPatch_addressSpace * appSpace, BPatch_function * curFunc,
        const char *funcName, BPatch_function * instIncFunc,
        int funcId)
{
    /* Find the instrumentation points */
    vector < BPatch_point * >*funcEntry = curFunc->findPoint (BPatch_entry);
    if (NULL == funcEntry) {
        cerr << "Failed to find entry for function " << funcName << endl;
        return false;
    }

    cout << "Inserting instrumention at function entry of " << funcName << endl;
    /* Create a vector of arguments to the function
     * incCoverage function takes the function name as argument */
    BPatch_Vector < BPatch_snippet * >instArgs;

    const char *calltype = "none";//sql_call";
    BPatch_constExpr id2 (calltype);
    instArgs.push_back (&id2);
BPatch_constExpr id (funcName);
    instArgs.push_back (&id);
    BPatch_funcCallExpr instIncExpr (*instIncFunc, instArgs);
//cout << "before insert"<<endl;
    InstSpec is_init;
    is_init.trampGuard = false;
    is_init.redZone = false;
    is_init.saveRegs.push_back(x86_64::rax);
    is_init.saveRegs.push_back(x86_64::rdx);
    /* Insert the snippet at function entry */
    BPatchSnippetHandle *handle =
        appSpace->insertSnippet (instIncExpr, *funcEntry, BPatch_callBefore,
                BPatch_lastSnippet, &is_init);
    if (!handle) {
        cerr << "Failed to insert instrumention at function entry of " << funcName
            << endl;
        return false;
    }
    return true;

}

int main (int argc, char *argv[])
{
    if(argc < 2) {
        cerr << "./codeCoverage inputbinary" << endl;
        return EXIT_FAILURE;
    }
    char *inputprogram = argv[1];
    pid_t pid ;
    char line[512];
    FILE * command = NULL;
    char strPID[100];    
    sprintf(strPID,"pidof -s %s",inputprogram);
    command = popen(strPID,"r");
    fgets(line,512,command);
    pid = strtoul(line,NULL,10);
    pclose(command);

    const char *func1 = "main.mainAdmin2";
    const char *func2 = "main.testinstrument";

    const char *func10 = "main.mainAdmin1";

    const char *func100 = "main.Testfunc1";


    const char *func3 = "instrumentmyfunc";
    const char *func4 = "omalloctest";
    
    
    cerr << "pid of process " << pid << endl;

    BPatch_addressSpace* addrSpace = bpatch.processAttach (NULL, pid);
    if (addrSpace == NULL) {
        cerr << "Failed to attach to the process" << endl;
        return EXIT_FAILURE;
    }
    cerr << "pid of process complete" << pid << endl;
    BPatch_process* appProc = dynamic_cast<BPatch_process*>(addrSpace);   
    if (appProc == NULL ) {
        cerr << "BPatch_process is NULL" << endl;
        return EXIT_FAILURE;
    }
    /*const char *instLibraryo = "/usr/lib/mycavlib.a";
    cout<<"loadlib start "  <<endl;
    if (!addrSpace->loadLibrary (instLibraryo)) {
        cerr << "Failed to open instrumentation library" << endl;
        return EXIT_FAILURE;
    }
    cout<<"loadlib done "  <<endl;*/

    BPatch_image *appImage = appProc->getImage ();
    if (appImage == NULL) {
        cerr << "Failed to get image" << endl;
        return EXIT_FAILURE;
    }

    BPatch_function *inserthellofunc =NULL,*hellofunc=NULL,*hellofunc10=NULL,*hellofunc100=NULL;
    BPatch_function *insertinfunc =NULL,*insertedfunc=NULL;
    int funcIndex = 100;
    vector < BPatch_module * >*modules = appImage->getModules ();
    vector < BPatch_module * >::iterator moduleIter;
    for (moduleIter = modules->begin (); moduleIter != modules->end ();
            ++moduleIter) {
        char moduleName[1024];
        (*moduleIter)->getName (moduleName, 1024);
        
        vector < BPatch_function * >*allFunctions =
            (*moduleIter)->getProcedures ();
        vector < BPatch_function * >::iterator funcIter;

        for (funcIter = allFunctions->begin (); funcIter != allFunctions->end ();
                ++funcIter) {
            BPatch_function *curFunc = *funcIter;
            
            char funcName[1024];
            curFunc->getName (funcName, 1024);

            if(ends_with_function(funcName,func2))
            {
                cout<<"found " << func2 <<endl;
                inserthellofunc = curFunc;
            }
            if(ends_with_function(funcName,func1))
            {
                cout<<"found " << func1 <<endl;
                hellofunc = curFunc;
            }

            if(ends_with_function(funcName,func10))
            {
                cout<<"found " << func10 <<endl;
                hellofunc10 = curFunc;
            }

            if(ends_with_function(funcName,func100))
            {
                cout<<"found " << func100 <<endl;
                hellofunc100 = curFunc;
            }

            if(ends_with_function(funcName,func3))
            {
                cout<<"found " << func3 <<endl;
                insertinfunc = curFunc;
            }

            if(ends_with_function(funcName,func4))
            {
                cout<<"found " << func4 <<endl;
                insertedfunc = curFunc;
            }
        }
        //funcIndex++;
    }
    if(inserthellofunc && hellofunc)
    {
        funcIndex = 500001;
      cout<<"inserting " << func2  << " value " << funcIndex <<endl;
      insertFuncEntry (addrSpace, hellofunc, func1, inserthellofunc, funcIndex);
    }

    if(inserthellofunc && hellofunc10)
    {
        funcIndex = 1500001;
      cout<<"inserting " << func2  << " value " << funcIndex <<endl;
      insertFuncEntry (addrSpace, hellofunc10, func10, inserthellofunc, funcIndex);
    }

    if(inserthellofunc && hellofunc100)
    {
        funcIndex = 85000001;
      cout<<"inserting " << func2  << " value " << funcIndex <<endl;
      insertFuncEntry (addrSpace, hellofunc100, func100, inserthellofunc, funcIndex);
    }

    if(insertedfunc && insertinfunc)
    {
      cout<<"inserting " << func2  << " value " << funcIndex <<endl;
      insertFuncEntry (addrSpace, insertinfunc, func1, insertedfunc, funcIndex);
    }

    cout << "end main" <<endl;
    appProc->continueExecution();
    appProc->detach(true);
    return EXIT_SUCCESS;
}
