#include "core.h"
#include "common.h"
#include "graph.h"
#include "utils.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/IR/Instructions.h"
#include <sstream>
#include <iostream>
#include <filesystem>
#include <set>
#include <cxxabi.h>

static const char *temp_ll = "temp.ll";


AnalysisCore::AnalysisCore(
        const char *i,
        const char *o,
        const char *f,
        const char *e,
        unsigned int s
    ) 
    : _input(i), _output(o), _format(f), _external_opt(e), _std(s), _use_cpp(false)
{}



int AnalysisCore::run() {
    llvm::SMDiagnostic err;
    int res;
    std::stringstream ss;
    auto ref = llvm::StringRef(_input);
    if(ref.endswith(".cpp") || ref.endswith(".cxx")) {
        _use_cpp = true;
    }
    if(_use_cpp) {
        ss << "clang++-12";
    }
    else {
        ss << "clang-12";
    }
    ss <<" -S -emit-llvm ";
    if(_use_cpp) {
        ss << "-std=c++"<< _std;
    }
    ss << " " << _input;
    ss << " -o " << temp_ll;

    res = system(ss.str().c_str());
    if(res) {
        std::cerr << "Can not compile " << _input << "\n";
        return UNKNOWN_ERROR;
    }

    _caches.push_back(temp_ll);
    
    _module = llvm::parseIRFile(temp_ll, err, _context);
    if(!_module) {
        std::cerr << "Can not open IR file." << "\n";
        return UNKNOWN_ERROR;
    }

    res = analysis();

    // for (auto cache : _caches) {
    //     ::remove(cache.c_str());
    // }

    return OK;
}


int AnalysisCore::analysis() {
    for(auto &func : *_module) {
        addFunc(&func);
    }
    for (auto &func : *_module) {
        addEdges(&func);
    }
    return OK;
}

static const std::set<std::string> keyword_set = {
    "const", "delete", "new"
};

void AnalysisCore::addFunc(llvm::Function *func) {
    static unsigned index = 0;
    FuncVertex *vtx;
    if(func->isDeclaration()) return;
    if(_use_cpp) {
        if(func->getName().equals("main")) {
            vtx = new FuncVertex();
            vtx->_index = index++;
            vtx->_name = func->getName().str();
            vtx->_edges = {};
            _vertix_map.insert({vtx->_name, vtx});
        }
        else {
            auto func_name = getFuncName(func->getName());
            if(strcmp("", func_name.c_str()) == 0) return;
            auto ref = llvm::StringRef(func_name);
            if(ref.contains("__gnu_cxx::") || ref.contains("operator")) {
                return;
            }
            if(keyword_set.find(func_name) != keyword_set.end()) {
                return;
            }
            vtx = new FuncVertex();
            vtx->_index = index++;
            vtx->_name = func_name;
            vtx->_edges = {};
            if(func_name == ")") {
                int a = 0;
            }
            _vertix_map.insert({func_name, vtx});
        }
    }
    else {
        vtx = new FuncVertex();
        vtx->_index = index++;
        vtx->_name = func->getName().str();
        vtx->_edges = {};
        _vertix_map.insert({vtx->_name, vtx});
    }
}

void AnalysisCore::addEdges(llvm::Function *func) {
    if(func->isDeclaration()) {
        return;
    }
    FuncVertex *vtx;
    llvm::CallInst *call;
    llvm::InvokeInst *invoke;
    std::set<std::string> edgemap = {};
    Edge *edge;
    std::string func_name, f_name;
    if(_use_cpp) {
        if(func->getName().equals("main")) {
            func_name = "main";
            vtx = _vertix_map[func_name];
        }
        else {
            func_name = getFuncName(func->getName());
            if(_vertix_map.find(func_name) == _vertix_map.end()) return;
            vtx = _vertix_map[func_name];
        }
    }
    else {
        func_name = func->getName().str();
        if(_vertix_map.find(func_name) == _vertix_map.end()) return;
        vtx = _vertix_map[func_name];
    }
    if(!vtx) return;
    // if(strcmp("_ZN3AES12EncryptBlockEPKhPhS2_", func->getName().data()) == 0) {
    //     int a = 0;
    // }
    for(auto &bblk : *func) {
        for(auto &inst : bblk) {
            call = llvm::dyn_cast<llvm::CallInst>(&inst);
            f_name = "";
            if(call && call->getCalledFunction()) {
                if(_use_cpp) {
                    f_name = getFuncName(call->getCalledFunction()->getName());
                }
                else {
                    f_name = call->getCalledFunction()->getName().str();
                }
            }
            invoke = llvm::dyn_cast<llvm::InvokeInst>(&inst);
            if(invoke) {
                if(_use_cpp) {
                    f_name = getFuncName(invoke->getCalledFunction()->getName());
                }
                else {
                    f_name = invoke->getCalledFunction()->getName().str();
                }
            }
            if(strcmp(f_name.c_str(), "") == 0) continue;
            if(edgemap.find(f_name) != edgemap.end()) {
                continue;
            }
            edgemap.insert(f_name);
            if(_vertix_map.find(f_name) != _vertix_map.end()) {
                edge = new Edge();
                edge->u = vtx;
                edge->v = _vertix_map[f_name];
                vtx->_edges.push_back(edge);
            }
        }
    }
}


int AnalysisCore::dumpAll() {
    std::error_code ec;
    std::filesystem::path out = _output, out1 = out / "image";
    std::string cmd;
    if(!std::filesystem::exists(out)) {
        std::filesystem::create_directory(out, ec);
        if(ec) {
            std::cerr << "Can not create dir " << out << "\n";
            return UNKNOWN_ERROR;
        }
    }
    if(!std::filesystem::exists(out1)) {
        std::filesystem::create_directory(out1, ec);
        if(ec) {
            std::cerr << "Can not create dir " << out1 << "\n";
            return UNKNOWN_ERROR;
        }
    }
    
    std::filesystem::path p1, p2;
    std::vector<FuncVertex *> vtx;
    for(auto &it : _vertix_map) {
        vtx.push_back(it.second);
    }

    p1 = out / "res.dot";
    p2 = out1 / "res.png";
    if(auto ret = dump(p1.c_str(), vtx)) {
        return ret;
    }
    cmd = "dot -T";
    cmd.append(_format).append(" ").append(p1.string())
        .append(" -o ").append(p2.string());
    if(auto ret = system(cmd.c_str())) {
        return ret;
    }
    return OK;
}