/*
 * Copyright 2020 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_fslocalactionstorage.h>
#include <buildboxcasd_server.h>
#include <buildboxcommon_fslocalcas.h>

#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_temporarydirectory.h>

#include <benchmark/benchmark.h>
#include <cstdint>
#include <fstream>
#include <grpcpp/create_channel.h>
#include <grpcpp/impl/codegen/call.h>
#include <random>
#include <sys/resource.h>

using namespace buildboxcasd;
using namespace buildboxcommon;

class LocalFixture : public benchmark::Fixture {
  protected:
    LocalFixture()
        : TEST_CASD_SERVER_ADDRESS(
              "unix://" + std::string(local_storage_root_directory.name()) +
              "/casd.sock"),
          local_storage(std::make_shared<FsLocalCas>(
              local_storage_root_directory.name())),
          local_asset_storage(std::make_shared<FsLocalAssetStorage>(
              local_storage_root_directory.name())),
          local_action_storage(std::make_shared<FsLocalActionStorage>(
              local_storage_root_directory.name())),
          casd_channel(grpc::CreateChannel(
              TEST_CASD_SERVER_ADDRESS, grpc::InsecureChannelCredentials())),
          localcas_stub(LocalContentAddressableStorage::NewStub(casd_channel))
    {
        const auto instance_name = "";
        local_server = std::make_shared<Server>(
            local_storage, local_asset_storage, local_action_storage);

        // Building and starting the casd server:
        local_server->addLocalServerInstance(instance_name);
        local_server->addListeningPort(TEST_CASD_SERVER_ADDRESS);
        local_server->start();
    }

    void clearCache()
    {
        local_storage->listBlobs(
            [this](const Digest &digest, blob_time_type time) {
                local_storage->deleteBlob(digest);
            });
    }

    static void generateTestFiles(const std::string &path, int64_t offset,
                                  int64_t count)
    {
        for (int i = 0; i < count; i++) {
            std::ofstream file(path + "/" + std::to_string(i),
                               std::ofstream::binary);
            // Write unique content
            file << std::to_string(offset + i);
            file.close();
        }
    }

    static void generateTestDirectory(const std::string &path, int64_t offset,
                                      int64_t fileCount)
    {
        if (fileCount <= 256) {
            generateTestFiles(path, offset, fileCount);
        }
        else {
            // Create 8 directories per level
            const int directoryCount = 8;
            const int64_t filesPerDirectory = fileCount / directoryCount;
            for (int i = 0; i < directoryCount; i++) {
                const std::string directoryPath =
                    path + "/" + std::to_string(i);
                FileUtils::createDirectory(directoryPath.c_str());
                generateTestDirectory(directoryPath, offset,
                                      filesPerDirectory);
                offset += filesPerDirectory;
            }
        }
    }

    Digest captureTree(const std::string &path)
    {
        grpc::ClientContext client_context;
        CaptureTreeRequest request;
        CaptureTreeResponse response;

        request.add_path(path);
        const auto status =
            localcas_stub->CaptureTree(&client_context, request, &response);
        assert(status.ok());
        assert(response.responses().size() == 1);
        auto resp = response.responses()[0];

        Tree tree;
        const auto tree_serialized =
            local_storage->readBlob(resp.tree_digest());
        assert(tree_serialized != nullptr);
        const bool ok = tree.ParseFromString(*tree_serialized);
        assert(ok);
        const auto root_directory_serialized = tree.root().SerializeAsString();
        return buildboxcommon::DigestGenerator::hash(
            root_directory_serialized);
    }

    TemporaryDirectory local_storage_root_directory;

    const std::string TEST_CASD_SERVER_ADDRESS;

    std::shared_ptr<FsLocalCas> local_storage;
    std::shared_ptr<FsLocalAssetStorage> local_asset_storage;
    std::shared_ptr<FsLocalActionStorage> local_action_storage;
    std::shared_ptr<Server> local_server;

    // gRPC stubs (allow to invoke gRPC calls):
    std::shared_ptr<grpc::Channel> casd_channel;
    std::unique_ptr<LocalContentAddressableStorage::Stub> localcas_stub;
};

BENCHMARK_DEFINE_F(LocalFixture, FetchTree)(benchmark::State &state)
{
    auto temp_dir = local_storage->createStagingDirectory();
    generateTestDirectory(temp_dir.strname(), 0, state.range(0));
    const auto digest = captureTree(temp_dir.strname());

    for (auto _ : state) {
        (void)_;
        grpc::ClientContext client_context;
        FetchTreeRequest request;
        FetchTreeResponse response;

        request.mutable_root_digest()->CopyFrom(digest);
        request.set_fetch_file_blobs(true);
        const auto status =
            localcas_stub->FetchTree(&client_context, request, &response);
        assert(status.ok());
    }
}

BENCHMARK_DEFINE_F(LocalFixture, StageTree)(benchmark::State &state)
{
    auto temp_dir = local_storage->createStagingDirectory();
    generateTestDirectory(temp_dir.strname(), 0, state.range(0));

    for (auto _ : state) {
        (void)_;
        state.PauseTiming();
        clearCache();
        const auto digest = captureTree(temp_dir.strname());
        state.ResumeTiming();

        grpc::ClientContext stage_client_context;
        StageTreeRequest stage_request;
        StageTreeResponse stage_response;

        stage_request.mutable_root_digest()->CopyFrom(digest);
        auto reader_writer = localcas_stub->StageTree(&stage_client_context);
        bool ok = reader_writer->Write(stage_request);
        ok &= reader_writer->Read(&stage_response);

        state.PauseTiming();
        if (ok) {
            // Ask the server to clean up:
            ok &= reader_writer->Write(StageTreeRequest());

            // Receive its last empty reply:
            ok &= reader_writer->Read(&stage_response);

            ok &= reader_writer->Finish().ok();
        }
        state.ResumeTiming();

        assert(ok);
    }
}

BENCHMARK_DEFINE_F(LocalFixture, CaptureTree)(benchmark::State &state)
{
    auto empty_dir = local_storage->createStagingDirectory();

    for (auto _ : state) {
        (void)_;
        state.PauseTiming();
        clearCache();
        const auto digest = captureTree(empty_dir.strname());

        grpc::ClientContext stage_client_context;
        StageTreeRequest stage_request;
        StageTreeResponse stage_response;

        stage_request.mutable_root_digest()->CopyFrom(digest);
        auto reader_writer = localcas_stub->StageTree(&stage_client_context);
        bool ok = reader_writer->Write(stage_request);
        ok &= reader_writer->Read(&stage_response);

        if (ok) {
            const auto staging_path = stage_response.path();
            const auto output_path = staging_path + "/subdir";
            FileUtils::createDirectory(output_path.c_str());
            generateTestDirectory(output_path, 0, state.range(0));
            state.ResumeTiming();

            captureTree(output_path);

            // Ask the server to clean up:
            ok &= reader_writer->Write(StageTreeRequest());

            // Receive its last empty reply:
            ok &= reader_writer->Read(&stage_response);
            ok &= reader_writer->Finish().ok();
        }
        assert(ok);
    }
}

int main(int argc, char **argv)
{
    // Increase open file limit to avoid running out of file descriptors
    struct rlimit rlim = {};
    getrlimit(RLIMIT_NOFILE, &rlim);
    rlim.rlim_cur = rlim.rlim_max;
    setrlimit(RLIMIT_NOFILE, &rlim);

    // Don't use BENCHMARK_REGISTER_F macro outside main() to
    // avoid static initialization order fiasco.
    // https://github.com/google/benchmark/issues/498

    // Range: Number of files to fetch
    BENCHMARK_REGISTER_F(LocalFixture, FetchTree)
        ->Range(64, 32768)
        ->Unit(benchmark::kMillisecond)
        ->UseRealTime();

    // Range: Number of files to stage
    BENCHMARK_REGISTER_F(LocalFixture, StageTree)
        ->Ranges({{4096, 262144}})
        ->Unit(benchmark::kMillisecond)
        ->Iterations(5);

    // Range: Number of files to capture
    BENCHMARK_REGISTER_F(LocalFixture, CaptureTree)
        ->Ranges({{64, 32768}})
        ->Unit(benchmark::kMillisecond)
        ->Iterations(5);

    benchmark::Initialize(&argc, argv);
    if (benchmark::ReportUnrecognizedArguments(argc, argv)) {
        return 1;
    }
    benchmark::RunSpecifiedBenchmarks();
}
