/*
 * 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 <buildboxrun_oci.h>

#include <buildboxcommon_merklize.h>
#include <buildboxcommon_remoteexecutionclient.h>

#include <build/bazel/remote/execution/v2/remote_execution_mock.grpc.pb.h>
#include <build/buildgrid/local_cas_mock.grpc.pb.h>
#include <google/bytestream/bytestream_mock.grpc.pb.h>
#include <grpcpp/test/mock_stream.h>

#include <filesystem>
#include <gtest/gtest.h>

#include <memory>
#include <string>
#include <vector>

using namespace build::buildgrid;
using namespace build::bazel::remote::execution::v2;
using namespace buildboxcommon;
using namespace buildboxrun::oci;
using namespace testing;

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

const int64_t MAX_BATCH_SIZE_BYTES = 64;

class OCIRunnerFixture : public OCIRunner, public ::testing::Test {};

std::shared_ptr<CASClient>
createCASClient(const TemporaryDirectory &staged_directory)
{
    auto bytestreamClient =
        std::make_shared<google::bytestream::MockByteStreamStub>();
    auto casClient = std::make_shared<MockContentAddressableStorageStub>();
    auto localCasClient =
        std::make_shared<MockLocalContentAddressableStorageStub>();
    auto capabilitiesClient = std::make_shared<MockCapabilitiesStub>();

    auto grpcClient = std::make_shared<GrpcClient>();
    auto client = std::make_shared<CASClient>(grpcClient);
    client->init(bytestreamClient, casClient, localCasClient, nullptr,
                 MAX_BATCH_SIZE_BYTES);

    grpc::testing::MockClientReaderWriter<
        typename build::buildgrid::StageTreeRequest,
        typename build::buildgrid::StageTreeResponse> *reader_writer =
        new grpc::testing::MockClientReaderWriter<
            typename build::buildgrid::StageTreeRequest,
            typename build::buildgrid::StageTreeResponse>();

    EXPECT_CALL(*localCasClient.get(), StageTreeRaw(_))
        .WillOnce(Return(reader_writer));

    // The client will issue 2 requests: the actual `StageTreeRequest` and
    // an empty message to indicate to the server that it can clean up.
    EXPECT_CALL(*reader_writer, Write(_, _))
        .Times(2)
        .WillRepeatedly(Return(true));

    EXPECT_CALL(*reader_writer, WritesDone());

    StageTreeResponse response;
    // Returning a valid directory:
    response.set_path(staged_directory.name());
    EXPECT_CALL(*reader_writer, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(response), Return(true)));

    return client;
}

TEST_F(OCIRunnerFixture, TestMakeDigest)
{
    const std::string property_name = "chrootRootDigest";
    Platform platform;
    Platform_Property &property = *platform.add_properties();
    property.set_name(property_name);
    property.set_value("abc/1234");

    const auto digest = this->getBundleDigest(platform);
    ASSERT_EQ(digest.hash(), "abc");
    ASSERT_EQ(digest.size_bytes(), 1234);
}

TEST_F(OCIRunnerFixture, TestMakeDigest_InvalidAssetQualifier)
{
    static const std::string assetQualifierPropName = "chrootAssetQualifier";

    Platform platform;
    Platform_Property &property = *platform.add_properties();
    property.set_name(assetQualifierPropName);
    property.set_value("abc/1234");

    EXPECT_THROW({ this->getBundleDigest(platform); }, std::runtime_error);
}

TEST_F(OCIRunnerFixture, TestMakeDigest_NoAssetUrl_FallbackToDigest)
{
    static const std::string assetQualifierPropName = "chrootAssetQualifier";
    static const std::string digestPropName = "chrootRootDigest";

    Platform platform;
    Platform_Property &property = *platform.add_properties();
    property.set_name(assetQualifierPropName);
    property.set_value("abc/1234");
    Platform_Property &digestProperty = *platform.add_properties();
    digestProperty.set_name(digestPropName);
    digestProperty.set_value("abc/1234");

    const auto digest = this->getBundleDigest(platform);
    ASSERT_EQ(digest.hash(), "abc");
    ASSERT_EQ(digest.size_bytes(), 1234);
}

Command createCommandPaths()
{
    Command cmd;
    cmd.add_output_paths("foo24");
    cmd.add_output_paths("bar1/bar2/bar3/bar4");
    cmd.add_output_paths("bar1/bar2/bar3/bar4");
    cmd.add_output_paths("dir3/dir1");

    return cmd;
};

class ContainerHostFixture : public ::testing::Test {};

// Test 1 for prepareFilesystem(): checks to make sure that all parent
// directories have been created
TEST_F(ContainerHostFixture, TestPrepareFilesystem_DoNotExist)
{
    TemporaryDirectory staged_directory;
    auto cmd = createCommandPaths();
    OCIRunner::ContainerHost host(nullptr, cmd, DigestGenerator::hash(""),
                                  createCASClient(staged_directory),
                                  StageTreeRequest_StagingMode_DEFAULT);

    std::string path(staged_directory.name());
    host.prepareFilesystem();

    const std::filesystem::path path1 = path + "/rootfs/bar1/bar2/bar3";
    const std::filesystem::path path2 = path + "/rootfs/dir3";
    ASSERT_TRUE(exists(path1)) << "prepareFilesystem(): A parent directory "
                                  "was unexpectedly not created";
    ASSERT_TRUE(exists(path2)) << "prepareFilesystem(): A parent directory "
                                  "was unexpectedly not created";

    const std::filesystem::path path3 = path + "/rootfs/foo24";
    ASSERT_FALSE(exists(path3))
        << "prepareFilesystem(): Leaf path was unexpectedly created";
}

Command createCommandFilesAndDirs()
{
    Command cmd;
    cmd.add_output_paths("fo24");
    cmd.add_output_paths("fo42");
    cmd.add_output_paths("ba1/ba2/ba3/ba4");
    cmd.add_output_paths("di3/di2/di1");

    return cmd;
};

// Test 2 for prepareFilesystem(): checks to make sure the leaf child
// directories have not been created
TEST_F(ContainerHostFixture, TestPrepareFilesystem_Exist)
{
    TemporaryDirectory staged_directory;
    auto cmd = createCommandFilesAndDirs();
    OCIRunner::ContainerHost host(nullptr, cmd, DigestGenerator::hash(""),
                                  createCASClient(staged_directory),
                                  StageTreeRequest_StagingMode_DEFAULT);

    std::string path(staged_directory.name());
    host.prepareFilesystem();

    const std::filesystem::path file1 = path + "/rootfs/ba1/ba2/ba3";
    const std::filesystem::path dir1 = path + "/rootfs/di3/di2";
    ASSERT_TRUE(exists(file1)) << "prepareFilesystem(): A parent directory "
                                  "was unexpectedly not created";
    ASSERT_TRUE(exists(dir1)) << "prepareFilesystem(): A parent directory "
                                 "was unexpectedly not created";

    const std::filesystem::path file2 = path + "/rootfs/fo24";
    const std::filesystem::path dir2 = path + "/rootfs/fo42";
    ASSERT_FALSE(exists(file2)) << "prepareFilesystem(): Leaf file in the "
                                   "path was unexpectedly created";
    ASSERT_FALSE(exists(dir2)) << "prepareFilesystem(): Leaf directory in the "
                                  "path was unexpectedly created";
}

Command createCommandInvalid()
{
    Command cmd;
    cmd.add_output_paths("bar0");
    cmd.add_output_paths(
        "bar6/             _-o#&&*''''?d:>b\\_\n"
        "          _o/\"`''  '',, dMF9MMMMMHo_\n"
        "       .o&#'        `\"MbHMMMMMMMMMMMHo.\n"
        "     .o\"\" '         vodM*$&&HMMMMMMMMMM?.\n"
        "    ,'              $M&ood,~'`(&##MMMMMMH\\\n"
        "   /               ,MMMMMMM#b?#bobMMMMHMMML\n"
        "  &              ?MMMMMMMMMMMMMMMMM7MMM$R*Hk\n"
        " ?$.            :MMMMMMMMMMMMMMMMMMM/HMMM|`*L\n"
        "|               |MMMMMMMMMMMMMMMMMMMMbMH'   T,\n"
        "$H#:            `*MMMMMMMMMMMMMMMMMMMMb#}'  `?\n"
        "]MMH#             \"\"*\"\"\"\"*#MMMMMMMMMMMMM'    -\n"
        "MMMMMb_                   |MMMMMMMMMMMP'     :\n"
        "HMMMMMMMHo                 `MMMMMMMMMT       .\n"
        "?MMMMMMMMP                  9MMMMMMMM}       -\n"
        "-?MMMMMMM                  |MMMMMMMMM?,d-    '\n"
        " :|MMMMMM-                 `MMMMMMMT .M|.   :\n"
        "  .9MMM[                    &MMMMM*' `'    .\n"
        "   :9MMk                    `MMM#\"        -\n"
        "     &M}                     `          .-\n"
        "      `&.                             .\n"
        "        `~,   .                     ./\n"
        "            . _                  .-\n"
        "              '`--._,dd###pp=\"\"'/bar9");
    return cmd;
};

// Test 3 for prepareFilesystem(): Checks exception handling when attempting to
// create an invalid file name
TEST_F(ContainerHostFixture, TestPrepareFilesystem_InvalidFile)
{
    TemporaryDirectory staged_directory;
    auto cmd = createCommandInvalid();
    OCIRunner::ContainerHost host(nullptr, cmd, DigestGenerator::hash(""),
                                  createCASClient(staged_directory),
                                  StageTreeRequest_StagingMode_DEFAULT);

    EXPECT_ANY_THROW(host.prepareFilesystem())
        << "prepareFilesystem(): Unexpectedly did not throw an exception when "
           "failing to create the file";
}

Command createCommandInvalid2()
{
    Command cmd;
    cmd.add_output_paths("bar10");
    cmd.add_output_paths(
        "bar15/             _-o#&&*''''?d:>b\\_\n"
        "          _o/\"`''  '',, dMF9MMMMMHo_\n"
        "       .o&#'        `\"MbHMMMMMMMMMMMHo.\n"
        "     .o\"\" '         vodM*$&&HMMMMMMMMMM?.\n"
        "    ,'              $M&ood,~'`(&##MMMMMMH\\\n"
        "   /               ,MMMMMMM#b?#bobMMMMHMMML\n"
        "  &              ?MMMMMMMMMMMMMMMMM7MMM$R*Hk\n"
        " ?$.            :MMMMMMMMMMMMMMMMMMM/HMMM|`*L\n"
        "|               |MMMMMMMMMMMMMMMMMMMMbMH'   T,\n"
        "$H#:            `*MMMMMMMMMMMMMMMMMMMMb#}'  `?\n"
        "]MMH#             \"\"*\"\"\"\"*#MMMMMMMMMMMMM'    -\n"
        "MMMMMb_                   |MMMMMMMMMMMP'     :\n"
        "HMMMMMMMHo                 `MMMMMMMMMT       .\n"
        "?MMMMMMMMP                  9MMMMMMMM}       -\n"
        "-?MMMMMMM                  |MMMMMMMMM?,d-    '\n"
        " :|MMMMMM-                 `MMMMMMMT .M|.   :\n"
        "  .9MMM[                    &MMMMM*' `'    .\n"
        "   :9MMk                    `MMM#\"        -\n"
        "     &M}                     `          .-\n"
        "      `&.                             .\n"
        "        `~,   .                     ./\n"
        "            . _                  .-\n"
        "              '`--._,dd###pp=\"\"'/bar18");
    return cmd;
};

// Test 4 for prepareFilesystem(): Checks exception handling when attempting to
// create an invalid path name
TEST_F(ContainerHostFixture, TestPrepareFilesystem_InvalidPath)
{
    TemporaryDirectory staged_directory;
    auto cmd = createCommandInvalid2();
    OCIRunner::ContainerHost host(nullptr, cmd, DigestGenerator::hash(""),
                                  createCASClient(staged_directory),
                                  StageTreeRequest_StagingMode_DEFAULT);

    EXPECT_ANY_THROW(host.prepareFilesystem())
        << "prepareFilesystem(): Unexpectedly did not throw an exception when "
           "failing to create the path";
}

// Test for fetchDockerImage()
TEST_F(OCIRunnerFixture, TestFetchDockerImage)
{
    if (!std::getenv("RUN_INTERNET_TESTS")) {
        SUCCEED();
        std::cout << "Unable to connect to Docker registry; skipping "
                     "TestFetchDockerImage tests\n";
    }
    else {
        const std::string str1 = "hello-world:linux";
        EXPECT_NO_THROW(fetchDockerImage(str1))
            << "Unexpectedly threw an exception when attempting to fetch a "
               "valid "
               "Docker image";

        const std::string str2 = "alpine:3.17";
        EXPECT_NO_THROW(fetchDockerImage(str2))
            << "Unexpectedly threw an exception when attempting to fetch a "
               "valid "
               "Docker image";

        const std::string str3 = "invalid-image";
        EXPECT_ANY_THROW(fetchDockerImage(str3))
            << "Unexpectedly did not throw an exception when attempting to "
               "fetch "
               "an invalid Docker image";
    }
}

// Test for extractFileSystemFromImage()
TEST_F(OCIRunnerFixture, TestExtractFileSystemFromImage)
{

    if (!std::getenv("RUN_INTERNET_TESTS")) {
        SUCCEED();
        std::cout << "Unable to connect to Docker registry; skipping "
                     "TestExtractFileSystemFromImage tests\n";
    }
    else {
        // Testing extracting the tar to an invalid directory
        const std::string helloImage = "hello-world:linux";
        const std::string invalidPath = "invalid-path";
        EXPECT_ANY_THROW(extractFileSystemFromImage(helloImage, invalidPath))
            << "Unexpectedly did not threw an exception when attempting to "
               "export "
               "to an invalid file location";

        // Testing extracting the tar to a valid directory
        TemporaryDirectory tmpDir;
        const std::string pathName = std::string(tmpDir.name());
        EXPECT_NO_THROW(extractFileSystemFromImage(helloImage, pathName))
            << "Unexpectedly threw an exception when attempting to export a "
               "valid "
               "image to a tar file in the given location";

        const std::filesystem::path tarPath = pathName;
        EXPECT_TRUE(!is_empty(tarPath))
            << "Failed to extract tarball to the specified file location";

        const std::filesystem::path helloPath = pathName + "/hello";
        EXPECT_TRUE(exists(helloPath));

        // Testing extracting the file system of an invalid image
        const std::string invalidImage = "not-real-image-4831377";
        EXPECT_ANY_THROW(extractFileSystemFromImage(invalidImage, pathName));
    }
}
