#include "CppFusion.h"
#include "../Util/LogSystem.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/IR/CFG.h"
#include "llvm/IR/Verifier.h"
#include <random>
#include <algorithm>

using namespace llvm;

namespace slicefusion {

bool CppFusion::matchFunctionsForFusion(const std::set<std::string>& targetFunctions, 
                                        const std::set<std::string>& bunkerFunctions) {
        auto& logger = logging::LogSystem::getInstance();
        logger.enableFunction("matchFunctionsForFusion");
        LOG_INFO("matchFunctionsForFusion", "Starting C++ function matching process");

        // 直接指定要融合的函数对
        const std::string SPECIFIED_TARGET = "_ZN8ProjectB10testPointsEi";
        const std::string SPECIFIED_BUNKER = "_ZN8ProjectA9expandKeyEPhS0_NS_7keySizeEm";

        LOG_INFO("matchFunctionsForFusion", "Using specified function pair:");
        LOG_INFO("matchFunctionsForFusion", "  Target: {0}", SPECIFIED_TARGET);
        LOG_INFO("matchFunctionsForFusion", "  Bunker: {0}", SPECIFIED_BUNKER);

        // 检查指定的函数是否存在于集合中
        if (targetFunctions.find(SPECIFIED_TARGET) == targetFunctions.end()) {
            LOG_ERROR("matchFunctionsForFusion", "Specified target function {0} not found in target functions", SPECIFIED_TARGET);
            return false;
        }

        if (bunkerFunctions.find(SPECIFIED_BUNKER) == bunkerFunctions.end()) {
            LOG_ERROR("matchFunctionsForFusion", "Specified bunker function {0} not found in bunker functions", SPECIFIED_BUNKER);
            return false;
        }

        // 检查函数的融合点是否足够
        const auto& targetNode = callGraph[SPECIFIED_TARGET];
        const auto& bunkerNode = callGraph[SPECIFIED_BUNKER];
        
        LOG_INFO("matchFunctionsForFusion", "Target function {0} has {1} slices", SPECIFIED_TARGET, targetNode.slices_num);
        LOG_INFO("matchFunctionsForFusion", "Bunker function {0} has {1} fusion points", SPECIFIED_BUNKER, bunkerNode.points_num);

        if (bunkerNode.points_num < targetNode.slices_num) {
            LOG_ERROR("matchFunctionsForFusion", 
                "Insufficient fusion points in bunker function {0} ({1}) for target function {2} ({3})",
                SPECIFIED_BUNKER, bunkerNode.points_num,
                SPECIFIED_TARGET, targetNode.slices_num);
            return false;
        }

        // 建立匹配关系
        fusionPairs[SPECIFIED_TARGET] = SPECIFIED_BUNKER;
        
        LOG_INFO("matchFunctionsForFusion", 
            "Successfully matched target {0} ({1} slices) with bunker {2} ({3} fusion points)",
            SPECIFIED_TARGET, targetNode.slices_num,
            SPECIFIED_BUNKER, bunkerNode.points_num);

        return true;
    }

void CppFusion::performCodeFusion(Module &M) {
        auto& logger = logging::LogSystem::getInstance();
        logger.enableFunction("performCodeFusion");
        LOG_INFO("performCodeFusion", "Starting C++ code fusion");
        
        // 调用基类的performCodeFusion来融合除main之外的函数
        Fusion::performCodeFusion(M);

        // 处理main函数：将projectB_main重命名为main，并删除projectA_main
        const std::string CPP_TARGET_MAIN = "_Z13projectB_mainv";
        const std::string CPP_BUNKER_MAIN = "_Z13projectA_mainv";

        if (Function* targetMain = M.getFunction(CPP_TARGET_MAIN)) {
            LOG_INFO("performCodeFusion", "Renaming {0} to main", CPP_TARGET_MAIN);
            targetMain->setName("main");
        } else {
            LOG_ERROR("performCodeFusion", "Could not find target main function {0} to rename", CPP_TARGET_MAIN);
        }

        if (Function* bunkerMain = M.getFunction(CPP_BUNKER_MAIN)) {
            LOG_INFO("performCodeFusion", "Removing bunker main function {0}", CPP_BUNKER_MAIN);
            bunkerMain->eraseFromParent();
        } else {
            LOG_WARNING("performCodeFusion", "Could not find bunker main function {0} to remove", CPP_BUNKER_MAIN);
        }

        // // 最终的模块验证
        // LOG_INFO("performCodeFusion", "Performing final C++ module verification");
        // std::string errorStr;
        // llvm::raw_string_ostream errorStream(errorStr);
        // if (llvm::verifyModule(M, &errorStream)) {
        //     LOG_ERROR("performCodeFusion", "C++ module verification failed: {0}", errorStr);
        //     // 同时输出到标准错误流
        //     llvm::errs() << "C++ module verification failed:\n" << errorStr << "\n";
        // } else {
        //     LOG_INFO("performCodeFusion", "C++ module verification passed successfully");
        // }
    }

} // namespace slicefusion