// Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
// This source file is part of the Cangjie project, licensed under Apache-2.0
// with Runtime Library Exception.
//
// See https://cangjie-lang.cn/pages/LICENSE for license information.

/**
 * @file
 *
 * This file implements the Driver.
 */

#include "cangjie/Driver/Driver.h"

#include <string>
#include <utility>
#include <vector>

#include "llvm/Bitcode/BitcodeReader.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/MemoryBuffer.h"

#include "cangjie/Basic/Print.h"
#include "cangjie/Basic/Version.h"
#include "cangjie/Driver/TempFileManager.h"
#include "cangjie/Driver/Utils.h"
#include "cangjie/FrontendTool/DefaultCompilerInstance.h"
#include "cangjie/FrontendTool/IncrementalCompilerInstance.h"
#include "cangjie/FrontendTool/FrontendTool.h"
#include "cangjie/Utils/CheckUtils.h"
#include "cangjie/Utils/FileUtil.h"
#include "cangjie/Utils/ProfileRecorder.h"
// DO NOT remove this header file, otherwise, the signal handling of Driver will fail:
#include "cangjie/Utils/Signal.h"
#include "cangjie/Utils/Utils.h"
#include "Job.h"

using namespace Cangjie;

Driver::Driver(const std::vector<std::string>& args, DiagnosticEngine& diag, const std::string& exeName)
    : args(args), diag(diag), executableName(exeName)
{
    optionTable = CreateOptionTable(false);
    CJC_ASSERT(optionTable && "create option table failed");
    argList = std::make_unique<ArgList>();
    driverOptions = std::make_unique<DriverOptions>();

    // Get absolute path of `cjc` executable program.
    driverOptions->executablePath = executableName;
    cangjieHome = FileUtil::GetDirPath(FileUtil::GetDirPath(FileUtil::GetAbsPath(exeName) | FileUtil::IdenticalFunc));
}

namespace {
bool StartsWith(const std::string str, const std::string prefix)
{
    return str.rfind(prefix, 0) == 0;
}

bool IsWarningArg(const std::string& arg)
{
    return StartsWith(arg, "-Woff") || StartsWith(arg, "--warn-off") || StartsWith(arg, "-Won") ||
        StartsWith(arg, "--warn-on");
}

void FrontendCmdPrint(const GlobalOptions& globalOptions, const std::string& cangjieHome,
    const std::vector<std::string>& args)
{
    CJC_ASSERT(args.size() != 0);
    std::string binPath = FileUtil::JoinPath(cangjieHome, "bin");
    std::string exeExtension =
#ifdef _WIN32
        ".exe";
#else
        "";
#endif
    binPath = FileUtil::JoinPath(binPath, "cjc-frontend" + exeExtension);
    std::string frontendCmd = GetCommandLineArgumentQuoted(binPath);
    std::vector<std::string> optionList = globalOptions.GenerateFrontendOptions();
    for (auto& str : optionList) {
        frontendCmd += " " + GetCommandLineArgumentQuoted(str);
    }
    Println(FileUtil::Normalize(frontendCmd));
}

bool DeleteInstance(DefaultCompilerInstance* instance)
{
    Utils::ProfileRecorder recorder("DeleteInstance", "CompleteTime");
    delete instance;
    Utils::FreeIdleMemoryToOS();
    return true;
}
} // namespace

bool Driver::ParseArgs()
{
    CJC_NULLPTR_CHECK(optionTable);
    if (!optionTable->ParseArgs(args, *argList)) {
        return false;
    }

    // All the warning options need to be preceded by other options.
    // Otherwise, the warnings generated by some options cannot be applied to.
    std::vector<std::unique_ptr<ArgInstance>> warningArgs;
    std::vector<std::unique_ptr<ArgInstance>> otherArgs;
    CJC_NULLPTR_CHECK(argList);
    for (auto& arg : argList->args) {
        if (IsWarningArg(arg->str)) {
            warningArgs.push_back(std::move(arg));
        } else {
            otherArgs.push_back(std::move(arg));
        }
    }
    argList->args.clear();
    argList->args.resize(warningArgs.size() + otherArgs.size());
    std::transform(
        warningArgs.begin(), warningArgs.end(), argList->args.begin(), [](auto& it) { return std::move(it); });
    std::transform(otherArgs.begin(), otherArgs.end(), argList->args.begin() + static_cast<long>(warningArgs.size()),
        [](auto& it) { return std::move(it); });

    CJC_NULLPTR_CHECK(driverOptions);
    if (!driverOptions->ParseFromArgs(*argList)) {
        return false;
    }

    driverOptions->SetCompilationCachedPath();

    if (!ExtractBCPackageNames()) {
        return false;
    }

    return true;
}

bool Driver::ExtractBCPackageNames()
{
    for (const auto& bcPath : driverOptions->bcInputFiles) {
        auto mbOrErr = llvm::MemoryBuffer::getFile(bcPath, false, false);
        if (auto ec = mbOrErr.getError()) {
            diag.DiagnoseRefactor(DiagKindRefactor::bc_file_open_failed, DEFAULT_POSITION, bcPath, ec.message());
            return false;
        }
        
        llvm::MemoryBufferRef mbref = (*mbOrErr)->getMemBufferRef();
        llvm::LLVMContext context;
        auto moduleOrErr = llvm::getLazyBitcodeModule(mbref, context, false);
        if (!moduleOrErr) {
            diag.DiagnoseRefactor(DiagKindRefactor::bc_parse_failed, DEFAULT_POSITION, bcPath);
            return false;
        }

        std::unique_ptr<llvm::Module> module = std::move(*moduleOrErr);
        llvm::StringRef pkgName = module->getSourceFileName();
        unsigned subCHIRPackageIdx = 0;
        unsigned Radix = 10;
        if (pkgName.consumeInteger(Radix, subCHIRPackageIdx) ||
            !pkgName.consume_front("-") || pkgName.empty()) {
                diag.DiagnoseRefactor(DiagKindRefactor::bc_invalid_package_name, DEFAULT_POSITION, bcPath);
                return false;
        }
        driverOptions->bcPackageNames.push_back(std::string(pkgName));
    }

    return true;
}

void Driver::EnvironmentSetup(const std::unordered_map<std::string, std::string>& environmentVars)
{
    CJC_NULLPTR_CHECK(driverOptions);
    driverOptions->ReadPathsFromEnvironmentVars(environmentVars);
    driverOptions->cangjieHome = driverOptions->environment.cangjieHome.value_or(cangjieHome);
}

bool Driver::ExecuteCompilation() const
{
    if (driverOptions == nullptr) {
        return false;
    }
    diag.RegisterHandler(driverOptions->diagFormat);
    // Do nothing else but show usage.
    if (driverOptions->showUsage) {
        const std::set<Options::Group> groups{Options::Group::GLOBAL, Options::Group::DRIVER};
        optionTable->Usage(driverOptions->GetOptionsBackend(), groups, driverOptions->experimentalMode);
        if (!argList->GetInputs().empty()) {
            Warningln("cjc doesn't do compilation job in show usage mode.");
        }
        return true;
    }

    if (driverOptions->enableVerbose || driverOptions->printVersionOnly) {
        Cangjie::PrintVersion();
    }

    if (driverOptions->printVersionOnly) {
        if (!argList->GetInputs().empty()) {
            Warningln("cjc doesn't do compilation job in version check mode.");
        }
        return true;
    }
    if (driverOptions->enableTimer) {
        Utils::ProfileRecorder::Enable(true, Utils::ProfileRecorder::Type::TIMER);
    }
    if (driverOptions->enableMemoryCollect) {
        Utils::ProfileRecorder::Enable(true, Utils::ProfileRecorder::Type::MEMORY);
    }
    if (!TempFileManager::Instance().Init(*driverOptions, false)) {
        return false;
    }
    CompilerInvocation compilerInvocation;

    compilerInvocation.globalOptions = *driverOptions;
    compilerInvocation.globalOptions.executablePath = executableName;
    compilerInvocation.globalOptions.environment = driverOptions->environment;

    // In Driver mode, FrontendOptions is never parsed, we need to do setup manually.
    if (compilerInvocation.globalOptions.scanDepPkg) {
        compilerInvocation.frontendOptions.dumpAction = FrontendOptions::DumpAction::DUMP_DEP_PKG;
    }

    diag.SetErrorCountLimit(compilerInvocation.globalOptions.errorCountLimit);
    diag.RegisterHandler(driverOptions->diagFormat);
    DefaultCompilerInstance* instance = nullptr;
    if (NeedCreateIncrementalCompilerInstance(compilerInvocation.globalOptions)) {
        instance = new IncrementalCompilerInstance(compilerInvocation, diag);
    } else {
        instance = new DefaultCompilerInstance(compilerInvocation, diag);
    }
    if (!ExecuteFrontendByDriver(*instance, *this)) {
        DeleteInstance(instance);
        return false;
    }
    if (driverOptions->enableVerbose) {
        FrontendCmdPrint(*driverOptions.get(), cangjieHome, args);
    }
    // cjc main.cj or cjc main.o
    bool noBackendInputs = driverOptions->frontendOutputFiles.empty() && driverOptions->inputObjs.empty();
    if (noBackendInputs || driverOptions->IsEmitCHIREnable() ||
        driverOptions->outputMode == GlobalOptions::OutputMode::CHIR) {
        DeleteInstance(instance);
        return true;
    }

    std::future<bool> future = std::async(DeleteInstance, instance);
    bool res = InvokeCompileToolchain();
    Utils::ProfileRecorder::Start("Main Stage", "DeleteInstanceLeftTime");
    future.get();
    Utils::ProfileRecorder::Stop("Main Stage", "DeleteInstanceLeftTime");
    return res;
}

bool Driver::InvokeCompileToolchain() const
{
#ifdef SIGNAL_TEST
    // The interrupt signal triggers the function. In normal cases, this function does not take effect.
    Cangjie::SignalTest::ExecuteSignalTestCallbackFunc(Cangjie::SignalTest::TriggerPointer::DRIVER_POINTER);
#endif
    // Job Initialization.
    std::unique_ptr<Job> job = std::make_unique<Job>();
    if (!job->Assemble(*driverOptions.get(), *this)) {
        return false;
    }

    // Execute Job.
    bool executeResult = job->Execute();
    return executeResult;
}