/*
 * Copyright 2018 Bloomberg Finance LP
 *
 * 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_remoteexecutionclient.h>
#include <buildboxrun_hosttools.h>
#include <buildboxrun_hosttools_changedirectoryguard.h>
#include <buildboxrun_hosttools_pathprefixutils.h>

#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_runner.h>
#include <buildboxcommon_scopeguard.h>
#include <buildboxcommon_stringutils.h>
#include <buildboxcommon_systemutils.h>

#include <fstream>
#include <sstream>
#include <stdlib.h>
#include <unistd.h>

extern "C" char **environ;

namespace buildboxcommon {
namespace buildboxrun {
namespace hosttools {

std::string join(const std::vector<std::string> &tokens,
                 const std::string &spacer)
{
    if (tokens.empty()) {
        return "";
    }

    std::ostringstream ss;
    bool first = true;
    for (const std::string &token : tokens) {
        if (first) {
            ss << token;
            first = false;
        }
        else {
            ss << spacer << token;
        }
    }
    return ss.str();
}

std::unordered_map<std::string, std::string>
parseEnvFile(const std::string &path)
{
    std::ifstream config(path);
    if (!config.good()) {
        std::cerr << "Failed to open env file: " << path << "\n";
        exit(1);
    }

    std::string line;
    std::unordered_map<std::string, std::string> environment;

    while (getline(config, line)) {
        // Remove comment
        size_t hashPos = line.find('#');
        if (hashPos != std::string::npos) {
            line = line.substr(0, hashPos);
        }

        line = buildboxcommon::StringUtils::trim(line);

        if (line.empty()) {
            continue;
        }

        const auto assignPos = line.find('=');
        if (assignPos == std::string_view::npos) {
            BUILDBOX_LOG_ERROR("Invalid line in environment file: " << line);
            continue;
        }

        environment.emplace(line.substr(0, assignPos),
                            line.substr(assignPos + 1));
    }

    return environment;
}

void HostToolsRunner::setUpEnvironment(const Command &command) const
{
    if (this->d_environment.has_value()) {
        // Clear environment
        environ = nullptr;

        for (const auto &envVar : *this->d_environment) {
            if (setenv(envVar.first.c_str(), envVar.second.c_str(), 1) == -1) {
                BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                    std::system_error, errno, std::system_category,
                    "Unable to set base environment " << envVar.first << "="
                                                      << envVar.second);
            }
        }
    }

    for (const auto &envVar : command.environment_variables()) {
        if (setenv(envVar.name().c_str(), envVar.value().c_str(), 1) == -1) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Unable to set environment " << envVar.name() << "="
                                             << envVar.value());
        }
    }

    if (this->d_actionDigestEnv.has_value()) {
        const std::string actionDigestStr = toString(this->d_action_digest);
        if (setenv(this->d_actionDigestEnv->c_str(), actionDigestStr.c_str(),
                   1) == -1) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Unable to set action digest environment variable "
                    << *this->d_actionDigestEnv << "=" << actionDigestStr);
        }
    }
}

void HostToolsRunner::createParentDirectories(
    const Command &command, const std::string &workingDir) const
{
    BUILDBOX_RUNNER_LOG(
        DEBUG, "Creating parent directories for Command output paths");
    for (const std::string &path : command.output_paths()) {
        if (path.find('/') != std::string::npos) {
            FileUtils::createDirectory(
                (workingDir + "/" + path.substr(0, path.rfind('/'))).c_str());
        }
    }
}

std::vector<std::string>
HostToolsRunner::generateCommandLine(const Command &command,
                                     const std::string &stage_path) const
{
    std::vector<std::string> commandLine(command.arguments().cbegin(),
                                         command.arguments().cend());

    // Perform PATH lookup as specified in REAPI v2.3
    const std::string resolvedCommand =
        SystemUtils::getPathToCommand(commandLine[0]);
    if (!resolvedCommand.empty()) {
        commandLine[0] = resolvedCommand;
    }

    if (this->d_prefix) {
        return PathPrefixUtils::prefixAbspathsWithStagedDir(commandLine,
                                                            stage_path);
    }

    return commandLine;
}

std::optional<std::unordered_map<std::string, std::string>>
HostToolsRunner::getEnvironment() const
{
    return this->d_environment;
}

const std::optional<std::string> &HostToolsRunner::getActionDigestEnv() const
{
    return this->d_actionDigestEnv;
}

ActionResult HostToolsRunner::execute(const Command &command,
                                      const Digest &inputRootDigest,
                                      const Platform &platform)
{
    // Ensure that the command is in our allow list
    const std::string &executable = command.arguments(0);
    if (!executableInAllowedExecutables(executable)) {
        std::ostringstream errorMessage;
        errorMessage << executable
                     << " not found in allowed executable list: ";
        errorMessage << "[";
        std::string seperator = "";
        for (const std::string &item : this->d_allowedExecutables) {
            errorMessage << seperator << item;
            seperator = ", ";
        }
        errorMessage << "]";

        BUILDBOX_RUNNER_LOG(ERROR, errorMessage.str());
        writeErrorStatusFile(grpc::StatusCode::PERMISSION_DENIED,
                             errorMessage.str());
        exit(1);
    }

    ActionResult result;
    auto *result_metadata = result.mutable_execution_metadata();

    Runner::metadata_mark_input_download_start(result_metadata);
    const auto stagedDir = this->stageDirectory(inputRootDigest);
    const std::string stagedDirPath(stagedDir->getPath());

    // If `stage()` fails, it'll log the error and throw.
    Runner::metadata_mark_input_download_end(result_metadata);

    std::ostringstream workingDir;
    workingDir << stagedDirPath << "/" << command.working_directory();

    {
        ChangeDirectoryGuard g(workingDir.str());
        setUpEnvironment(command);
        createParentDirectories(command, workingDir.str());

        const std::vector<std::string> commandLine =
            generateCommandLine(command, stagedDirPath);

        BUILDBOX_RUNNER_LOG(DEBUG, "Executing " << join(commandLine, " "));
        executeAndStore(commandLine, &result);
    }

    if (!getSignalStatus()) {
        BUILDBOX_RUNNER_LOG(DEBUG, "Capturing command outputs from \""
                                       << workingDir.str() << "\"...");
        Runner::metadata_mark_output_upload_start(result_metadata);
        stagedDir->captureAllOutputs(command, &result);
        Runner::metadata_mark_output_upload_end(result_metadata);
        BUILDBOX_RUNNER_LOG(DEBUG, "Finished capturing command outputs");
    }

    return result;
}

bool HostToolsRunner::parseArg(std::string_view arg_view)
{
    if (arg_view.size() >= 2 && arg_view.substr(0, 2) == "--") {
        arg_view.remove_prefix(2);
        size_t assign_pos = arg_view.find('=');
        if (assign_pos != std::string_view::npos) {
            std::string_view key = arg_view.substr(0, assign_pos);
            std::string value(arg_view.substr(assign_pos + 1));
            if (key == "allow-executable") {
                // Eventually this should be an actual list, but for now we
                // just need to allow one
                this->d_allowedExecutables.insert(value);
                return true;
            }
            else if (key == "env-file") {
                if (this->d_environment.has_value()) {
                    std::cerr << "Only one of --clearenv and --env-file may "
                                 "be specified\n";
                    exit(1);
                }
                this->d_environment = parseEnvFile(value);
                return true;
            }
            else if (key == "action-digest-env") {
                this->d_actionDigestEnv = value;
                return true;
            }
        }
        else {
            if (arg_view == "prefix-staged-dir") {
                this->d_prefix = true;
                return true;
            }
            else if (arg_view == "clearenv") {
                if (this->d_environment.has_value()) {
                    std::cerr << "Only one of --clearenv and --env-file may "
                                 "be specified\n";
                    exit(1);
                }
                this->d_environment =
                    std::unordered_map<std::string, std::string>();
                return true;
            }
        }
    }
    return false;
}

void HostToolsRunner::printSpecialUsage()
{
    std::clog << "    --action-digest-env=VAR_NAME Set the action digest in "
                 "the specified environment variable.\n";
    std::clog << "    --allow-executable=ALLOWED_EXECUTABLE_PATH Only "
                 "commands which run this executable are allowed.\n";
    std::clog << "    --clearenv          Unset all environment variables.\n";
    std::clog << "    --env-file=PATH     Replace environment with the "
                 "contents of the specified file.\n";
}

bool HostToolsRunner::executableInAllowedExecutables(
    const std::string &executable) const
{
    return this->d_allowedExecutables.empty() ||
           this->d_allowedExecutables.find(executable) !=
               this->d_allowedExecutables.end();
}

} // namespace hosttools
} // namespace buildboxrun
} // namespace buildboxcommon
