#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;

bool insertFuncEntry (BPatch_addressSpace * appSpace, BPatch_function * curFunc,
        char *funcName, BPatch_function * instIncFunc,
        int funcId)
{
    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;

    BPatch_Vector < BPatch_snippet * >instArgs;
    // BPatch_constExpr fname (funcName);
    // instArgs.push_back (&fname);
    // const char *secondarg = "none";
    // BPatch_constExpr fcallout (secondarg);
    // instArgs.push_back (&fcallout);
    // 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);


    BPatch_funcCallExpr instIncExpr (*instIncFunc, instArgs);

    BPatchSnippetHandle *handle =
        appSpace->insertSnippet (instIncExpr, *funcEntry, BPatch_callBefore,
                BPatch_lastSnippet);
    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 *wrappingFuncName = NULL, *wrapppedFuncName = NULL,*strFntoInstrument = NULL, *strInstrumentedfunc = NULL;
    BPatch_process* appProc = NULL;
    pid_t pid ;
    char line[512];
    FILE * command = NULL;
    char strPID[100];
    char * nmf1 = "github.com/testforchi.(*httphandler).ServeHTTP";
    inputProgram = argv[1];
    wrappingFuncName = "github.com/go-chi/chi.NewMux";
    wrapppedFuncName = "WrapChi";
    strFntoInstrument = "github.com/qor/qor-example/app/products.Controller.Show";
    strInstrumentedfunc = "Getlibsymbol";

    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_function *nmf1Func=NULL;
    BPatch_addressSpace* app2 = bpatch.processAttach (NULL, pid);
    //BPatch_addressSpace* app2 = bpatch.processCreate (inputProgram, NULL,NULL);
    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/libtestgolib.so";//"libgithub.com-testforchi.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 *wrappedFunc=NULL,*wrappingFunc=NULL, *Instrumentedfunc=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;
    char symbolname[1024];
    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,wrapppedFuncName))
            {
                cout<<"found " << funcName <<endl;
                wrappedFunc = curFunc;
            }
            if(ends_with_function(funcName,strInstrumentedfunc))
            {
                cout<<"found " << strInstrumentedfunc <<endl;
                Instrumentedfunc = curFunc;
            }
            if(ends_with_function(funcName,strFntoInstrument))
            {
                cout<<"found " << strFntoInstrument <<endl;
                funcToInstrument = curFunc;
            }
            if(ends_with_function(funcName,nmf1))
            {
                cout<<"found " << funcName << " in mod " << moduleName<< endl;
                if(strstr(moduleName,"testforchi") !=0) {
                    cout<<"found use " << funcName <<endl;
                    nmf1Func = curFunc;
                }
            }
            char *dot = NULL;
            dot = strstr(funcName, wrappingFuncName);
            if(dot != NULL && strcmp(dot,wrappingFuncName) == 0)            
            {
                fprintf(stderr, "found %s in module %s\n", funcName,moduleName);           
                if(strcmp(moduleName,inputProgram) == 0)
                {
                    cout<<"found omalloc " << funcName <<endl;
                    wrappingFunc = curFunc;
                }
                else if(strstr(moduleName,"testgolib") !=0)
                {
                    curFunc->getName (symbolname, 1024);
                    cout<<"found symbol in testgolib" <<endl;
                }
            }
        }
    }

   if(funcToInstrument && Instrumentedfunc)
    {
        cerr << "inserting in " << strFntoInstrument << endl;
        insertFuncEntry (app2, funcToInstrument,strFntoInstrument, Instrumentedfunc,2);
    }

    if(wrappingFunc && wrappedFunc)
    {
        Dyninst::SymtabAPI::Symtab *symtab = Dyninst::SymtabAPI::convert(wrappedFunc->getModule()->getObject());
        std::vector<Dyninst::SymtabAPI::Symbol*> syms;
        symtab->findSymbol(syms, symbolname, Dyninst::SymtabAPI::Symbol::ST_UNKNOWN,Dyninst::SymtabAPI::mangledName, false,false,true);
        cout << " sym size " << syms.size() << " name  " << symbolname << endl;
       
        for(int i = 0; i <(int)syms.size();i++)
        {
            if(syms[i]->getModule() !=NULL) {
               // if(nmf1Func)
                   // app2->AddSecondFunc(nmf1Func);
                cout << " wrapping 1 " << syms[i] << " off " << syms[i]->getOffset()<< endl;
               // app2->wrapFunction(wrappingFunc, wrappedFunc,syms[i]);
            }
        }
    }

    app2->InjectUnwinding("runtime.findfunc",inputProgram,1);
    app2->InjectUnwinding("runtime.pcvalue",inputProgram,4);
    
    app2->FinalizeUnwinding();
    cout<<"end mainfunc "  <<endl;
    
    appProc->detach(true);
    appProc->continueExecution();

    return EXIT_SUCCESS;
}
