// 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_mergeutil.h>

#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_exception.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_merklize.h>

#include <algorithm>
#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

namespace buildboxcommon {

namespace {

class NodeMetaData {
  private:
    const std::string d_path;

  protected:
    explicit NodeMetaData(std::string path) : d_path(std::move(path)) {}
    ~NodeMetaData() {}
    // Delete copy constructor
    NodeMetaData(const NodeMetaData &) = delete;

    // Delete copy assignment operator
    NodeMetaData &operator=(const NodeMetaData &) = delete;

    // Delete move constructor
    NodeMetaData(NodeMetaData &&) = delete;

    // Delete move assignment operator
    NodeMetaData &operator=(NodeMetaData &&) = delete;

  public:
    const std::string &path() const { return d_path; }
    virtual const Digest &digest() const = 0;
    virtual void addToNestedDirectory(NestedDirectory *nd) const = 0;
    virtual bool isExecutable() const { return false; }
    virtual NodeProperties nodeProperties() const { return {}; }
    virtual void print(std::ostream &out) const = 0;
};

bool operator==(const NodeMetaData &obj, const FileNode &fileNode)
{
    return (obj.digest() == fileNode.digest() &&
            obj.isExecutable() == fileNode.is_executable() &&
            google::protobuf::util::MessageDifferencer::Equals(
                obj.nodeProperties(), fileNode.node_properties()));
}

bool operator!=(const NodeMetaData &obj, const FileNode &fileNode)
{
    return !(obj == fileNode);
}

class FileNodeMetaData : public NodeMetaData {
  private:
    const File d_file;

  public:
    virtual ~FileNodeMetaData() = default;
    FileNodeMetaData(const std::string &path, const Digest &digest,
                     const bool is_executable,
                     const NodeProperties &nodeProperties)
        : NodeMetaData(path), d_file(digest, is_executable, nodeProperties)
    {
    }

    const Digest &digest() const override { return d_file.d_digest; }

    void addToNestedDirectory(NestedDirectory *nd) const override
    {
        nd->add(d_file, path().c_str());
    }

    bool isExecutable() const override { return d_file.d_executable; }

    NodeProperties nodeProperties() const override
    {
        return d_file.d_nodeProperties;
    }

    void print(std::ostream &out) const override
    {
        out << "file:    " << path() << " [" << digest()
            << ", executable = " << std::boolalpha << isExecutable() << "]\n";
    }
};

class SymlinkNodeMetaData : public NodeMetaData {
  private:
    const std::string d_symlinkTarget;

  public:
    virtual ~SymlinkNodeMetaData() = default;
    SymlinkNodeMetaData(const std::string &symlinkName,
                        std::string symlinkTarget)
        : NodeMetaData(symlinkName), d_symlinkTarget(std::move(symlinkTarget))
    {
    }

    const std::string &target() const { return d_symlinkTarget; }

    // symlinks have no digest, so return an empty one
    // see
    // https://gitlab.com/BuildGrid/buildbox/buildbox/blob/master/protos/build/bazel/remote/execution/v2/remote_execution.proto#L658
    const Digest &digest() const override
    {
        static Digest d;
        return d;
    }

    void addToNestedDirectory(NestedDirectory *nd) const override
    {
        nd->addSymlink(d_symlinkTarget, path().c_str());
    }

    void print(std::ostream &out) const override
    {
        out << "symlink: " << path() << ", " << d_symlinkTarget << "\n";
    }
};

class DirNodeMetaData : public NodeMetaData {
  private:
    const Digest d_digest;

  public:
    virtual ~DirNodeMetaData() = default;
    DirNodeMetaData(const std::string &path, Digest digest)
        : NodeMetaData(path), d_digest(std::move(digest))
    {
    }

    const Digest &digest() const override { return d_digest; }

    void addToNestedDirectory(NestedDirectory *nd) const override
    {
        nd->addDirectory(path().c_str());
    }

    void print(std::ostream &out) const override
    {
        out << "dir:     " << path() << " [" << digest() << "]\n";
    }
};

typedef std::unordered_map<std::string, std::shared_ptr<NodeMetaData>>
    PathNodeMetaDataMap;

std::ostream &operator<<(std::ostream &out, const NodeMetaData &obj)
{
    obj.print(out);
    return out;
}

inline std::string genNewPath(const std::string &dirName,
                              const std::string &nodeName)
{
    std::string result = dirName.empty() ? nodeName : dirName + "/" + nodeName;
    return result;
}

/**
 * Create a string, one per file or directory by recursively iterating
 * over a chain of directories and subdirectores until arriving at the leaf
 * node. For example, the following directory tree:
 *   src/
 *       headers/
 *               foo.h
 *       cpp/
 *           foo.cpp
 *           foo1.cpp -> foo.cpp
 *   local/
 *         lib/
 *             libc.so
 *   var/
 *
 *  can be expressed as a series of path names which
 *  we can pass to the NestedDirectory::add(), NestedDirectory::addSymlink()
 *  and NestedDirectory::addDirectory() methods:
 *
 *  NestedDirectory result;
 *  result.add("src/headers/foo.h");
 *  result.add("src/cpp/foo.cpp");
 *  result.addSymlink("src/cpp/foo1.cpp");
 *  result.add("local/lib/libc.so");
 *  result.addDirectory("var");
 */
void buildFlattenedPath(PathNodeMetaDataMap *map,
                        const buildboxcommon::Directory &directory,
                        const buildboxcommon::DirectoryTree &tree,
                        const std::string &dirName = "")
{
    // files
    for (const auto &node : directory.files()) {
        const std::string newFile = genNewPath(dirName, node.name());

        // collision detection for files is defined as
        // same file name but different digest or 'is_executable' flag
        const auto it = map->find(newFile);
        if (it != map->end() && (*it->second != node)) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "file collision: existing file ["
                    << it->second->path() << ":" << it->second->digest() << ":"
                    << std::boolalpha << it->second->isExecutable() << "]"
                    << " detected while attempting to add new file ["
                    << newFile << ":" << node.digest() << ":" << std::boolalpha
                    << node.is_executable());
        }
        map->emplace(newFile, std::make_shared<FileNodeMetaData>(
                                  newFile, node.digest(), node.is_executable(),
                                  node.node_properties()));
    }

    // symlinks
    for (const auto &node : directory.symlinks()) {
        const std::string newSymlinkName = genNewPath(dirName, node.name());

        // collision detection for symlinks is defined as
        // same name but different target, ie;
        // 1. /some/path/name1 -> ../target1   # OK
        // 2. /some/path/name1 -> ../target2   # BAD
        // 3. /some/path/name2 -> ../target1   # OK
        const auto it = map->find(newSymlinkName);
        if (it != map->end()) {
            // same path/name but different target - not allowed
            const auto *metaNode =
                dynamic_cast<SymlinkNodeMetaData *>(it->second.get());
            if (metaNode != nullptr && metaNode->target() != node.target()) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "error processing symlink: existing symlink ["
                        << it->first << " -> " << metaNode->target()
                        << "] has the same name but different target than new "
                           "symlink ["
                        << newSymlinkName << " -> " << node.target() << "]");
            }
        }
        map->emplace(newSymlinkName, std::make_shared<SymlinkNodeMetaData>(
                                         newSymlinkName, node.target()));
    }

    // subdirectories
    // no collision detection is needed at this level because we allow
    // directories with the same name in the merged output; if there are
    // collisions in the subdirectory data, it will be detected at the file and
    // symlink level
    for (const auto &node : directory.directories()) {
        const std::string newDirectoryPath = genNewPath(dirName, node.name());

        map->emplace(newDirectoryPath, std::make_shared<DirNodeMetaData>(
                                           newDirectoryPath, node.digest()));

        const auto subDirIt = tree.map.find(node.digest());
        if (subDirIt == tree.map.end()) {
            BUILDBOX_LOG_ERROR("error finding digest " << node.digest());
            continue;
        }
        buildFlattenedPath(map, subDirIt->second, tree, newDirectoryPath);
    }
}

void buildDigestDirectorySet(const buildboxcommon::DirectoryTree &tree,
                             std::unordered_set<Digest> *dsSet)
{
    for (const auto &[digest, _] : tree.map) {
        dsSet->insert(digest);
    }
}

// Given two directories, merge their node properties.
// To be binary-compatible, if both nodes don't have properties,
// return std::nullopt
std::optional<NodeProperties>
mergeDirectoryNodeProperties(const Directory &prev, const Directory &next)
{

    if (!prev.has_node_properties() && !next.has_node_properties()) {
        return std::nullopt;
    }
    if (!prev.has_node_properties()) {
        return next.node_properties();
    }
    if (!next.has_node_properties()) {
        return prev.node_properties();
    }

    NodeProperties result = prev.node_properties();
    if (next.node_properties().has_mtime()) {
        *result.mutable_mtime() = next.node_properties().mtime();
    }
    if (next.node_properties().has_unix_mode()) {
        *result.mutable_unix_mode() = next.node_properties().unix_mode();
    }

    // Merge K=V properties
    result.clear_properties();
    std::set<std::string> allPropertyNames;
    std::map<std::string, std::set<std::string>> prevProperties,
        nextProperties, mergedProperties;
    // Collect KVs into multimap. std::multimap is not used as values also need
    // to be sorted
    for (const auto &property : prev.node_properties().properties()) {
        allPropertyNames.emplace(property.name());
        prevProperties[property.name()].insert(property.value());
    }
    for (const auto &property : next.node_properties().properties()) {
        allPropertyNames.emplace(property.name());
        nextProperties[property.name()].insert(property.value());
    }
    for (const auto &name : allPropertyNames) {
        if (nextProperties.find(name) != nextProperties.end()) {
            mergedProperties[name] = nextProperties.at(name);
        }
        else {
            mergedProperties[name] = prevProperties.at(name);
        }
    }

    for (const auto &[name, values] : mergedProperties) {
        for (const auto &value : values) {
            NodeProperty property;
            property.set_name(name);
            property.set_value(value);
            *result.mutable_properties()->Add() = property;
        }
    }
    return result;
}

} // namespace

bool MergeUtil::createMergedDigest(const DirectoryTree &inputTree,
                                   const DirectoryTree &templateTree,
                                   Digest *rootDigest,
                                   digest_string_map *newDirectoryBlobs,
                                   DigestVector *mergedDirectoryList)
{
    const std::vector<DirectoryTree> treesToMerge = {inputTree, templateTree};
    return createMergedDigest(treesToMerge, rootDigest, newDirectoryBlobs,
                              mergedDirectoryList);
}

bool MergeUtil::createMergedDigest(
    const std::vector<DirectoryTree> &treesToMerge, Digest *rootDigest,
    digest_string_map *newDirectoryBlobs, DigestVector *mergedDirectoryList)
{
    if (treesToMerge.empty()) {
        BUILDBOX_LOG_ERROR(
            "Invalid arguments, no DirectoryTree's specified to merge");
        return false;
    }

    // Build up a set of all Directory Digests (dsSet) and a map of full path
    // names to node metadata (nodeMap)
    // Building up the nodeMap also detects collisions which we define
    // as files/directories with the same name but with different digests
    // eg:
    // proj/src/file.cpp
    // proj/headers/file.h
    //
    // If a collision is found we handle the exception by returning false
    std::unordered_set<Digest> dsSet;
    PathNodeMetaDataMap nodeMap;
    try {
        for (const DirectoryTree &tree : treesToMerge) {
            buildDigestDirectorySet(tree, &dsSet);
            buildFlattenedPath(&nodeMap, tree.map.at(tree.root), tree);
        }
    }
    catch (const std::runtime_error &e) {
        BUILDBOX_LOG_WARNING("Error encountered in no conflict merge\n" +
                             std::string(e.what()));
        return false;
    }

    // Iterate over the list of file/directory paths
    // and use the NestedDirectory component to
    // build a merged directory tree
    NestedDirectory result;
    for (const auto &it : nodeMap) {
        it.second->addToNestedDirectory(&result);
    }

    // Iterate over all the dirs/file and generate a new
    // merged root digest. Store all digest in newDirectoryBlob
    *rootDigest = result.to_digest(newDirectoryBlobs);

    // Place newly created merged directories into mergedDirectoryList
    if (mergedDirectoryList != nullptr) {
        for (const auto &it : *newDirectoryBlobs) {
            if (!dsSet.contains(it.first)) {
                mergedDirectoryList->emplace_back(it.first);
            }
        }
    }

    return true;
}

std::ostream &operator<<(std::ostream &out, const DirectoryTree &tree)
{
    PathNodeMetaDataMap map;
    buildFlattenedPath(&map, tree.map.at(tree.root), tree);
    for_each(map.begin(), map.end(),
             [&out](const PathNodeMetaDataMap::value_type &p) {
                 out << *p.second;
             });
    return out;
}

std::string
toString(const ::google::protobuf::RepeatedPtrField<buildboxcommon::Directory>
             &tree)
{
    std::ostringstream out;

    for (int i = 0; i < tree.size(); ++i) {
        const auto &directory = tree.Get(i);
        const auto digest =
            DigestGenerator::hash(directory.SerializeAsString());

        // files
        const auto &fileNodes = directory.files();
        for (int j = 0; j < fileNodes.size(); ++j) {
            out << "Directory[" << i << "](" << digest << ") --> FileNode["
                << j << "]: name = \"" << fileNodes[j].name()
                << "\", digest = \"" << fileNodes[j].digest()
                << "\", executable = " << std::boolalpha
                << fileNodes[j].is_executable() << "\n";
        }

        // symlinks
        const auto &symNodes = directory.symlinks();
        for (int j = 0; j < symNodes.size(); ++j) {
            out << "Directory[" << i << "](" << digest << ") --> SymlinkNode["
                << j << "]: name = \"" << symNodes[j].name()
                << "\", target = \"" << symNodes[j].target() << "\""
                << "\n";
        }

        // sub-directories
        const auto &dirNodes = directory.directories();
        for (int j = 0; j < dirNodes.size(); ++j) {
            out << "Directory[" << i << "](" << digest
                << ") --> DirectoryNode[" << j << "]: name = \""
                << dirNodes[j].name() << "\", digest = \""
                << dirNodes[j].digest() << "\""
                << "\n";
        }
    }

    return out.str();
}

namespace layer {

// Helper struct similar to `Directory` but with extra information to
// handle whiteouts
struct LayerDirInfo {
    std::map<std::string, DirectoryNode> directories;
    std::map<std::string, FileNode> files;
    std::map<std::string, SymlinkNode> symlinks;
    std::set<std::string> removals;
    bool opaqueRemoval{false};

    LayerDirInfo(const Directory &dir)
    {
        for (const auto &subdir : dir.directories()) {
            directories.emplace(subdir.name(), subdir);
        }
        for (const auto &symlink : dir.symlinks()) {
            symlinks.emplace(symlink.name(), symlink);
        }
        for (const auto &file : dir.files()) {
            if (file.name() == ".wh..wh..opq") {
                opaqueRemoval = true;
            }
            else if (file.name().find(".wh.") == 0) {
                removals.emplace(file.name().substr(4));
            }
            else {
                files.emplace(file.name(), file);
            }
        }
    }
};

Directory formatLayerDirectory(const Directory &dir)
{
    Directory result(dir);
    // remove whiteout files
    result.clear_files();
    for (const auto &file : dir.files()) {
        if (file.name().find(".wh.") != 0) {
            result.mutable_files()->Add()->CopyFrom(file);
        }
    }
    return result;
}

// Helper function to output an filesystem layer directory
// 1. Format it to remove whiteout files
// 2. Convert to Digest/Blob pair
// 3. Update lookup map and blob map
Digest outputLayerDirectory(const Directory &dir,
                            MergeUtil::DigestDirectoryMap &digestDirectoryMap,
                            digest_string_map &newDirectoriesBlobs)
{
    const Directory formattedDir = layer::formatLayerDirectory(dir);
    const std::string blob = formattedDir.SerializeAsString();
    Digest digest = DigestGenerator::hash(formattedDir);
    if (digestDirectoryMap.find(digest) == digestDirectoryMap.end()) {
        digestDirectoryMap.emplace(digest, formattedDir);
    }
    if (newDirectoriesBlobs.find(digest) == newDirectoriesBlobs.end()) {
        newDirectoriesBlobs.emplace(digest, blob);
    }
    return digest;
}

} // namespace layer

bool MergeUtil::createMergedLayersDigest(
    const std::vector<DirectoryTree> &treesToMerge, Digest *rootDigest,
    digest_string_map *newDirectoryBlobs)
{

    // Create a lookup map from digest to directory
    DigestDirectoryMap digestDirectoryMap;
    const auto emptyDigest = emptyDirDigest();
    digestDirectoryMap[emptyDigest] = s_emptyDir;
    std::vector<Digest> rootDigests;
    for (const DirectoryTree &tree : treesToMerge) {
        if (tree.root.hash().empty()) {
            BUILDBOX_LOG_ERROR(
                "Error merging layers: Tree without root digest");
            return false;
        }
        rootDigests.emplace_back(tree.root);
        for (const auto &[digest, dir] : tree.map) {
            digestDirectoryMap[digest] = dir;
        }
    }

    if (rootDigests.empty()) {
        // Nothing to merge
        return false;
    }
    try {
        // Ensure the .wh files in the first tree is handled and the root is
        // put into the output map
        Digest mergedDigest =
            MergeUtil::mergeLayers({emptyDigest}, {rootDigests[0]},
                                   digestDirectoryMap, *newDirectoryBlobs);
        // Merge remaining trees from left to right
        for (size_t idx = 1; idx < rootDigests.size(); idx++) {
            mergedDigest =
                MergeUtil::mergeLayers({mergedDigest}, {rootDigests[idx]},
                                       digestDirectoryMap, *newDirectoryBlobs);
        }

        rootDigest->CopyFrom(mergedDigest);
    }
    catch (const MergeFailure &e) {
        BUILDBOX_LOG_ERROR("Error merging layers: " << e.what());
        return false;
    }
    return true;
}

Digest MergeUtil::mergeLayers(const Digest &prev, const Digest &next,
                              DigestDirectoryMap &digestDirectoryMap,
                              digest_string_map &newDirectoriesBlobs)
{
    if (digestDirectoryMap.find(prev) == digestDirectoryMap.end()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            MergeFailure, "unable to find directory_digest="
                              << prev.hash() << "/" << prev.size_bytes()
                              << " in the digest directory map");
    }
    if (digestDirectoryMap.find(next) == digestDirectoryMap.end()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            MergeFailure, "unable to find directory_digest="
                              << next.hash() << "/" << next.size_bytes()
                              << " in the digest directory map");
    }

    Directory &prevDir = digestDirectoryMap.at(prev),
              &nextDir = digestDirectoryMap.at(next);

    auto emptyDigest = emptyDirDigest();

    // Shortcuts
    if (prev == emptyDigest && next == emptyDigest) {
        return emptyDigest;
    }
    if (prev == emptyDigest &&
        (newDirectoriesBlobs.find(next) != newDirectoriesBlobs.end())) {
        return next;
    }
    if (next == emptyDigest &&
        (newDirectoriesBlobs.find(prev) != newDirectoriesBlobs.end())) {
        return prev;
    }
    if (prev == next) {
        // If two non-trivial directories are the same, only recursively
        // visit one of them
        return mergeLayers(prev, emptyDigest, digestDirectoryMap,
                           newDirectoriesBlobs);
    }

    const layer::LayerDirInfo prevInfo(prevDir), nextInfo(nextDir);

    // If there is an opaque removal, we will throw away the previous node
    // completely.
    if (nextInfo.opaqueRemoval) {
        const Directory formattedNextDir =
            layer::formatLayerDirectory(nextDir);
        const Digest formattedNextDirDigest =
            DigestGenerator::hash(formattedNextDir);
        digestDirectoryMap[formattedNextDirDigest] = formattedNextDir;
        return mergeLayers(emptyDigest, formattedNextDirDigest,
                           digestDirectoryMap, newDirectoriesBlobs);
    }

    // collect all names
    std::set<std::string> allFiles;
    std::transform(prevInfo.files.cbegin(), prevInfo.files.cend(),
                   std::inserter(allFiles, allFiles.end()),
                   [](const auto &it) { return it.first; });
    std::transform(nextInfo.files.cbegin(), nextInfo.files.cend(),
                   std::inserter(allFiles, allFiles.end()),
                   [](const auto &it) { return it.first; });
    std::set<std::string> allDirs;
    std::transform(prevInfo.directories.cbegin(), prevInfo.directories.cend(),
                   std::inserter(allDirs, allDirs.end()),
                   [](const auto &it) { return it.first; });
    std::transform(nextInfo.directories.cbegin(), nextInfo.directories.cend(),
                   std::inserter(allDirs, allDirs.end()),
                   [](const auto &it) { return it.first; });
    std::set<std::string> allSymlinks;
    std::transform(prevInfo.symlinks.cbegin(), prevInfo.symlinks.cend(),
                   std::inserter(allSymlinks, allSymlinks.end()),
                   [](const auto &it) { return it.first; });
    std::transform(nextInfo.symlinks.cbegin(), nextInfo.symlinks.cend(),
                   std::inserter(allSymlinks, allSymlinks.end()),
                   [](const auto &it) { return it.first; });

    Directory result;
    // Node property
    const auto mergedProperties =
        mergeDirectoryNodeProperties(prevDir, nextDir);
    if (mergedProperties.has_value()) {
        *result.mutable_node_properties() = mergedProperties.value();
    }

    // Merge
    for (const auto &name : allFiles) {
        if (nextInfo.removals.find(name) != nextInfo.removals.end() ||
            nextInfo.directories.find(name) != nextInfo.directories.end() ||
            nextInfo.symlinks.find(name) != nextInfo.symlinks.end()) {
            continue;
        }
        if (nextInfo.files.find(name) != nextInfo.files.end()) {
            result.mutable_files()->Add()->CopyFrom(nextInfo.files.at(name));
        }
        else {
            result.mutable_files()->Add()->CopyFrom(prevInfo.files.at(name));
        }
    }
    for (const auto &name : allSymlinks) {
        if (nextInfo.removals.find(name) != nextInfo.removals.end() ||
            nextInfo.files.find(name) != nextInfo.files.end() ||
            nextInfo.directories.find(name) != nextInfo.directories.end()) {
            continue;
        }
        if (nextInfo.symlinks.find(name) != nextInfo.symlinks.end()) {
            result.mutable_symlinks()->Add()->CopyFrom(
                nextInfo.symlinks.at(name));
        }
        else {
            result.mutable_symlinks()->Add()->CopyFrom(
                prevInfo.symlinks.at(name));
        }
    }
    for (const auto &name : allDirs) {
        if (nextInfo.removals.find(name) != nextInfo.removals.end() ||
            nextInfo.files.find(name) != nextInfo.files.end() ||
            nextInfo.symlinks.find(name) != nextInfo.symlinks.end()) {
            continue;
        }
        // Merge subdirs
        Digest prevSubdirDigest = emptyDigest, nextSubdirDigest = emptyDigest;
        if (prevInfo.directories.find(name) != prevInfo.directories.end()) {
            prevSubdirDigest = prevInfo.directories.at(name).digest();
        }
        if (nextInfo.directories.find(name) != nextInfo.directories.end()) {
            nextSubdirDigest = nextInfo.directories.at(name).digest();
        }
        Digest subdirDigest =
            MergeUtil::mergeLayers(prevSubdirDigest, nextSubdirDigest,
                                   digestDirectoryMap, newDirectoriesBlobs);
        DirectoryNode *subdirNode = result.mutable_directories()->Add();
        subdirNode->set_name(name);
        subdirNode->mutable_digest()->CopyFrom(subdirDigest);
    }

    return layer::outputLayerDirectory(result, digestDirectoryMap,
                                       newDirectoriesBlobs);
}

} // namespace buildboxcommon
