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

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

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

#include "BPatch_object.h"

#include "Type.h"

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

using namespace Dyninst;

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

bool deleteFuncEntry (BPatch_addressSpace * appSpace, BPatch_function * curFunc,
                const char *funcName, BPatch_function * instIncFunc, bool bentry) {
    bool status = false;
    vector < BPatch_point * >*funcEntry;
    BPatch_callWhen callWhen = BPatch_callBefore;
    if (bentry) {
        funcEntry = curFunc->findPoint (BPatch_entry);
    }
    else {
        funcEntry = curFunc->findPoint (BPatch_exit);
        callWhen = BPatch_callAfter;
    }
    if (NULL == funcEntry) {
        cerr << "Failed to find entry for function " << endl;
        return false;
    }

    BPatch_Vector < BPatch_snippet * >instArgs;
    // BPatch_constExpr fname (funcName);
    // instArgs.push_back (&fname);
    // const char *secondarg = "none";
    // if(callout)
    //     secondarg = callout;
    // BPatch_constExpr fcallout (secondarg);
    // instArgs.push_back (&fcallout);
    BPatch_funcCallExpr instIncExpr (*instIncFunc, instArgs);

    // InstSpec is_init;
    // is_init.trampGuard = false;
    // is_init.redZone = false;
    // is_init.saveRegs.push_back(x86_64::rdi);
    // is_init.saveRegs.push_back(x86_64::rsi);
    fprintf(stderr, "deleting snippet\n");
    appSpace->DeletingSnippet (instIncExpr, *funcEntry, callWhen,
                            BPatch_firstSnippet,NULL);

    return status;
}
bool insertFuncEntry (BPatch_addressSpace * appSpace, BPatch_function * curFunc,
        char *funcName, BPatch_function * instIncFunc, bool bentry)
{
    vector < BPatch_point * >*funcEntry;
    BPatch_callWhen callWhen = BPatch_callBefore;

    char f1[1024],f2[1024];
            curFunc->getName (f1, 1024);
            instIncFunc->getName (f2, 1024);
            fprintf(stderr, "f1f2 %d, %s  %s\n", bentry,f1,f2);
            
    if (bentry) {
        funcEntry = curFunc->findPoint (BPatch_entry);
    }
    else {
        funcEntry = curFunc->findPoint (BPatch_exit);
        callWhen = BPatch_callAfter;
    }
    if (NULL == funcEntry) {
        cerr << "Failed to find entry for function " << funcName << endl;
        return false;
    }

    cout << "Inserting instrumention at function entry of " << funcName << endl;

    BPatch_Vector < BPatch_snippet * >instArgs;

    // const char *firstarg = "first";
    // BPatch_constExpr farg1 (firstarg);
    // instArgs.push_back (&farg1);
    // const char *secondarg = "second";
    // BPatch_constExpr farg2 (secondarg);
    // instArgs.push_back (&farg2);
    // is_init.trampGuard = false;
    // is_init.redZone = false;
    // is_init.saveRegs.push_back(x86_64::rdi);
    //is_init.saveRegs.push_back(x86_64::rsi);


    BPatch_funcCallExpr instIncExpr (*instIncFunc, instArgs);

    BPatchSnippetHandle *handle =
        appSpace->insertSnippet (instIncExpr, *funcEntry, callWhen,
                BPatch_firstSnippet);
    if (!handle) {
        cerr << "Failed to insert instrumention at function entry of " << funcName
            << endl;
        return false;
    }
    return true;
}

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

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

int main (int argc, char *argv[])
{
    char *inputProgram = NULL;
    char *strFntoInstrument = NULL, *strInstrumentedfunc = NULL,*strInstrumentedfunc2 = NULL;
    BPatch_process* appProc = NULL;
    pid_t pid ;
    char line[512];
    FILE * command = NULL;
    char strPID[100];
    inputProgram = "student";
    //int doinstr = atoi(argv[1]);
    //fprintf(stderr, "doinstr %d\n", doinstr);
    strFntoInstrument = "pg.Connect";
    strInstrumentedfunc = "json_value_init_object";
    strInstrumentedfunc2 = "Method_myexitWrap";
    sprintf(strPID,"pidof -s %s",inputProgram);
    command = popen(strPID,"r");
    fgets(line,512,command);
    pid = strtoul(line,NULL,10);
    pclose(command);

    cerr << "pid of process " << pid << endl;
    BPatch_addressSpace* app2 = bpatch.processAttach (NULL, pid);
    if (app2 == NULL) {
        cerr << "Failed to open process" << endl;
        return EXIT_FAILURE;
    }

    appProc = dynamic_cast<BPatch_process*>(app2);   
    if (appProc == NULL ) {
        cerr << "Failed to open binary" << endl;
        return EXIT_FAILURE;
    }

    const char *instLibraryw = "/usr/lib/libusejson.so";
    cout<<"loadlib start "  <<endl;
    if (!app2->loadLibrary (instLibraryw)) {
        cerr << "Failed to open instrumentation library" << endl;
        return EXIT_FAILURE;
    }
    // cout<<"loadlib done "  <<endl;

    BPatch_function *Instrumentedfunc=NULL,*Instrumentedfunc2=NULL,*funcToInstrument=NULL;
    BPatch_image *appImage = appProc->getImage ();

    if (appImage == NULL) {
        cerr << "Failed to get image" << endl;
        return EXIT_FAILURE;
    }
    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);
        //cout << "Instrumenting module: " << moduleName << endl;
        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,strInstrumentedfunc))
            {
                cout<<"found " << strInstrumentedfunc <<endl;
                Instrumentedfunc = curFunc;
            }
            if(ends_with_function(funcName,strInstrumentedfunc2))
            {
                cout<<"found " << strInstrumentedfunc2 <<endl;
                Instrumentedfunc2 = curFunc;
            }
            if(ends_with_function(funcName,strFntoInstrument))
            {
                cout<<"found " << strFntoInstrument <<endl;
                funcToInstrument = curFunc;
            }
        }
    }

   if(funcToInstrument && Instrumentedfunc && Instrumentedfunc)
    {
        //if(doinstr == 1) {
            cerr << "inserting in " << strFntoInstrument << endl;
            insertFuncEntry (app2, funcToInstrument,strFntoInstrument, Instrumentedfunc,true);
            //insertFuncEntry (app2, funcToInstrument,strFntoInstrument, Instrumentedfunc2,false);
        //}
        // else if(doinstr == 0){
        //     deleteFuncEntry(app2, funcToInstrument,strFntoInstrument, Instrumentedfunc,true);
        //     deleteFuncEntry(app2, funcToInstrument,strFntoInstrument, Instrumentedfunc2,false);
        // }
    }
    // if(doinstr != 2) {
    //     app2->InjectUnwinding("runtime.findfunc",inputProgram,1);
    //     app2->InjectUnwinding("runtime.pcvalue",inputProgram,4);
        
    //     app2->FinalizeUnwinding();
    // }
    cout<<"end mainfunc num " <<endl;
    
    appProc->detach(true);
    appProc->continueExecution();

    return EXIT_SUCCESS;
}
