/*
 * Copyright 2021 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 <buildboxrun_container_host.h>

#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_mergeutil.h>
#include <buildboxcommon_runner.h>
#include <buildboxcommon_systemutils.h>

#include <array>
#include <exception>
#include <filesystem>
#include <functional>
#include <iostream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

namespace buildboxcommon {
namespace buildboxrun {
namespace oci {

ActionResult OCIRunner::execute(const Command &command,
                                const Digest &inputRootDigest,
                                const Platform &platform)
{
    const Digest bundleDigest = getBundleDigest(platform);

    DirectoryTree inputRootTree;
    DirectoryTree bundleTree;

    try {
        d_casClient->fetchTree(inputRootDigest, true);
        d_casClient->fetchTree(bundleDigest, true);

        inputRootTree = d_casClient->getTree(inputRootDigest);
        bundleTree = d_casClient->getTree(bundleDigest);

        const Digest mergedDigest =
            mergeFilesystemDigest(inputRootTree, bundleTree);
        const Digest rootDigest = prepareNestedDigest(mergedDigest);
        return ContainerHost(this, command, rootDigest, d_casClient,
                             d_stagingMode)
            .execute();
    }
    catch (const std::exception &e) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Errors when fetching and getting directory tree from CAS: "
                << e.what());
    }
}

bool OCIRunner::parseArg(std::string_view arg)
{
    const size_t delimiter_idx = arg.find('=');
    if (delimiter_idx == std::string::npos || delimiter_idx == arg.size() - 1)
        return false;

    const std::string value(arg.substr(delimiter_idx + 1));
    const std::string_view key = arg.substr(0, delimiter_idx);
    try {
        if (key == "--memory-limit") {
            d_resourceLimits.emplace("RLIMIT_AS", std::stoull(value));
        }
        else if (key == "--cpu-time") {
            d_resourceLimits.emplace("RLIMIT_CPU", std::stoull(value));
        }
        else {
            return false;
        }

        return true;
    }
    catch (...) {
        return false;
    }
}

Digest OCIRunner::mergeFilesystemDigest(const DirectoryTree &inputRootTree,
                                        const DirectoryTree &bundleTree)
{
    // We start merging.
    std::unordered_map<Digest, std::string> blobs;
    std::vector<Digest> newDigestList;

    Digest mergedDigest;
    if (!MergeUtil::createMergedDigest(inputRootTree, bundleTree,
                                       &mergedDigest, &blobs, &newDigestList))
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Error merging input digest with bundle filesystem digest.");

    std::vector<CASClient::UploadRequest> uploadRequests;
    for (const Digest &missingDigest :
         d_casClient->findMissingBlobs(newDigestList))
        uploadRequests.emplace_back(missingDigest, blobs.at(missingDigest));

    const std::vector<CASClient::UploadResult> uploadResults =
        d_casClient->uploadBlobs(uploadRequests);

    std::ostringstream err;
    if (std::count_if(uploadResults.cbegin(), uploadResults.cend(),
                      [&err](const CASClient::UploadResult &result) {
                          if (!result.status.ok())
                              err << "Failed to upload a merged digest("
                                  << result.digest << "), status = ["
                                  << result.status.error_code() << ": \""
                                  << result.status.error_message() << "\"]\n";

                          return !result.status.ok();
                      }))
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error, err.str());

    return mergedDigest;
}

Digest OCIRunner::prepareNestedDigest(const Digest &fsDigest)
{
    Directory root;
    DirectoryNode *rootfsNode = root.add_directories();
    rootfsNode->set_name("rootfs");
    *rootfsNode->mutable_digest() = fsDigest;

    const std::string rootBlob = root.SerializeAsString();
    Digest rootDigest = DigestGenerator::hash(rootBlob);

    const std::vector<CASClient::UploadResult> uploadResults =
        d_casClient->uploadBlobs({{rootDigest, rootBlob}});
    if (!uploadResults.empty()) {
        const CASClient::UploadResult &uploadResult = uploadResults.front();
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Failed to upload nested digest("
                << uploadResult.digest << "), status = ["
                << uploadResult.status.error_code() << ": \""
                << uploadResult.status.error_message() << "\"]\n");
    }

    return rootDigest;
}

Digest OCIRunner::getBundleDigest(const Platform &platform)
{
    static const std::string assetUrlPropName = "chrootAssetUrl";
    static const std::string assetQualifierPropName = "chrootAssetQualifier";
    static const std::string digestPropName = "chrootRootDigest";
    static const std::string containerImageIDPropName = "container-image";

    std::vector<std::string> urls;
    std::vector<std::pair<std::string, std::string>> qualifiers;

    for (const Platform_Property &prop : platform.properties()) {
        if (d_assetClient) {
            if (prop.name() == assetUrlPropName) {
                urls.push_back(prop.value());
            }
            else if (prop.name() == assetQualifierPropName) {
                const std::size_t delimiterPos =
                    prop.value().find_first_of('=');
                if (delimiterPos == std::string::npos) {
                    BUILDBOXCOMMON_THROW_EXCEPTION(
                        std::runtime_error,
                        "Invalid " << assetQualifierPropName
                                   << " property: " << prop.value()
                                   << "; No equal sign found.");
                }

                qualifiers.emplace_back(prop.value().substr(0, delimiterPos),
                                        prop.value().substr(delimiterPos + 1));
            }
        }

        if (prop.name() == containerImageIDPropName) {
            std::string imageID = prop.value();

            const std::string sha256Str = "@sha256:";
            if (imageID.find(sha256Str) == std::string::npos) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error, "Invalid container image ID "
                                            << imageID
                                            << "; image ID must be in the "
                                               "form of a sha256 digest");
            }

            // Removing the "docker://" prefix if it exists
            const std::string dockerPrefixStr = "docker://";
            if (imageID.substr(0, dockerPrefixStr.length()) ==
                dockerPrefixStr) {
                imageID = imageID.substr(dockerPrefixStr.length());
            }

            return fetchAndUploadImage(imageID);
        }
    }

    if (d_assetClient && !urls.empty()) {
        return d_assetClient->fetchDirectory(urls, qualifiers).digest;
    }

    // For whatever reason, the asset client path didn't work, and no container
    // image was specified. We switch back to the old digest info path.
    const auto containerDigestProp = std::find_if(
        std::cbegin(platform.properties()), std::cend(platform.properties()),
        [&](const Platform_Property &property) {
            return property.name() == digestPropName;
        });

    if (containerDigestProp == std::cend(platform.properties()))
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Cannot find metadata for the OCI Spec file.");

    // The property value should be formatted as <spec hash>/<spec size>
    const std::string &propValue = containerDigestProp->value();
    const std::size_t delimiterPos = propValue.find('/');
    if (delimiterPos == std::string::npos)
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Property "
                << digestPropName
                << " doesn't contain a valid value; missing delimiter '/'.");

    Digest digest;
    digest.set_hash(propValue.substr(0, delimiterPos));
    digest.set_size_bytes(std::stoll(propValue.substr(delimiterPos + 1)));

    return digest;
}

void OCIRunner::fetchDockerImage(const std::string &imageName) const
{
    std::string dockerPath = SystemUtils::getPathToCommand("docker");
    if (dockerPath.empty()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Unable to pull image"
                                           << imageName
                                           << "; failed to locate docker.");
    }
    std::vector<std::string> cmd = {dockerPath, "pull", "-q", imageName};
    if (SystemUtils::executeCommandAndWait(cmd)) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Unable to pull docker image "
                                    << imageName << " from the registry.");
    }
}

void OCIRunner::extractFileSystemFromImage(const std::string &imageName,
                                           const std::string &filePath) const
{
    // Checking if filePath is a valid file location
    if (filePath.empty() || !exists(std::filesystem::path(filePath))) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Specified file path to extract tar to is invalid: " << filePath);
    }

    // Getting path to Docker installation
    std::string dockerLocation = SystemUtils::getPathToCommand("docker");
    if (dockerLocation.empty()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Unable to create container for image "
                                           << imageName
                                           << "; failed to locate docker.");
    }

    // Getting path to tar installation
    std::string tarLocation = SystemUtils::getPathToCommand("tar");
    if (tarLocation.empty()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Unable to extract tar file"
                                    << imageName << "; failed to locate tar.");
    }

    // Creating a container with the image
    std::vector<std::string> createCmd = {dockerLocation, "create", imageName};
    auto createRes = SystemUtils::executeCommandWithResult(createCmd, true);
    if (createRes.d_exitCode) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Unable to create container for image: " << imageName);
    }

    // Creating a temporary directory to export the tar to
    TemporaryDirectory tmpDir;
    std::string tarPath = tmpDir.strname() + "/container-image.tar";
    std::string outputTar = "--output=" + tarPath;

    // Exporting the image's file system into a tarball in the temporary
    // directory
    std::vector<std::string> exportCmd = {
        dockerLocation, "export", outputTar,
        createRes.d_stdOut.substr(0, createRes.d_stdOut.length() - 2)};
    auto resExport = SystemUtils::executeCommandWithResult(exportCmd);
    if (resExport.d_exitCode) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Unable to export container for image: " << imageName);
    }

    // Extracting the tar file into the specified filePath
    std::vector<std::string> tarCommand = {tarLocation, "-xf", tarPath, "-C",
                                           filePath};
    auto res = SystemUtils::executeCommandWithResult(tarCommand);
    if (res.d_exitCode) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Unable to extract tar to directory "
                                    << filePath << " for image: " << imageName
                                    << "; tar extract command failed");
    }
}

const Digest OCIRunner::uploadImageFileSystem(const std::string &path) const
{
    Digest root_directory_digest;
    Tree tree;

    const auto failed_uploads =
        d_casClient->uploadDirectory(path, &root_directory_digest, &tree);

    if (!failed_uploads.empty()) {
        const auto &status = failed_uploads[0].status;
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Failed to upload "
                                    << failed_uploads.size()
                                    << " blobs to CAS:" << status.error_code()
                                    << ": " << status.error_message());
    }

    return root_directory_digest;
}

const Digest OCIRunner::fetchAndUploadImage(const std::string &imageName) const
{
    fetchDockerImage(imageName);
    TemporaryDirectory tmpDir;
    extractFileSystemFromImage(imageName, tmpDir.name());
    return uploadImageFileSystem(tmpDir.name());
}

} // namespace oci
} // namespace buildboxrun
} // namespace buildboxcommon
