// Copyright 2018-2023 Bloomberg Finance L.P
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_localexecutionclient.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_runner.h>
#include <buildboxcommon_stringutils.h>
#include <buildboxcommon_systemutils.h>
#include <buildboxcommon_temporaryfile.h>

#include <google/rpc/code.pb.h>
#include <google/rpc/status.pb.h>

#include <csignal>
#include <fcntl.h>
#include <functional>
#include <optional>
#include <sys/wait.h>
#include <thread>
#include <unistd.h>
#include <unordered_map>
#include <vector>

#ifdef __linux__
#include <sys/prctl.h>
#endif

#define POLL_WAIT std::chrono::milliseconds(500)

namespace buildboxcommon {

void LocalExecutionClient::setRunner(
    const std::string &runnerCommand,
    const std::vector<std::string> &extraRunArgs)
{
    this->d_runnerPath =
        buildboxcommon::SystemUtils::getPathToCommand(runnerCommand);
    if (this->d_runnerPath.empty()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Could not find runner command \"" +
                                           runnerCommand + "\"");
    }

    d_extraRunArgs.clear();
    d_extraRunArgs.insert(d_extraRunArgs.end(), extraRunArgs.cbegin(),
                          extraRunArgs.cend());
}

void LocalExecutionClient::setCacheFailures(bool cacheFailures)
{
    this->d_cacheFailures = cacheFailures;
}

void LocalExecutionClient::init(
    std::shared_ptr<ActionCache::StubInterface> actionCacheStub)
{
    this->ExecutionClient::init(actionCacheStub);

    if (d_runnerPath.empty()) {
        this->setRunner("buildbox-run");
    }
}

void LocalExecutionClient::init()
{
    this->ExecutionClient::init();

    if (d_runnerPath.empty()) {
        this->setRunner("buildbox-run");
    }
}

CASClient *LocalExecutionClient::getCASClient()
{
    std::lock_guard<std::mutex> lock(this->d_mutex);

    if (!d_casClient) {
        auto grpcClient = std::make_shared<GrpcClient>();
        grpcClient->init(d_casServer);
        d_casClient = std::make_unique<CASClient>(grpcClient);
        d_casClient->init();
    }

    return d_casClient.get();
}

std::vector<std::string> buildRunnerCommand(
    const std::string &runnerBinaryPath,
    const std::vector<std::string> &runnerExtraArguments,
    const buildboxcommon::ConnectionOptions &casConnectionOptions,
    const buildboxcommon::ConnectionOptions &assetConnectionOptions,
    const std::string &actionFilePath, const std::string &actionResultFilePath,
    const std::string &stdoutFilePath, const std::string &stderrFilePath,
    const std::optional<std::string> &partialExecutionMetadataFile)
{
    std::vector<std::string> command = {runnerBinaryPath};

    command.insert(command.end(), runnerExtraArguments.cbegin(),
                   runnerExtraArguments.cend());
    casConnectionOptions.putArgs(&command);
    assetConnectionOptions.putArgs(&command, "ra-");

    command.emplace_back("--action=" + actionFilePath);
    command.emplace_back("--action-result=" + actionResultFilePath);
    command.emplace_back("--stdout-file=" + stdoutFilePath);
    command.emplace_back("--stderr-file=" + stderrFilePath);
    command.emplace_back(
        "--digest-function=" +
        DigestFunction_Value_Name(DigestGenerator::digestFunction()));
    if (partialExecutionMetadataFile) {
        command.emplace_back("--partial-execution-metadata-file=" +
                             partialExecutionMetadataFile.value());
    }

    auto &logger = logging::Logger::getLoggerInstance();

    const auto it = logging::logLevelToStringMap().find(logger.getLogLevel());
    if (it != logging::logLevelToStringMap().end()) {
        command.emplace_back("--log-level=" + it->second);
    }

    const std::string logDirectory = logger.getOutputDirectory();
    if (!logDirectory.empty()) {
        command.emplace_back("--log-directory=" + logDirectory);
    }

    return command;
}

bool LocalExecutionClient::validateRunner()
{
    TemporaryFile dummyActionFile;
    TemporaryFile dummyActionResultFile;
    TemporaryFile stdoutFile;
    TemporaryFile stderrFile;
    TemporaryFile partialExecutionMetadataFile;

    auto command = buildRunnerCommand(
        this->d_runnerPath, this->d_extraRunArgs, this->d_casServer,
        this->d_assetServer, dummyActionFile.strname(),
        dummyActionResultFile.strname(), stdoutFile.strname(),
        stderrFile.strname(), partialExecutionMetadataFile.strname());
    command.emplace_back("--validate-parameters");

    try {
        const int exitCode = SystemUtils::executeCommandAndWait(command);
        return (exitCode == 0);
    }
    catch (const std::exception &e) {
        BUILDBOX_LOG_ERROR(
            "Failed to launch runner command for parameter validation ["
            << buildboxcommon::logging::printableCommandLine(command)
            << "]: " << e.what())
        return false;
    }
}

/**
 * Start a subprocess with the given command and return its PID.
 *
 * The first element in `command` must be a path to an executable file.
 */
pid_t runCommandInSubprocess(const std::vector<std::string> &command)
{
    const pid_t pid = fork();

    if (pid == -1) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category, "error in fork()");
    }

    else if (pid == 0) {
        // (runs only in the child)

#ifdef __linux__
        // Terminate runner if parent dies unexpectedly.
        if (prctl(PR_SET_PDEATHSIG, SIGTERM) < 0) {
            perror("prctl(PR_SET_PDEATHSIG, SIGTERM) failed");
        }
#endif

        // Signals should not be ignored across exec
        struct sigaction sa{};
        sa.sa_handler = SIG_DFL;
        sigemptyset(&sa.sa_mask);
        sa.sa_flags = 0;
        if (sigaction(SIGPIPE, &sa, nullptr) == -1) {
            perror("Unable to reset signal handler for SIGPIPE in subprocess");
            _Exit(1);
        }

        const int devNullFD = open("/dev/null", O_RDWR | O_CLOEXEC);
        if (devNullFD == -1) {
            perror("/dev/null unavailable");
            _Exit(1);
        }

        // Set the group pid of the subprocess to its pid.
        if (setpgid(0, 0) != 0) {
            perror("Unable to set group process id to subprocess");
            _Exit(1);
        }

        // Prevent action commands from reading our stdin
        dup2(devNullFD, STDIN_FILENO);
        close(devNullFD);

        const auto exit_code =
            buildboxcommon::SystemUtils::executeCommand(command, true);

        perror(nullptr);
        _Exit(exit_code);
    }

    else {
        // Set the group pid of the subprocess to its pid.
        // This is to prevent a race condition in which the child processes
        // pgid is checked in the parent before it is set in the child. To
        // prevent this, we also set it in the parent.
        if (setpgid(pid, 0) != 0) {
            // EACCES would be set if the child had exec'd. That is okay in
            // this case, as the pgid is set in the child as well.
            if (errno != EACCES) {
                BUILDBOX_LOG_ERROR(
                    "Unable to set group process id to subprocess, due to: "
                    << strerror(errno));
            }
        }
        BUILDBOX_LOG_DEBUG("Forked runner process: pid ["
                           << pid << "], group-pid [" << getpgid(pid) << "]");

        return pid;
    }
}

// Wait for the runner subprocess to exit and return its exit code
// throws a std::system_error exception if waitpid has an error
int waitForRunner(pid_t subprocessPid, bool wait)
{
    while (true) { // Waiting until the process exits or get signaled...
        int status = 0;
        const pid_t pid = waitpid(subprocessPid, &status, wait ? 0 : WNOHANG);

        if (pid == -1) {
            const auto waitpid_error = errno;

            if (waitpid_error == EINTR) {
                continue; // Not an error, we can keep waiting.
            }

            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Error while waiting for runner subprocess");
        }
        if (!wait && pid == 0) {
            // Process still running
            return -1;
        }
        if (WIFSIGNALED(status)) {
            BUILDBOX_LOG_INFO("Runner received signal " << WTERMSIG(status));
            const int SIGNAL_OFFSET = 128;
            return SIGNAL_OFFSET + WTERMSIG(status);
        }
        if (WIFEXITED(status)) {
            BUILDBOX_LOG_INFO("Runner exited with code "
                              << WEXITSTATUS(status));
            return WEXITSTATUS(status);
        }
    }
}

buildboxcommon::ActionResult readActionResultFile(const std::string &path)
{
    const std::string actionResultRaw =
        buildboxcommon::FileUtils::getFileContents(path.c_str());
    if (actionResultRaw.empty()) {
        // The ActionResult file is empty, and an empty ActionResult isn't
        // valid
        const std::string error_message =
            "ActionResultFile " + path + " is an empty file";
        BUILDBOX_LOG_WARNING(error_message);
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error, error_message);
    }

    buildboxcommon::ActionResult actionResult;
    if (!actionResult.ParseFromString(actionResultRaw)) {
        const std::string error_message = "Unable to parse ActionResultFile " +
                                          path + ": content of file was [" +
                                          actionResultRaw + "]";
        BUILDBOX_LOG_ERROR(error_message);
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error, error_message);
    }

    return actionResult;
}

Digest serializeAction(const Action &action, const std::string &actionFile)
{
    BUILDBOX_LOG_DEBUG("Writing Action to [" << actionFile << "]");

    const std::string serializedAction = action.SerializeAsString();
    Digest digest = DigestGenerator::hash(serializedAction);

    FileUtils::writeFileAtomically(actionFile, serializedAction);

    return digest;
}

google::longrunning::Operation LocalExecutionClient::asyncExecuteAction(
    const Action &action, const std::string &_stdoutFile,
    const std::string &_stderrFile,
    const std::string &partialExecutionMetadataFile)
{
    const std::string uuid = StringUtils::getUUIDString();

    google::longrunning::Operation operation;
    operation.set_name(uuid);

    const std::string prefix = d_tempDirectory.strname() + "/" + uuid;
    const std::string actionFile = prefix + ".action";
    const std::string actionResultFile = prefix + ".actionresult";
    const std::string stdoutFile =
        !_stdoutFile.empty() ? _stdoutFile : prefix + ".stdout";
    const std::string stderrFile =
        !_stderrFile.empty() ? _stderrFile : prefix + ".stderr";

    InternalOperation internalOperation;
    internalOperation.action = action;
    internalOperation.actionDigest = serializeAction(action, actionFile);

    std::vector<std::string> command = buildRunnerCommand(
        this->d_runnerPath, this->d_extraRunArgs, this->d_casServer,
        this->d_assetServer, actionFile, actionResultFile, stdoutFile,
        stderrFile, partialExecutionMetadataFile);

    BUILDBOX_LOG_DEBUG(
        "Launching runner process: "
        << buildboxcommon::logging::printableCommandLine(command));
    internalOperation.pid = runCommandInSubprocess(command);

    std::lock_guard<std::mutex> lock(this->d_mutex);
    d_runningOperations.emplace(uuid, internalOperation);

    return operation;
}

google::longrunning::Operation LocalExecutionClient::asyncExecuteAction(
    const Digest &actionDigest, const std::atomic_bool & /* stop_requested */,
    bool skipCache, const ExecutionPolicy * /* executionPolicy */)
{
    if (!skipCache && this->getActionCacheStub()) {
        ExecuteResponse executeResponse;
        if (fetchFromActionCache(actionDigest, {},
                                 executeResponse.mutable_result())) {
            executeResponse.set_cached_result(true);

            BUILDBOX_LOG_DEBUG("Action Cache hit for [" << actionDigest
                                                        << "]");

            google::longrunning::Operation operation;
            operation.set_name(StringUtils::getUUIDString());
            operation.set_done(true);
            operation.mutable_response()->PackFrom(executeResponse);

            return operation;
        }
    }

    auto casClient = this->getCASClient();
    auto action = casClient->fetchMessage<Action>(actionDigest);

    return asyncExecuteAction(action);
}

ActionResult LocalExecutionClient::executeAction(
    const Digest &actionDigest, const std::atomic_bool &stopRequested,
    bool skipCache, const ExecutionPolicy *executionPolicy)
{
    auto operation = asyncExecuteAction(actionDigest, stopRequested, skipCache,
                                        executionPolicy);

    while (!operation.done()) {
        if (stopRequested) {
            BUILDBOX_LOG_WARNING(
                "Cancelling job, operation name: " << operation.name());
            cancelOperation(operation.name());
            operation = waitExecution(operation.name());
            break;
        }
        std::this_thread::sleep_for(POLL_WAIT);
        operation = getOperation(operation.name());
    }

    return getActionResult(operation);
}

google::longrunning::Operation
LocalExecutionClient::getOperationInternal(const std::string &operationName,
                                           bool wait)
{
    google::longrunning::Operation operation;
    operation.set_name(operationName);
    pid_t pid = 0;

    {
        std::lock_guard<std::mutex> lock(this->d_mutex);
        pid = d_runningOperations.at(operationName).pid;
    }

    const std::string prefix = d_tempDirectory.strname() + "/" + operationName;
    const std::string actionResultFile = prefix + ".actionresult";

    ExecuteResponse executeResponse;
    google::rpc::Status *status = executeResponse.mutable_status();
    buildboxcommon::ActionResult *actionResult =
        executeResponse.mutable_result();

    const int rc = waitForRunner(pid, wait);
    if (rc < 0) {
        // Not yet complete
        return operation;
    }

    // Get internal operation object after the wait as `cancelled` may have
    // been updated.
    InternalOperation internalOperation;
    {
        std::lock_guard<std::mutex> lock(this->d_mutex);
        internalOperation = d_runningOperations.at(operationName);
    }

    status->set_code(google::rpc::Code::OK);

    if (rc != 0) {
        const std::string msg = "Runner exited with non-zero exit code [" +
                                std::to_string(rc) + "]";
        BUILDBOX_LOG_WARNING(msg);
        constexpr int SIGNAL_OFFSET = 128;
        if (internalOperation.cancelled && rc == SIGNAL_OFFSET + SIGTERM) {
            status->set_code(google::rpc::Code::CANCELLED);
            status->set_message("The operation was cancelled");
        }
        else {
            // The runner might have written a Status file detailing what went
            // wrong. We'll attempt to read it:
            const std::string statusFilePath =
                buildboxcommon::Runner::errorStatusCodeFilePath(
                    actionResultFile);
            if (!buildboxcommon::Runner::readStatusFile(statusFilePath,
                                                        status)) {
                status->set_code(google::rpc::Code::INTERNAL);
                status->set_message(msg);
            }
        }
    }

    try {
        *actionResult = readActionResultFile(actionResultFile);
        BUILDBOX_LOG_DEBUG("ActionResult " << actionResult->DebugString());
        if (status->code() == google::rpc::Code::OK &&
            this->getActionCacheStub() &&
            this->isActionCacheUpdatesAllowed() &&
            !internalOperation.action.do_not_cache() &&
            (this->d_cacheFailures || actionResult->exit_code() == 0)) {
            updateActionCache(internalOperation.actionDigest, *actionResult);
        }
    }
    catch (const std::runtime_error &e) {
        // ActionResult is required if runner exits with 0.
        // It may or may not be available in case of failure.
        // E.g. stdout/stderr should be available on action timeout.
        if (rc == 0) {
            BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                           "Error reading ActionResult file "
                                           "after runner exited with code 0: "
                                               << e.what());
        }
    }

    operation.set_done(true);
    operation.mutable_response()->PackFrom(executeResponse);

    std::lock_guard<std::mutex> lock(this->d_mutex);
    d_runningOperations.erase(operationName);

    return operation;
}

google::longrunning::Operation
LocalExecutionClient::getOperation(const std::string &operationName)
{
    return getOperationInternal(operationName, false /* wait */);
}

google::longrunning::Operation
LocalExecutionClient::waitExecution(const std::string &operationName)
{
    return getOperationInternal(operationName, true /* wait */);
}

bool LocalExecutionClient::cancelOperation(const std::string &operationName)
{
    std::lock_guard<std::mutex> lock(this->d_mutex);

    auto it = d_runningOperations.find(operationName);
    if (it == d_runningOperations.end()) {
        return false;
    }

    const auto subprocessPid = it->second.pid;

    if (killpg(subprocessPid, SIGTERM) != 0) {
        if (errno != ESRCH) {
            BUILDBOX_LOG_ERROR("Unable to send SIGTERM to subprocess : ["
                               << subprocessPid
                               << "]. Due to: " << strerror(errno));
        }
        return false;
    }

    it->second.cancelled = true;

    return true;
}

void LocalExecutionClient::cancelAllOperations()
{
    std::lock_guard<std::mutex> lock(this->d_mutex);

    for (auto &operation : d_runningOperations) {
        const auto subprocessPid = operation.second.pid;
        if (killpg(subprocessPid, SIGTERM) == 0) {
            operation.second.cancelled = true;
        }
        else if (errno != ESRCH) {
            BUILDBOX_LOG_ERROR("Unable to send SIGTERM to subprocess : ["
                               << subprocessPid
                               << "]. Due to: " << strerror(errno));
        }
    }
}

google::longrunning::ListOperationsResponse
LocalExecutionClient::listOperations(const std::string & /* name */,
                                     const std::string & /* filter */,
                                     int /* page_size */,
                                     const std::string & /* page_token */)
{
    BUILDBOXCOMMON_THROW_EXCEPTION(
        std::runtime_error,
        "LocalExecutionClient::listOperations() is not supported");
}

} // namespace buildboxcommon
