// Copyright 2019 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_casclient.h>
#include <buildboxcommon_commandline.h>
#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_merklize.h>
#include <buildboxcommon_stringutils.h>
#include <buildboxcommon_version.h>

#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <grpcpp/create_channel.h>
#include <grpcpp/security/credentials.h>
#include <grpcpp/support/status.h>
#include <iostream>
#include <memory>
#include <processargs.h>
#include <span>
#include <sstream>
#include <stdexcept>
#include <sys/stat.h>
#include <ThreadPool.h>
#include <utility>

void uploadResources(
    int dirfd, const buildboxcommon::digest_string_map &blobs,
    const buildboxcommon::digest_string_map &digest_to_filepaths,
    const std::unique_ptr<buildboxcommon::CASClient> &casClient,
    ThreadPool *uploadThreadPool)
{
    std::vector<buildboxcommon::Digest> digestsToUpload;
    for (const auto &i : blobs) {
        digestsToUpload.push_back(i.first);
    }
    for (const auto &i : digest_to_filepaths) {
        digestsToUpload.push_back(i.first);
    }

    const auto missingDigests = casClient->findMissingBlobs(digestsToUpload);

    std::vector<buildboxcommon::CASClient::UploadRequest> upload_requests;
    upload_requests.reserve(missingDigests.size());
    for (const auto &digest : missingDigests) {
        // Finding the data in one of the source maps:
        if (blobs.count(digest)) {
            upload_requests.emplace_back(digest, blobs.at(digest));
        }
        else if (digest_to_filepaths.count(digest)) {
            upload_requests.emplace_back(
                buildboxcommon::CASClient::UploadRequest::from_path(
                    digest, dirfd, digest_to_filepaths.at(digest)));
        }
        else {
            throw std::runtime_error(
                "FindMissingBlobs returned non-existent digest");
        }
    }

    const auto failedToUpload =
        casClient->uploadBlobs(upload_requests, uploadThreadPool);
    if (!failedToUpload.empty()) {
        const auto &firstError = failedToUpload.front().status;
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Failed to upload "
                << failedToUpload.size()
                << " blobs. First error: code=" << firstError.error_code()
                << " message=" << firstError.error_message());
    }
}

void uploadDirectory(
    int dirfd, const std::string &path, const buildboxcommon::Digest &digest,
    const buildboxcommon::digest_string_map &digestsToBlobs,
    const buildboxcommon::digest_string_map &directoryDigestToFilepaths,
    const std::unique_ptr<buildboxcommon::CASClient> &casClient,
    ThreadPool *uploadThreadPool)
{
    assert(casClient != nullptr);

    try {
        BUILDBOX_LOG_DEBUG("Starting to upload merkle tree...");
        uploadResources(dirfd, digestsToBlobs, directoryDigestToFilepaths,
                        casClient, uploadThreadPool);
        BUILDBOX_LOG_INFO("Uploaded \"" << path << "\": " << digest.hash()
                                        << "/" << digest.size_bytes());
    }
    catch (const std::runtime_error &e) {
        BUILDBOX_LOG_ERROR("Uploading " << path
                                        << " failed with error: " << e.what());
        exit(1);
    }
}

std::pair<buildboxcommon::Digest, // rootDigest
          buildboxcommon::Digest  // treeDigest
          >
processDirectory(
    const std::string &path, const bool followSymlinks,
    const std::unique_ptr<buildboxcommon::CASClient> &casClient,
    const std::shared_ptr<std::vector<buildboxcommon::IgnorePattern>>
        &ignorePatterns,
    const std::vector<std::string> &captureProperties,
    const mode_t unixModeMask,
    const std::map<std::string, std::string> &nodeProperties,
    const bool allowChmodToRead, ThreadPool *uploadThreadPool,
    ThreadPool *digestThreadPool,
    const buildboxcommon::Command_OutputDirectoryFormat dirFormat,
    const casupload::LocalCasOption &localCasOption)
{
    // Note: LocalCAS::CaptureTree doesn't support following symlink yet
    if (localCasOption.d_useLocalCas && !followSymlinks &&
        casClient != nullptr) {
        const auto captureResponse =
            casClient
                ->captureTree({std::filesystem::absolute(path)},
                              captureProperties,
                              localCasOption.d_bypassLocalCache, unixModeMask,
                              nullptr, dirFormat, allowChmodToRead)
                .responses()
                .at(0);
        if (captureResponse.status().code() != grpc::StatusCode::OK) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "Failed to capture directory \""
                    << path << "\" code=" << captureResponse.status().code()
                    << " message=" << captureResponse.status().message());
        }
        const auto &rootDigest = captureResponse.root_directory_digest();
        const auto &treeDigest = captureResponse.tree_digest();

        std::ostringstream ss;
        ss << "Captured \"" << path << "\":";
        if (rootDigest.ByteSizeLong() != 0) {
            ss << " rootDirectoryDigest=" << rootDigest.hash() << "/"
               << rootDigest.ByteSizeLong();
        }
        if (treeDigest.ByteSizeLong() != 0) {
            ss << " treeDigest=" << treeDigest.hash() << "/"
               << treeDigest.ByteSizeLong();
        }
        BUILDBOX_LOG_INFO(ss.str())

        return std::make_pair(rootDigest, treeDigest);
    }

    const auto ignoreMatcher =
        ignorePatterns == nullptr
            ? nullptr
            : std::make_shared<buildboxcommon::IgnoreMatcher>(path,
                                                              ignorePatterns);

    buildboxcommon::Merklizer merklizer(followSymlinks, captureProperties,
                                        ignoreMatcher, digestThreadPool);
    buildboxcommon::FileDescriptor dirfd(
        open(path.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC));
    if (dirfd.get() < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Failed to open directory \"" << path << "\"");
    }
    auto merklizeResult = merklizer.merklize(
        dirfd.get(), "", buildboxcommon::Merklizer::hashFile,
        buildboxcommon::unixModeMaskUpdater(unixModeMask), nodeProperties);
    const auto rootDigest = merklizeResult.d_rootDigest;

    buildboxcommon::digest_string_map digestsToBlobs;
    if (dirFormat != buildboxcommon::Command_OutputDirectoryFormat::
                         Command_OutputDirectoryFormat_TREE_ONLY) {
        digestsToBlobs = std::move(merklizeResult.d_digestToDirectoryBlob);
    }

    buildboxcommon::Digest treeDigest;
    if (dirFormat != buildboxcommon::Command_OutputDirectoryFormat::
                         Command_OutputDirectoryFormat_DIRECTORY_ONLY) {
        auto tree = merklizeResult.tree().SerializeAsString();
        treeDigest = buildboxcommon::DigestGenerator::hash(tree);
        BUILDBOX_LOG_DEBUG("Computed tree digest for \""
                           << path << "\": " << treeDigest.hash() << "/"
                           << treeDigest.size_bytes());
        digestsToBlobs.emplace(treeDigest, std::move(tree));
    }

    BUILDBOX_LOG_DEBUG("Finished building nested directory from \""
                       << path << "\": " << rootDigest.hash() << "/"
                       << rootDigest.size_bytes());

    if (casClient == nullptr) {
        BUILDBOX_LOG_INFO("Computed directory digest for \""
                          << path << "\": " << rootDigest.hash() << "/"
                          << rootDigest.size_bytes());
    }
    else {
        uploadDirectory(dirfd.get(), path, rootDigest, digestsToBlobs,
                        merklizeResult.d_digestToPath, casClient,
                        uploadThreadPool);
    }

    return std::make_pair(rootDigest, std::move(treeDigest));
}

std::pair<buildboxcommon::Digest, // rootDigest
          buildboxcommon::Digest  // treeDigest
          >
processPath(const std::string &path, const bool followSymlinks,
            buildboxcommon::NestedDirectory *nestedDirectory,
            buildboxcommon::digest_string_map *digestToFilePaths,
            const std::unique_ptr<buildboxcommon::CASClient> &casClient,
            const std::shared_ptr<std::vector<buildboxcommon::IgnorePattern>>
                &ignorePatterns,
            const std::vector<std::string> &captureProperties,
            const mode_t unixModeMask,
            const std::map<std::string, std::string> &nodeProperties,
            const bool allowChmodToRead, ThreadPool *uploadThreadPool,
            ThreadPool *digestThreadPool,
            const buildboxcommon::Command_OutputDirectoryFormat dirFormat,
            const casupload::LocalCasOption &localCasOption)
{
    BUILDBOX_LOG_DEBUG("Starting to process \""
                       << path << "\", followSymlinks = " << std::boolalpha
                       << followSymlinks << std::noboolalpha);

    const std::filesystem::file_status stat =
        followSymlinks ? std::filesystem::status(path)
                       : std::filesystem::symlink_status(path);

    if (stat.type() == std::filesystem::file_type::directory) {
        return processDirectory(
            path, followSymlinks, casClient, ignorePatterns, captureProperties,
            unixModeMask, nodeProperties, allowChmodToRead, uploadThreadPool,
            digestThreadPool, dirFormat, localCasOption);
    }
    else if (stat.type() == std::filesystem::file_type::symlink) {
        const auto target = std::filesystem::read_symlink(path);
        nestedDirectory->addSymlink(target, path.c_str());
    }
    else if (stat.type() == std::filesystem::file_type::regular) {
        const buildboxcommon::File file(
            path.c_str(), captureProperties,
            buildboxcommon::unixModeMaskUpdater(unixModeMask), nodeProperties,
            allowChmodToRead);
        nestedDirectory->add(file, path.c_str());
        digestToFilePaths->emplace(file.d_digest, path);
    }
    else {
        BUILDBOX_LOG_DEBUG("Encountered unsupported file \""
                           << path << "\", skipping...");
    }
    return std::make_pair(buildboxcommon::Digest(), buildboxcommon::Digest());
}

int main(int argc, char *argv[])
{
    auto cliArgs = std::span(argv, argc);
    buildboxcommon::logging::Logger::getLoggerInstance().initialize(
        cliArgs[0]);

    auto args = casupload::processArgs(argc, argv);
    if (args.d_processed) {
        return 0;
    }
    if (!args.d_valid) {
        return 2;
    }

    BUILDBOX_LOG_SET_LEVEL(args.d_logLevel);

    buildboxcommon::DigestGenerator::init(args.d_digestFunctionValue);

    if (args.d_paths.size() > 1 && (!args.d_outputRootDigestFile.empty() ||
                                    !args.d_outputTreeDigestFile.empty())) {
        const auto &args_const_ref = args;
        // Making sure that we'll produce a single directory digest:
        const auto directoryPath = std::find_if(
            args.d_paths.cbegin(), args.d_paths.cend(),
            [&args_const_ref](const std::string &path) {
                return args_const_ref.d_followSymlinks
                           ? buildboxcommon::FileUtils::isDirectory(
                                 path.c_str())
                           : buildboxcommon::FileUtils::isDirectoryNoFollow(
                                 path.c_str());
            });

        if (directoryPath != args.d_paths.cend()) {
            BUILDBOX_LOG_ERROR(
                "`--output-digest-file` and `--output-tree-digest-file` can "
                "be used only when uploading a single directory structure");
            return 2;
        }
    }

    std::shared_ptr<std::vector<buildboxcommon::IgnorePattern>> ignorePatterns;
    if (args.d_ignoreFile != "") {
        std::ifstream ifs(args.d_ignoreFile);
        if (!ifs.good()) {
            BUILDBOX_LOG_ERROR(
                "Unable to read from file specified in --ignore-file");
            return 2;
        }
        ignorePatterns =
            buildboxcommon::IgnoreMatcher::parseIgnorePatterns(ifs);
    }

    // CAS client object (we don't initialize it if `dryRunMode` is set):
    std::unique_ptr<buildboxcommon::CASClient> casClient;

    try {
        if (!args.d_dryRunMode) {
            auto grpcClient = std::make_shared<buildboxcommon::GrpcClient>();
            grpcClient->init(args.d_casConnectionOptions);
            grpcClient->setToolDetails("casupload", buildboxcommon::VERSION);
            const std::string toolInvocationId =
                buildboxcommon::StringUtils::getUUIDString();
            const std::string correlatedInvocationsId =
                buildboxcommon::StringUtils::getUUIDString();
            grpcClient->setRequestMetadata("", toolInvocationId,
                                           correlatedInvocationsId);
            casClient =
                std::make_unique<buildboxcommon::CASClient>(grpcClient);
            casClient->init();
        }
    }
    catch (std::runtime_error &e) {
        std::cerr << "Failed to connect to CAS server: " << e.what() << '\n';
        return 1;
    }

    std::vector<std::string> captureProperties;
    if (args.d_captureMtime) {
        captureProperties.emplace_back("mtime");
    }
    if (args.d_captureUnixMode) {
        captureProperties.emplace_back("unix_mode");
    }

    buildboxcommon::NestedDirectory nestedDirectory;
    buildboxcommon::digest_string_map digestToFilePaths;

    const auto &localCasOption = args.d_localCasOption;

    // Thread pools
    // Disabled if local CAS is used
    std::unique_ptr<ThreadPool> uploadThreadPool = nullptr,
                                digestThreadPool = nullptr;
    if (!localCasOption.d_useLocalCas && args.d_numUploadThreads > 0) {
        uploadThreadPool = std::make_unique<ThreadPool>(
            args.d_numUploadThreads, "casupload.upload");
    }
    if (!localCasOption.d_useLocalCas && args.d_numDigestThreads > 0) {
        digestThreadPool = std::make_unique<ThreadPool>(
            args.d_numDigestThreads, "casupload.digest");
    }

    buildboxcommon::Digest dirRootDigest;
    buildboxcommon::Digest treeDigest;

    // Upload directories individually, and aggregate files to upload as single
    // merkle tree
    std::vector<buildboxcommon::Digest> directory_digests;
    for (const auto &path : args.d_paths) {
        std::tie(dirRootDigest, treeDigest) = processPath(
            path, args.d_followSymlinks, &nestedDirectory, &digestToFilePaths,
            casClient, ignorePatterns, captureProperties, args.d_fileUmask,
            args.d_nodeProperties, args.d_allowChmodToRead,
            uploadThreadPool.get(), digestThreadPool.get(),
            args.d_directoryFormat, localCasOption);
    }

    // Upload single files
    if (!digestToFilePaths.empty()) {
        BUILDBOX_LOG_DEBUG("Building nested directory structure...");
        buildboxcommon::digest_string_map blobs;
        dirRootDigest = nestedDirectory.to_digest(&blobs);

        // Compute the tree digest and add it to the blobs to be uploaded
        auto tree = nestedDirectory.to_tree().SerializeAsString();
        treeDigest = buildboxcommon::DigestGenerator::hash(tree);
        blobs.emplace(treeDigest, std::move(tree));

        BUILDBOX_LOG_INFO("Computed directory digest: "
                          << dirRootDigest.hash() << "/"
                          << dirRootDigest.size_bytes());

        BUILDBOX_LOG_INFO("Computed tree digest: " << treeDigest.hash() << "/"
                                                   << treeDigest.size_bytes());

        if (!args.d_dryRunMode) {
            try {
                uploadResources(AT_FDCWD, blobs, digestToFilePaths, casClient,
                                uploadThreadPool.get());
                BUILDBOX_LOG_DEBUG("Files uploaded successfully");
            }
            catch (const std::runtime_error &e) {
                BUILDBOX_LOG_ERROR(
                    "Uploading files failed with error: " << e.what());
                return 1;
            }
        }
    }

    if (!args.d_outputTreeDigestFile.empty() &&
        (args.d_directoryFormat == buildboxcommon::Command::TREE_ONLY ||
         args.d_directoryFormat ==
             buildboxcommon::Command::TREE_AND_DIRECTORY)) {

        // Write the tree digest to the outputTreeDigestFile
        BUILDBOX_LOG_DEBUG("Writing tree digest ["
                           << treeDigest << "] to '"
                           << args.d_outputTreeDigestFile << "'");
        buildboxcommon::FileUtils::writeDigestToFile(
            treeDigest, args.d_outputTreeDigestFile);
    }

    if (!args.d_outputRootDigestFile.empty() &&
        (args.d_directoryFormat == buildboxcommon::Command::DIRECTORY_ONLY ||
         args.d_directoryFormat ==
             buildboxcommon::Command::TREE_AND_DIRECTORY)) {

        // Write the root digest to the outputRootDigestFile
        BUILDBOX_LOG_DEBUG("Writing root digest ["
                           << dirRootDigest << "] to '"
                           << args.d_outputRootDigestFile << "'");
        buildboxcommon::FileUtils::writeDigestToFile(
            dirRootDigest, args.d_outputRootDigestFile);
    }

    return 0;
}
