/*
 * Copyright 2019 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 <buildboxcasd_server.h>

#include <buildboxcasd_casinstance.h>
#include <buildboxcasd_proxymodefixture.h>
#include <buildboxcommon_mergeutil.h>
#include <buildboxcommon_temporarydirectory.h>

#include <gmock/gmock.h>
#include <grpcpp/client_context.h>
#include <grpcpp/create_channel.h>
#include <grpcpp/impl/codegen/call.h>
#include <grpcpp/server_context.h>
#include <grpcpp/support/sync_stream.h>
#include <gtest/gtest.h>

#include <fcntl.h>
#include <thread>

using namespace buildboxcasd;
using buildboxcommon::DigestGenerator;

namespace {
const auto digestFunctionInitializer = []() {
    buildboxcommon::DigestGenerator::init();
    return 0;
}();
} // namespace

class MergeFixture : public CasProxyModeFixture {
  protected:
    Digest d_badDigest; // remains uninitialized
    Digest d_emptyInputDigest;
    Digest d_inputTreeWithExecutableTrue;
    Digest d_inputTreeWithExecutableFalse;
    Digest d_inputDigestWithOverlapWithoutConflict;
    Digest d_inputDigestWithOverlapWithConflict;
    Digest d_chrootTemplateDigest;

    MergeFixture()
    {
        prepareEmptyInputDigest(&d_emptyInputDigest);
        prepareInputDigest(&d_inputTreeWithExecutableTrue);
        prepareInputDigest(&d_inputTreeWithExecutableFalse, false);
        prepareInputDigestWithOverlap(
            &d_inputDigestWithOverlapWithoutConflict);
        prepareInputDigestWithOverlap(
            &d_inputDigestWithOverlapWithConflict,
            "lib_so_contents_but_with_different_data");
        prepareTemplateDigest(&d_chrootTemplateDigest);
    }

    void prepareEmptyInputDigest(Digest *digest)
    {
        Directory empty_directory;
        const auto empty_directory_serialized =
            empty_directory.SerializeAsString();
        const auto empty_directory_digest =
            DigestGenerator::hash(empty_directory_serialized);

        // .
        Directory root_directory;
        DirectoryNode *emptyNode = root_directory.add_directories();
        emptyNode->set_name("empty");
        emptyNode->mutable_digest()->CopyFrom(empty_directory_digest);
        const auto root_directory_serialized =
            root_directory.SerializeAsString();
        const auto root_directory_digest =
            DigestGenerator::hash(root_directory_serialized);

        // set the output variable
        *digest = root_directory_digest;

        // store all in remote CAS
        remote_storage->writeBlob(root_directory_digest,
                                  root_directory_serialized);
    }

    void prepareInputDigest(Digest *digest, const bool isExecutable = true)
    {
        /* Creates the following directory structure:
         *
         * ./
         *   src/
         *       build.sh*
         *       headers/
         *               file1.h
         *               file2.h
         *               file3.h
         *       cpp/
         *           file1.cpp
         *           file2.cpp
         *           file3.cpp
         *           foo.cpp --> file.cpp
         *       shared/
         *              file4.cpp --> ../cpp/file3.cpp
         */

        // ./src/headers
        Directory headers_directory;
        std::vector<std::string> headerFiles = {"file1.h", "file2.h",
                                                "file3.h"};
        for (const auto &file : headerFiles) {
            FileNode *fileNode = headers_directory.add_files();
            fileNode->set_name(file);
            fileNode->set_is_executable(false);
            fileNode->mutable_digest()->CopyFrom(
                DigestGenerator::hash(file + "_contents"));
        }
        const auto headers_directory_serialized =
            headers_directory.SerializeAsString();
        const auto headers_directory_digest =
            DigestGenerator::hash(headers_directory_serialized);

        // ./src/cpp
        Directory cpp_directory;
        std::vector<std::string> cppFiles = {"file1.cpp", "file2.cpp",
                                             "file3.cpp"};
        for (const auto &file : cppFiles) {
            FileNode *fileNode = cpp_directory.add_files();
            fileNode->set_name(file);
            fileNode->set_is_executable(false);
            fileNode->mutable_digest()->CopyFrom(
                DigestGenerator::hash(file + "_contents"));
        }
        SymlinkNode *fooNode = cpp_directory.add_symlinks();
        fooNode->set_name("foo.cpp");
        fooNode->set_target("file.cpp");
        const auto cpp_directory_serialized =
            cpp_directory.SerializeAsString();
        const auto cpp_directory_digest =
            DigestGenerator::hash(cpp_directory_serialized);

        // ./src/shared
        Directory shared_directory;
        SymlinkNode *symNode = shared_directory.add_symlinks();
        symNode->set_name("file4.cpp");
        symNode->set_target("../cpp/file3.cpp");
        const auto shared_directory_serialized =
            shared_directory.SerializeAsString();
        const auto shared_directory_digest =
            DigestGenerator::hash(shared_directory_serialized);

        // ./src
        Directory src_directory;

        // add subdir 'headers'
        DirectoryNode *headersNode = src_directory.add_directories();
        headersNode->set_name("headers");
        headersNode->mutable_digest()->CopyFrom(headers_directory_digest);

        // add subdir 'cpp'
        DirectoryNode *cppNode = src_directory.add_directories();
        cppNode->set_name("cpp");
        cppNode->mutable_digest()->CopyFrom(cpp_directory_digest);

        // add subdir 'shared'
        DirectoryNode *sharedNode = src_directory.add_directories();
        sharedNode->set_name("shared");
        sharedNode->mutable_digest()->CopyFrom(shared_directory_digest);

        // add local file
        FileNode *fileNode = src_directory.add_files();
        fileNode->set_name("build.sh");
        fileNode->set_is_executable(isExecutable);
        fileNode->mutable_digest()->CopyFrom(
            DigestGenerator::hash("build.sh_contents"));

        // get the digest of 'src' directory
        const auto src_directory_serialized =
            src_directory.SerializeAsString();
        const auto src_directory_digest =
            DigestGenerator::hash(src_directory_serialized);

        // .
        Directory root_directory;
        DirectoryNode *srcNode = root_directory.add_directories();
        srcNode->set_name("src");
        srcNode->mutable_digest()->CopyFrom(src_directory_digest);
        const auto root_directory_serialized =
            root_directory.SerializeAsString();
        const auto root_directory_digest =
            DigestGenerator::hash(root_directory_serialized);

        // set the output variable
        *digest = root_directory_digest;

        // store all in remote CAS
        remote_storage->writeBlob(headers_directory_digest,
                                  headers_directory_serialized);
        remote_storage->writeBlob(cpp_directory_digest,
                                  cpp_directory_serialized);
        remote_storage->writeBlob(shared_directory_digest,
                                  shared_directory_serialized);
        remote_storage->writeBlob(src_directory_digest,
                                  src_directory_serialized);
        remote_storage->writeBlob(root_directory_digest,
                                  root_directory_serialized);
    }

    void prepareInputDigestWithOverlap(
        Digest *digest,
        const std::string &forcedCollisionData = "libc_so_contents")
    {
        /* Creates the following directory structure:
         *
         * ./
         *   src/
         *       headers/
         *               foo.h
         *       cpp/
         *           foo.cpp
         *   local/
         *         lib/
         *             libc.so
         */

        // ./src/headers
        Directory headers_directory;
        FileNode *headersFileNodes = headers_directory.add_files();
        headersFileNodes->set_name("foo.h");
        headersFileNodes->set_is_executable(false);
        headersFileNodes->mutable_digest()->CopyFrom(
            DigestGenerator::hash("foo_h_contents"));
        const auto headers_directory_serialized =
            headers_directory.SerializeAsString();
        const auto headers_directory_digest =
            DigestGenerator::hash(headers_directory_serialized);

        // ./src/cpp
        Directory cpp_directory;
        FileNode *cppFileNodes = cpp_directory.add_files();
        cppFileNodes->set_name("foo.cpp");
        cppFileNodes->set_is_executable(false);
        cppFileNodes->mutable_digest()->CopyFrom(
            DigestGenerator::hash("foo_cpp_contents"));
        const auto cpp_directory_serialized =
            cpp_directory.SerializeAsString();
        const auto cpp_directory_digest =
            DigestGenerator::hash(cpp_directory_serialized);

        // ./src
        Directory src_directory;
        DirectoryNode *headersNode = src_directory.add_directories();
        headersNode->set_name("headers");
        headersNode->mutable_digest()->CopyFrom(headers_directory_digest);
        DirectoryNode *cppNode = src_directory.add_directories();
        cppNode->set_name("cpp");
        cppNode->mutable_digest()->CopyFrom(cpp_directory_digest);
        const auto src_directory_serialized =
            src_directory.SerializeAsString();
        const auto src_directory_digest =
            DigestGenerator::hash(src_directory_serialized);

        // ./lib/libc.so
        Directory lib_directory;
        FileNode *libFileNode = lib_directory.add_files();
        libFileNode->set_name("libc.so");
        libFileNode->set_is_executable(false);
        libFileNode->mutable_digest()->CopyFrom(
            DigestGenerator::hash(forcedCollisionData));
        const auto lib_directory_serialized =
            lib_directory.SerializeAsString();
        const auto lib_directory_digest =
            DigestGenerator::hash(lib_directory_serialized);

        // ./local
        Directory local_directory;
        DirectoryNode *libNode = local_directory.add_directories();
        libNode->set_name("lib");
        libNode->mutable_digest()->CopyFrom(lib_directory_digest);
        const auto local_directory_serialized =
            local_directory.SerializeAsString();
        const auto local_directory_digest =
            DigestGenerator::hash(local_directory_serialized);

        // .
        Directory root_directory;
        DirectoryNode *srcNode = root_directory.add_directories();
        srcNode->set_name("src");
        srcNode->mutable_digest()->CopyFrom(src_directory_digest);
        DirectoryNode *localNode = root_directory.add_directories();
        localNode->set_name("local");
        localNode->mutable_digest()->CopyFrom(local_directory_digest);
        const auto root_directory_serialized =
            root_directory.SerializeAsString();
        const auto root_directory_digest =
            DigestGenerator::hash(root_directory_serialized);

        // set the output variable
        *digest = root_directory_digest;

        // store all in remote CAS
        remote_storage->writeBlob(headers_directory_digest,
                                  headers_directory_serialized);
        remote_storage->writeBlob(cpp_directory_digest,
                                  cpp_directory_serialized);
        remote_storage->writeBlob(src_directory_digest,
                                  src_directory_serialized);
        remote_storage->writeBlob(lib_directory_digest,
                                  lib_directory_serialized);
        remote_storage->writeBlob(local_directory_digest,
                                  local_directory_serialized);
        remote_storage->writeBlob(root_directory_digest,
                                  root_directory_serialized);
    }

    void prepareTemplateDigest(Digest *digest)
    {
        /* Creates the following directory structure:
         *
         * ./
         *   include/
         *           time.h
         *           sys/
         *               stat.h
         *   local/
         *         lib/
         *             libc.so
         *   var/
         */

        // ./include/sys
        Directory sys_directory;
        FileNode *sysFileNodes = sys_directory.add_files();
        sysFileNodes->set_name("stat.h");
        sysFileNodes->set_is_executable(false);
        sysFileNodes->mutable_digest()->CopyFrom(
            DigestGenerator::hash("stat_h_contents"));
        const auto sys_directory_serialized =
            sys_directory.SerializeAsString();
        const auto sys_directory_digest =
            DigestGenerator::hash(sys_directory_serialized);

        // ./include
        Directory include_directory;
        FileNode *includeFileNodes = include_directory.add_files();
        includeFileNodes->set_name("time.h");
        includeFileNodes->set_is_executable(false);
        includeFileNodes->mutable_digest()->CopyFrom(
            DigestGenerator::hash("time_h_contents"));
        DirectoryNode *sysNode = include_directory.add_directories();
        sysNode->set_name("sys");
        sysNode->mutable_digest()->CopyFrom(sys_directory_digest);
        const auto include_directory_serialized =
            include_directory.SerializeAsString();
        const auto include_directory_digest =
            DigestGenerator::hash(include_directory_serialized);

        // ./local/lib
        Directory lib_directory;
        FileNode *libFileNode = lib_directory.add_files();
        libFileNode->set_name("libc.so");
        libFileNode->set_is_executable(false);
        libFileNode->mutable_digest()->CopyFrom(
            DigestGenerator::hash("libc_so_contents"));
        const auto lib_directory_serialized =
            lib_directory.SerializeAsString();
        const auto lib_directory_digest =
            DigestGenerator::hash(lib_directory_serialized);

        // ./local
        Directory local_directory;
        DirectoryNode *libNode = local_directory.add_directories();
        libNode->set_name("lib");
        libNode->mutable_digest()->CopyFrom(lib_directory_digest);
        const auto local_directory_serialized =
            local_directory.SerializeAsString();
        const auto local_directory_digest =
            DigestGenerator::hash(local_directory_serialized);

        // ./var
        Directory var_directory;
        const auto var_directory_serialized =
            var_directory.SerializeAsString();
        const auto var_directory_digest =
            DigestGenerator::hash(var_directory_serialized);

        // .
        Directory root_directory;
        // add include to root
        DirectoryNode *includeNode = root_directory.add_directories();
        includeNode->set_name("include");
        includeNode->mutable_digest()->CopyFrom(include_directory_digest);

        // add local to root
        DirectoryNode *localNode = root_directory.add_directories();
        localNode->set_name("local");
        localNode->mutable_digest()->CopyFrom(local_directory_digest);

        // add var to root
        DirectoryNode *varNode = root_directory.add_directories();
        varNode->set_name("var");
        varNode->mutable_digest()->CopyFrom(var_directory_digest);

        const auto root_directory_serialized =
            root_directory.SerializeAsString();
        const auto root_directory_digest =
            DigestGenerator::hash(root_directory_serialized);

        // set the output variable
        *digest = root_directory_digest;

        // store all in remote CAS
        remote_storage->writeBlob(sys_directory_digest,
                                  sys_directory_serialized);
        remote_storage->writeBlob(include_directory_digest,
                                  include_directory_serialized);
        remote_storage->writeBlob(lib_directory_digest,
                                  lib_directory_serialized);
        remote_storage->writeBlob(local_directory_digest,
                                  local_directory_serialized);
        remote_storage->writeBlob(var_directory_digest,
                                  var_directory_serialized);
        remote_storage->writeBlob(root_directory_digest,
                                  root_directory_serialized);
    }

    bool batchUpdateBlobs(const buildboxcommon::digest_string_map &dsMap)
    {
        grpc::ClientContext context;
        BatchUpdateBlobsRequest updateRequest;
        BatchUpdateBlobsResponse updateResponse;
        for (const auto &it : dsMap) {
            auto *request = updateRequest.add_requests();
            request->mutable_digest()->CopyFrom(it.first);
            request->mutable_data()->assign(it.second);
        }
        cas_stub->BatchUpdateBlobs(&context, updateRequest, &updateResponse);
        for (const auto &response : updateResponse.responses()) {
            if (response.status().code(), grpc::OK) {
                return false;
            }
        }

        return true;
    }
};

TEST_P(MergeFixture, MergeFailBadInputTree)
{
    EXPECT_THROW(cas_client->getTree(d_badDigest), std::runtime_error);
}

TEST_P(MergeFixture, MergeSuccessEmptyInputTree)
{
    // input request
    ASSERT_TRUE(remote_storage->hasBlob(d_emptyInputDigest));
    DirectoryTree inputTree = cas_client->getTree(d_emptyInputDigest);
    ASSERT_EQ(inputTree.map.size(), 2);
    std::cout << "inputTree contains " << inputTree.map.size()
              << " directories\n"
              << inputTree << '\n';

    // template request
    ASSERT_TRUE(remote_storage->hasBlob(d_chrootTemplateDigest));
    DirectoryTree templateTree = cas_client->getTree(d_chrootTemplateDigest);
    ASSERT_EQ(templateTree.map.size(), 6);
    std::cout << "templateTree contains " << templateTree.map.size()
              << " directories\n"
              << templateTree << '\n';

    // merge
    Digest mergedRootDigest;
    buildboxcommon::digest_string_map dsMap;
    const bool result = MergeUtil::createMergedDigest(
        inputTree, templateTree, &mergedRootDigest, &dsMap);
    ASSERT_TRUE(result);

    // write the merged tree
    ASSERT_TRUE(batchUpdateBlobs(dsMap));

    // read the merged blob
    DirectoryTree mergedTree = cas_client->getTree(mergedRootDigest);
    ASSERT_EQ(mergedTree.map.size(), 6);
    std::cout << "mergedTree contains " << mergedTree.map.size()
              << " directories\n"
              << mergedTree << '\n';
}

TEST_P(MergeFixture, MergeSuccessNoOverlap)
{
    ASSERT_TRUE(remote_storage->hasBlob(d_inputTreeWithExecutableTrue));
    DirectoryTree inputTree =
        cas_client->getTree(d_inputTreeWithExecutableTrue);
    ASSERT_EQ(inputTree.map.size(), 5);
    std::cout << "inputTree contains " << inputTree.map.size()
              << " directories\n"
              << inputTree << '\n';

    // template request
    DirectoryTree templateTree = cas_client->getTree(d_chrootTemplateDigest);
    ASSERT_EQ(templateTree.map.size(), 6);
    std::cout << "templateTree contains " << templateTree.map.size()
              << " directories\n"
              << templateTree << '\n';

    // merge
    Digest mergedRootDigest;
    buildboxcommon::digest_string_map dsMap;
    const bool result = MergeUtil::createMergedDigest(
        inputTree, templateTree, &mergedRootDigest, &dsMap);
    ASSERT_TRUE(result);

    // write the merged tree
    ASSERT_TRUE(batchUpdateBlobs(dsMap));

    // read the merged blob
    DirectoryTree mergedTree = cas_client->getTree(mergedRootDigest);
    ASSERT_EQ(mergedTree.map.size(), 10);
    std::cout << "mergedTree contains " << mergedTree.map.size()
              << " directories\n"
              << mergedTree << '\n';
}

TEST_P(MergeFixture, MergeSuccessOverlapWithoutConflict)
{
    ASSERT_TRUE(
        remote_storage->hasBlob(d_inputDigestWithOverlapWithoutConflict));
    DirectoryTree inputTree =
        cas_client->getTree(d_inputDigestWithOverlapWithoutConflict);
    ASSERT_EQ(inputTree.map.size(), 6);
    std::cout << "inputTree contains " << inputTree.map.size()
              << " directories\n"
              << inputTree << '\n';

    // template request
    ASSERT_TRUE(remote_storage->hasBlob(d_chrootTemplateDigest));
    DirectoryTree templateTree = cas_client->getTree(d_chrootTemplateDigest);
    ASSERT_EQ(templateTree.map.size(), 6);
    std::cout << "templateTree contains " << templateTree.map.size()
              << " directories\n"
              << templateTree << '\n';

    // merge
    Digest mergedRootDigest;
    buildboxcommon::digest_string_map dsMap;
    const bool result = MergeUtil::createMergedDigest(
        inputTree, templateTree, &mergedRootDigest, &dsMap);
    ASSERT_TRUE(result);

    // write the merged tree
    ASSERT_TRUE(batchUpdateBlobs(dsMap));

    // read the merged blob
    DirectoryTree mergedTree = cas_client->getTree(mergedRootDigest);
    ASSERT_EQ(mergedTree.map.size(), 9);
    std::cout << "mergedTree contains " << mergedTree.map.size()
              << " directories\n"
              << mergedTree << '\n';
}

TEST_P(MergeFixture, MergeFailOverlapWithConflict)
{
    ASSERT_TRUE(remote_storage->hasBlob(d_inputDigestWithOverlapWithConflict));
    DirectoryTree inputTree =
        cas_client->getTree(d_inputDigestWithOverlapWithConflict);
    ASSERT_EQ(inputTree.map.size(), 6);
    std::cout << "inputTree contains " << inputTree.map.size()
              << " directories\n"
              << inputTree << '\n';

    // template request
    ASSERT_TRUE(remote_storage->hasBlob(d_chrootTemplateDigest));
    DirectoryTree templateTree = cas_client->getTree(d_chrootTemplateDigest);
    ASSERT_EQ(templateTree.map.size(), 6);
    std::cout << "templateTree contains " << templateTree.map.size()
              << " directories\n"
              << templateTree << '\n';

    // merge should fail
    Digest mergedRootDigest;
    buildboxcommon::digest_string_map dsMap;
    const bool result = MergeUtil::createMergedDigest(
        inputTree, templateTree, &mergedRootDigest, &dsMap);
    ASSERT_FALSE(result);
}

TEST_P(MergeFixture, MergeMismatchIsExecutable)
{
    // inputTrue request
    ASSERT_TRUE(remote_storage->hasBlob(d_inputTreeWithExecutableTrue));
    DirectoryTree inputTrueTree =
        cas_client->getTree(d_inputTreeWithExecutableTrue);
    ASSERT_EQ(inputTrueTree.map.size(), 5);
    std::cout << "inputTrueTree contains " << inputTrueTree.map.size()
              << " directories\n"
              << inputTrueTree << '\n';

    // inputFalse request
    ASSERT_TRUE(remote_storage->hasBlob(d_inputTreeWithExecutableFalse));
    DirectoryTree inputFalseTree =
        cas_client->getTree(d_inputTreeWithExecutableFalse);
    ASSERT_EQ(inputFalseTree.map.size(), 5);
    std::cout << "inputFalseTree contains " << inputFalseTree.map.size()
              << " directories\n"
              << inputFalseTree << '\n';

    // merge should fail
    Digest mergedRootDigest;
    buildboxcommon::digest_string_map dsMap;
    const bool result = MergeUtil::createMergedDigest(
        inputTrueTree, inputFalseTree, &mergedRootDigest, &dsMap);
    ASSERT_FALSE(result);
}

INSTANTIATE_TEST_SUITE_P(
    StaticInstance, MergeFixture,
    testing::Values(std::make_tuple(STATIC_INSTANCE,
                                    FINDMISSINGBLOBS_CACHE_OFF, std::string(),
                                    NON_THREADED, AC_VALIDATION_CACHE_OFF)));
