// 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_executionclient.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_permissions.h>
#include <buildboxcommon_stringutils.h>
#include <ThreadPool.h>

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

#include <optional>
#include <set>
#include <sys/stat.h>
#include <utility>

namespace buildboxcommon {

void ExecutionClient::init(
    std::shared_ptr<ActionCache::StubInterface> actionCacheStub)
{
    d_actionCacheStub = std::move(actionCacheStub);
}

void ExecutionClient::init()
{
    d_actionCacheStub =
        d_actionCacheGrpcClient
            ? ActionCache::NewStub(d_actionCacheGrpcClient->channel())
            : nullptr;
}

void ExecutionClient::disableActionCacheUpdates()
{
    d_actionCacheUpdatesAllowed = false;
}

bool ExecutionClient::fetchFromActionCache(
    const Digest &actionDigest, const std::set<std::string> &outputs,
    ActionResult *result)
{
    if (!d_actionCacheStub) {
        throw std::runtime_error("ActionCache Stub not Configured");
    }
    GetActionResultRequest actionRequest;
    actionRequest.set_instance_name(d_actionCacheGrpcClient->instanceName());

    actionRequest.set_inline_stdout(true);
    actionRequest.set_inline_stderr(true);
    for (const auto &o : outputs) {
        actionRequest.add_inline_output_files(o);
    }

    *actionRequest.mutable_action_digest() = actionDigest;

    ActionResult actionResult;
    bool found = false;

    auto getactionresult_lambda = [&](grpc::ClientContext &context) {
        grpc::Status status = d_actionCacheStub->GetActionResult(
            &context, actionRequest, &actionResult);

        if (status.ok()) {
            found = true;
        }
        else if (status.error_code() == grpc::StatusCode::NOT_FOUND) {
            /* Return OK as this shouldn't throw an exception */
            return grpc::Status::OK;
        }

        return status;
    };

    d_actionCacheGrpcClient->issueRequest(
        getactionresult_lambda, "ActionCache.GetActionResult()", nullptr);

    if (found && result != nullptr) {
        *result = actionResult;
    }

    return found;
}

void ExecutionClient::updateActionCache(const Digest &actionDigest,
                                        const ActionResult &result)
{
    if (!d_actionCacheStub) {
        throw std::runtime_error("ActionCache Stub not Configured");
    }

    if (!d_actionCacheUpdatesAllowed) {
        throw std::runtime_error("ActionCache updates not allowed");
    }

    UpdateActionResultRequest actionRequest;
    actionRequest.set_instance_name(d_actionCacheGrpcClient->instanceName());

    actionRequest.mutable_action_digest()->CopyFrom(actionDigest);
    actionRequest.mutable_action_result()->CopyFrom(result);

    auto updateactionresult_lambda = [&](grpc::ClientContext &context) {
        ActionResult actionResult;
        return d_actionCacheStub->UpdateActionResult(&context, actionRequest,
                                                     &actionResult);
    };

    d_actionCacheGrpcClient->issueRequest(updateactionresult_lambda,
                                          "ActionCache.UpdateActionResult()",
                                          nullptr);
}

void checkDownloadBlobsResult(const CASClient::DownloadBlobsResult &results)
{
    std::vector<std::string> missingBlobs;

    for (const auto &result : results) {
        const auto &status = result.second.first;
        if (status.code() == grpc::StatusCode::NOT_FOUND) {
            missingBlobs.push_back(result.first);
        }
        else if (status.code() != grpc::StatusCode::OK) {
            throw GrpcError(
                "Failed to download output blob " + result.first + ": " +
                    std::to_string(status.code()) + ": " + status.message(),
                grpc::Status(static_cast<grpc::StatusCode>(status.code()),
                             status.message()));
        }
    }

    if (!missingBlobs.empty()) {
        std::ostringstream error;
        error << missingBlobs.size()
              << " output blobs missing from ActionResult: ";
        bool first = true;
        for (const auto &hash : missingBlobs) {
            if (!first) {
                error << ", ";
            }
            error << hash;
            first = false;
        }
        throw GrpcError(error.str(), grpc::Status(grpc::StatusCode::NOT_FOUND,
                                                  error.str()));
    }
}

Digest addDirectoryToMap(std::unordered_map<Digest, Directory> *map,
                         const Directory &directory)
{
    auto digest = DigestGenerator::hash(directory);
    (*map)[digest] = directory;
    return digest;
}

// Helper function to recursively visit a directory tree to populate file
// digests and detect duplicate ones.
// This function should be only called for root directories.
void populateFileDigestsFromDirectory(
    const Directory &directory,
    const std::unordered_map<Digest, Directory> &digestDirectoryMap,
    std::unordered_set<Digest> *fileDigests,
    std::unordered_set<Digest> *duplicateFileDigests)
{
    for (const auto &fileNode : directory.files()) {
        const bool inserted = fileDigests->insert(fileNode.digest()).second;
        if (!inserted) {
            duplicateFileDigests->insert(fileNode.digest());
        }
    }
    for (const auto &subdirNode : directory.directories()) {
        populateFileDigestsFromDirectory(
            digestDirectoryMap.at(subdirNode.digest()), digestDirectoryMap,
            fileDigests, duplicateFileDigests);
    }
}

void createParentDirectory(int dirfd, const std::string &path)
{
    const auto pos = path.rfind('/');
    if (pos != std::string::npos) {
        const std::string parent_path = path.substr(0, pos);
        FileUtils::createDirectory(dirfd, parent_path.c_str());
    }
}
void stageSymlink(int dirfd, const std::string &path,
                  const std::string &target)
{

    // Remove the existing entity if it exists
    struct stat st = {};
    if (fstatat(dirfd, path.c_str(), &st, AT_SYMLINK_NOFOLLOW) == 0) {
        if (S_ISDIR(st.st_mode)) {
            FileUtils::deleteDirectory(dirfd, path.c_str());
        }
        else {
            if (unlinkat(dirfd, path.c_str(), 0) < 0) {
                BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                    std::system_error, errno, std::system_category,
                    "Failed to remove existing file at: " + path);
            }
        }
    }
    else {
        if (errno != ENOENT) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Failed to stat entity at: " + path);
        }
    }

    if (symlinkat(target.c_str(), dirfd, path.c_str()) < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(std::system_error, errno,
                                              std::system_category,
                                              "Failed to create symlink");
    }
}

void stageDownloadedFile(
    int dirfd, const std::string &path, const Digest &digest,
    bool isExecutable, const NodeProperties &nodeProperties, int temp_dirfd,
    const CASClient::DownloadBlobsResult &downloaded_files,
    const std::unordered_set<Digest> &duplicate_file_digests)
{
    auto temp_path = downloaded_files.at(digest.hash()).second;

    if (duplicate_file_digests.find(digest) != duplicate_file_digests.end()) {
        // Digest is used by multiple output files, create a copy
        auto temp_copy_path = temp_path + StringUtils::getRandomHexString();
        FileUtils::copyFile(temp_dirfd, temp_path.c_str(), temp_dirfd,
                            temp_copy_path.c_str());
        temp_path = temp_copy_path;
    }

    mode_t mode = nodeProperties.has_unix_mode()
                      ? static_cast<mode_t>(nodeProperties.unix_mode().value())
                      : PERMISSION_RWUSR_RALL;
    if (isExecutable) {
        mode |= S_IXUSR | S_IXGRP | S_IXOTH;
    }
    if (fchmodat(temp_dirfd, temp_path.c_str(), mode, 0) < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Failed to set file mode of downloaded file");
    }
    if (nodeProperties.has_mtime()) {
        const auto mtimePoint =
            buildboxcommon::TimeUtils::parse_timestamp(nodeProperties.mtime());
        buildboxcommon::FileUtils::setFileTimes(temp_dirfd, temp_path.c_str(),
                                                std::nullopt, mtimePoint);
    }
    if (renameat(temp_dirfd, temp_path.c_str(), dirfd, path.c_str()) < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Failed to move downloaded file to final location: " << path);
    }
}

void stageDownloadedDirectory(
    int dirfd, const std::string &path, const Digest &dir_digest,
    int temp_dirfd,
    const std::unordered_map<Digest, Directory> &digest_directory_map,
    const CASClient::DownloadBlobsResult &downloaded_files,
    const std::unordered_set<Digest> &duplicate_file_digests)
{
    const auto &directory = digest_directory_map.at(dir_digest);
    FileUtils::createDirectory(dirfd, path.c_str());

    FileDescriptor current_dirfd(
        openat(dirfd, path.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC));
    if (current_dirfd.get() < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Failed to open newly created subdirectory");
    }

    for (const auto &file_node : directory.files()) {
        stageDownloadedFile(current_dirfd.get(), file_node.name(),
                            file_node.digest(), file_node.is_executable(),
                            file_node.node_properties(), temp_dirfd,
                            downloaded_files, duplicate_file_digests);
    }
    for (const auto &dir_node : directory.directories()) {
        // Recursive call for subdirectory
        stageDownloadedDirectory(current_dirfd.get(), dir_node.name(),
                                 dir_node.digest(), temp_dirfd,
                                 digest_directory_map, downloaded_files,
                                 duplicate_file_digests);
    }
    for (const auto &symlink_node : directory.symlinks()) {
        stageSymlink(current_dirfd.get(), symlink_node.name(),
                     symlink_node.target());
    }

    const auto &nodeProperties = directory.node_properties();
    if (nodeProperties.has_unix_mode()) {
        const mode_t mode =
            static_cast<mode_t>(nodeProperties.unix_mode().value());
        if (fchmodat(dirfd, path.c_str(), mode, 0) < 0) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Failed to set file mode of downloaded file");
        }
    }
    if (nodeProperties.has_mtime()) {
        const auto mtimePoint =
            buildboxcommon::TimeUtils::parse_timestamp(nodeProperties.mtime());
        buildboxcommon::FileUtils::setFileTimes(dirfd, path.c_str(),
                                                std::nullopt, mtimePoint);
    }
}

void ExecutionClient::downloadOutputs(buildboxcommon::CASClient *casClient,
                                      const ActionResult &actionResult,
                                      int dirfd, ThreadPool *threadPool)
{
    std::unordered_set<Digest> tree_digests;
    for (const auto &dir : actionResult.output_directories()) {
        tree_digests.insert(dir.tree_digest());
    }

    const auto downloaded_trees = casClient->downloadBlobs(
        std::vector<Digest>(tree_digests.cbegin(), tree_digests.cend()),
        threadPool);
    checkDownloadBlobsResult(downloaded_trees);

    std::unordered_set<Digest> file_digests, duplicate_file_digests;
    std::unordered_map<Digest, Directory> digest_directory_map;
    // Map from Digest of Tree to Digest of root directory of that tree
    std::unordered_map<Digest, Digest> tree_digest_root_digest_map;

    for (const auto &file : actionResult.output_files()) {
        const bool inserted = file_digests.insert(file.digest()).second;
        if (!inserted) {
            duplicate_file_digests.insert(file.digest());
        }
    }
    for (const auto &dir : actionResult.output_directories()) {
        Tree tree;
        const auto serialized_tree =
            downloaded_trees.at(dir.tree_digest().hash()).second;
        if (!tree.ParseFromString(serialized_tree)) {
            throw std::runtime_error("Could not deserialize downloaded Tree");
        }

        const auto root_digest =
            addDirectoryToMap(&digest_directory_map, tree.root());
        tree_digest_root_digest_map[dir.tree_digest()] = root_digest;

        for (const auto &tree_child : tree.children()) {
            addDirectoryToMap(&digest_directory_map, tree_child);
        }
    }

    // Visit each tree recursively to detect duplications
    for (const auto &treeRootIter : tree_digest_root_digest_map) {
        const Digest &dirDigest = treeRootIter.second;
        populateFileDigestsFromDirectory(digest_directory_map.at(dirDigest),
                                         digest_directory_map, &file_digests,
                                         &duplicate_file_digests);
    }

    const auto tempDirectoryName =
        ".reclient-" + StringUtils::getRandomHexString();

    if (mkdirat(dirfd, tempDirectoryName.c_str(), PERMISSION_RWXUSR) < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Failed to create temporary directory");
    }
    try {
        FileDescriptor temp_dirfd(openat(dirfd, tempDirectoryName.c_str(),
                                         O_RDONLY | O_DIRECTORY | O_CLOEXEC));
        if (temp_dirfd.get() == -1) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Failed to open temporary directory");
        }

        const auto downloaded_files = casClient->downloadBlobsToDirectory(
            std::vector<Digest>(file_digests.cbegin(), file_digests.cend()),
            temp_dirfd.get(), threadPool);
        checkDownloadBlobsResult(downloaded_files);

        for (const auto &file : actionResult.output_files()) {
            createParentDirectory(dirfd, file.path());
            stageDownloadedFile(dirfd, file.path(), file.digest(),
                                file.is_executable(), file.node_properties(),
                                temp_dirfd.get(), downloaded_files,
                                duplicate_file_digests);
        }
        for (const auto &symlink : actionResult.output_symlinks()) {
            createParentDirectory(dirfd, symlink.path());
            stageSymlink(dirfd, symlink.path(), symlink.target());
        }
        for (const auto &dir : actionResult.output_directories()) {
            const auto dir_digest =
                tree_digest_root_digest_map[dir.tree_digest()];
            createParentDirectory(dirfd, dir.path());
            stageDownloadedDirectory(dirfd, dir.path(), dir_digest,
                                     temp_dirfd.get(), digest_directory_map,
                                     downloaded_files, duplicate_file_digests);
        }
    }
    catch (...) {
        FileUtils::deleteDirectory(dirfd, tempDirectoryName.c_str());
        throw;
    }
    FileUtils::deleteDirectory(dirfd, tempDirectoryName.c_str());
}

ActionResult ExecutionClient::getActionResult(const Operation &operation)
{
    if (!operation.done()) {
        throw std::logic_error(
            "Called getActionResult on an unfinished Operation");
    }
    else if (operation.has_error()) {
        throw GrpcError(
            "Operation failed: " + std::to_string(operation.error().code()) +
                ": " + operation.error().message(),
            grpc::Status(
                static_cast<grpc::StatusCode>(operation.error().code()),
                operation.error().message()));
    }
    else if (!operation.response().Is<ExecuteResponse>()) {
        throw std::runtime_error("Server returned invalid Operation result");
    }

    ExecuteResponse executeResponse;
    if (!operation.response().UnpackTo(&executeResponse)) {
        throw std::runtime_error("Operation response unpacking failed");
    }

    const auto executeStatus = executeResponse.status();
    if (executeStatus.code() != google::rpc::Code::OK) {
        throw GrpcError(
            "Execution failed: " + executeStatus.message(),
            grpc::Status(static_cast<grpc::StatusCode>(executeStatus.code()),
                         executeStatus.message()));
    }

    ActionResult actionResult = executeResponse.result();
    if (actionResult.exit_code() == 0) {
        BUILDBOX_LOG_DEBUG("Execute response message: " +
                           executeResponse.message());
    }
    else if (!executeResponse.message().empty()) {
        BUILDBOX_LOG_INFO("Remote execution message: " +
                          executeResponse.message());
    }

    return actionResult;
}

std::shared_ptr<buildboxcommon::GrpcClient>
ExecutionClient::getActionCacheGrpcClient() const
{
    return d_actionCacheGrpcClient;
}

std::shared_ptr<ActionCache::StubInterface>
ExecutionClient::getActionCacheStub() const
{
    return d_actionCacheStub;
}

bool ExecutionClient::isActionCacheUpdatesAllowed() const
{
    return d_actionCacheUpdatesAllowed;
}

} // namespace buildboxcommon
