/*
 * Copyright 2022 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 <filesystem>
#include <iterator>
#include <memory>
#include <ThreadPool.h>
#include <trexe_actionbuilder.h>
#include <trexe_actiondata.h>
#include <utility>
#include <vector>

#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_mergeutil.h>
#include <buildboxcommon_merklize.h>
#include <buildboxcommon_remoteexecutionclient.h>
#include <buildboxcommon_stringutils.h>

namespace trexe {

static std::pair<Digest, Command>
generateCommand(const std::vector<std::string> &argv,
                const std::map<std::string, std::string> &environment,
                const std::set<std::string> &outputPaths,
                const std::set<std::pair<std::string, std::string>> &platform,
                const std::string &workingDir,
                const std::set<std::string> &outputNodeProperties)
{
    BUILDBOX_LOG_DEBUG("Generating command.");
    Command command;
    for (const auto &arg : argv) {
        command.add_arguments(arg);
    }

    // environment variables
    for (const auto &var : environment) {
        auto envVar = command.add_environment_variables();
        envVar->set_name(var.first);
        envVar->set_value(var.second);
    }

    for (const auto &path : outputPaths) {
        command.add_output_paths(path);
    }

    // platform (deprecated but set to the same value as the Action for
    // redundancy)
    Platform *commandPlatform = commandMutablePlatformDeprecated(command);
    for (const auto &p : platform) {
        auto property = commandPlatform->add_properties();
        property->set_name(p.first);
        property->set_value(p.second);
    }

    command.set_working_directory(workingDir);

    // ouptut_node_properties
    for (const auto &property : outputNodeProperties) {
        command.add_output_node_properties(property);
    }

    auto digest = DigestGenerator::hash(command.SerializeAsString());

    BUILDBOX_LOG_DEBUG("Generated Command: Digest:" << toString(digest));

    return {std::move(digest), std::move(command)};
}

// Generates input tree and returns everything needed for
// FMB and uploads
// namely:
// - tree digest
// - digestToDiskPath (all needed file digests to path map)
// - digestToDirMessages (all needed dir proto digests to serialized proto map)
static std::tuple<Digest, digest_string_map, digest_string_map>
generateInputTree(std::shared_ptr<CASClient> &casClient,
                  const std::shared_ptr<Digest> &inputRootDigest,
                  const std::vector<InputPathOption> &inputPaths,
                  const bool followSymlinks, const size_t numDigestThreads)
{
    BUILDBOX_LOG_DEBUG("Generating input merkle tree.");

    std::unordered_map<buildboxcommon::Digest, std::string> digestToDiskPath;
    digest_string_map digestToSerializedProtos;

    // First handle shortcuts: input root without paths, no inputs at all.

    if (inputRootDigest && inputPaths.empty()) {
        return {*inputRootDigest, digestToDiskPath, digestToSerializedProtos};
    }

    if (!inputRootDigest && inputPaths.empty()) {
        Directory emptyDir;
        auto emptyDirMessage = emptyDir.SerializeAsString();
        auto emptyDigest = DigestGenerator::hash(emptyDirMessage);
        digestToSerializedProtos.emplace(emptyDigest,
                                         std::move(emptyDirMessage));

        return std::make_tuple(emptyDigest, digestToDiskPath,
                               digestToSerializedProtos);
    }

    std::vector<DirectoryTree> treesToMerge;

    if (inputRootDigest) {
        treesToMerge.push_back(casClient->getTree(*inputRootDigest));
    }

    // Create a thread pool to hash files
    std::unique_ptr<ThreadPool> threadPool;
    if (numDigestThreads > 0) {
        threadPool =
            std::make_unique<ThreadPool>(numDigestThreads, "trexe.digest");
    }

    for (const auto &input : inputPaths) {
        const auto &[localInput, remoteInput, permission, captureMtime,
                     properties, ignoreMatcher] = input;
        const auto localInputPath = std::filesystem::path(localInput);
        // If no remote mapping, the remote path is the filename
        // Implementation note: don't use move constructor of path
        // https://gitlab.com/BuildGrid/trexe/-/merge_requests/56
        std::filesystem::path remoteInputPath =
            remoteInput.has_value() ? std::filesystem::path(*remoteInput)
                                          .relative_path()
                                          .lexically_normal()
                                    : localInputPath.filename();
        // remove the trailing slash of a directory path
        if (!remoteInputPath.has_filename()) {
            remoteInputPath = remoteInputPath.parent_path();
        }

        std::vector<std::string> captureProperties;
        if (captureMtime) {
            captureProperties.emplace_back("mtime");
        }

        // c++20 allows for structured binding variables to be used in lambda
        // captures but this is currently a clang tidy fail. When this is no
        // longer an issue the unixModeUpdater lambda should be switched to
        // directly capture permission and the _permission variable can be
        // removed.
        const auto _permission = permission;

        // Create unix_mode updater based on permission options
        UnixModeUpdater unixModeUpdater;
        if (permission != InputPathOption::FilePermission::NONE) {
            captureProperties.emplace_back("unix_mode");
            unixModeUpdater = [&_permission](mode_t mode) {
                switch (_permission) {
                    case InputPathOption::FilePermission::READ_ONLY:
                        mode &= ~(S_IWUSR | S_IWGRP | S_IWOTH);
                        break;
                    case InputPathOption::FilePermission::READ_WRITE:
                        // don't add o+w for compatibility with buildbox-casd
                        mode |= (S_IWUSR | S_IWGRP);
                        break;
                    default:
                        break;
                }
                return mode;
            };
        }

        Merklizer merklizer(followSymlinks, captureProperties, ignoreMatcher,
                            threadPool.get());
        MerklizeResult merklizeResult;
        const auto inputStatus = std::filesystem::status(localInputPath);
        if (!std::filesystem::exists(inputStatus)) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error, "Input doesn't exist: " << localInputPath);
        }
        if (std::filesystem::is_directory(inputStatus)) {
            auto dirFd = FileDescriptor(open(
                localInputPath.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC));
            if (dirFd.get() < 0) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "Unable to open directory: " << localInputPath);
            }

            if (remoteInput.has_value() && remoteInputPath != ".") {
                MerklizeResult srcResult = merklizer.merklize(
                    dirFd.get(), localInputPath, Merklizer::hashFile,
                    unixModeUpdater, properties);
                merklizeResult = Merklizer::remapDirectory(
                    std::move(srcResult), remoteInputPath);
            }
            else {
                merklizeResult = merklizer.merklize(
                    dirFd.get(), localInputPath, Merklizer::hashFile,
                    unixModeUpdater, properties);
            }
        }
        else if (std::filesystem::is_regular_file(inputStatus)) {
            const auto file = File(localInputPath.c_str(), captureProperties,
                                   unixModeUpdater, properties);
            merklizeResult = Merklizer::remapFile(file, remoteInputPath);
            merklizeResult.d_digestToPath[file.d_digest] =
                localInputPath.c_str();
        }
        else if (std::filesystem::is_symlink(inputStatus)) {
            merklizeResult = Merklizer::remapSymlink(
                std::filesystem::read_symlink(localInputPath),
                remoteInputPath);
        }
        else {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "File type is not supported to be an input: "
                    << localInputPath);
        }
        // collect file paths
        digestToDiskPath.merge(std::move(merklizeResult.d_digestToPath));

        // collect tree
        DirectoryTree treeDirectory;
        {
            Tree rootTree = merklizeResult.tree();
            treeDirectory.root = DigestGenerator::hash(rootTree.root());
            treeDirectory.map.emplace(treeDirectory.root, rootTree.root());
            for (const auto &child : rootTree.children()) {
                const auto digest = DigestGenerator::hash(child);
                treeDirectory.map.emplace(digest, child);
            };
        }
        treesToMerge.emplace_back(std::move(treeDirectory));
    }

    Digest rootDigest;
    const bool result = MergeUtil::createMergedLayersDigest(
        treesToMerge, &rootDigest, &digestToSerializedProtos);
    if (!result) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Unable to merge input-path trees because of collision");
    }

    return {std::move(rootDigest), std::move(digestToDiskPath),
            std::move(digestToSerializedProtos)};
}

std::pair<Digest, Action>
generateAction(const Digest &commandDigest, const Digest &inputRootDigest,
               const int &execTimeout, const bool doNotCache,
               const std::string &salt,
               const std::set<std::pair<std::string, std::string>> &platform)
{
    BUILDBOX_LOG_DEBUG("Generating Action: ");

    Action action;
    action.mutable_command_digest()->CopyFrom(commandDigest);

    // input_root_digest
    action.mutable_input_root_digest()->CopyFrom(inputRootDigest);

    // timeout
    if (execTimeout != 0) { // optional, has server default if not specified
        google::protobuf::Duration actionTimeout;
        actionTimeout.set_seconds(execTimeout);
        action.mutable_timeout()->CopyFrom(actionTimeout);
    }

    // do_not_cache
    action.set_do_not_cache(doNotCache);

    // salt
    if (salt != "") {
        action.set_salt(salt);
    }

    // platform
    Platform *actionPlatform = action.mutable_platform();
    for (const auto &p : platform) {
        auto property = actionPlatform->add_properties();
        property->set_name(p.first);
        property->set_value(p.second);
    }

    auto digest = DigestGenerator::hash(action.SerializeAsString());

    return {std::move(digest), std::move(action)};
}

ActionData buildAction(
    std::shared_ptr<CASClient> casClient, const std::vector<std::string> &argv,
    const std::string &workingDir,
    const std::vector<InputPathOption> &inputPaths,
    const std::shared_ptr<Digest> &inputRootDigest,
    const std::set<std::string> &outputPaths,
    const std::set<std::pair<std::string, std::string>> &platform,
    const std::map<std::string, std::string> &environment,
    const int &execTimeout, const bool doNotCache, const bool followSymlinks,
    const std::string &salt, const std::set<std::string> &outputNodeProperties,
    const size_t numDigestThreads)
{
    auto [commandDigest, command] =
        generateCommand(argv, environment, outputPaths, platform, workingDir,
                        outputNodeProperties);

    auto [inputTreeDigest, inputDigestsToPaths,
          inputDigestsToSerializedProtos] =
        generateInputTree(casClient, inputRootDigest, inputPaths,
                          followSymlinks, numDigestThreads);

    BUILDBOX_LOG_DEBUG("Generated input root: "
                       << inputTreeDigest.hash() << "/"
                       << inputTreeDigest.size_bytes());

    // check that these digests are actually generated
    auto [actionDigest, action] =
        generateAction(commandDigest, inputTreeDigest, execTimeout, doNotCache,
                       salt, platform);

    BUILDBOX_LOG_DEBUG("Generated action: Digest: " << toString(actionDigest));

    return ActionData{
        .d_commandProto = std::move(command),
        .d_commandDigest = std::move(commandDigest),
        .d_actionProto = std::move(action),
        .d_actionDigest = std::move(actionDigest),
        .d_inputDigestsToPaths = std::move(inputDigestsToPaths),
        .d_inputDigestsToSerializedProtos =
            std::move(inputDigestsToSerializedProtos),
    };
}

} // namespace trexe
