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

#define TESTING_RUMBAD_BUFFEREDPUBLISHER

#include <filesystem>
#include <fstream>

// component under test
#include <rumbad_bufferedpublisher.h>
#include <rumbad_compositewriter.h>
#include <rumbad_filewriter.h>
#include <rumbad_stdoutwriter.h>

// third party includes
#include <gmock/gmock.h>
#include <gtest/gtest.h>

namespace rumbad {
// NOTE: We're not using the test namespace here because we need access to
// the private members of BufferedPublisher.

using namespace ::testing;

const int TEST_BATCH_SIZE = 10;
const std::chrono::seconds TEST_PUBLISH_INTERVAL(1);

TEST(BufferedPublisherTest, TestBuffering)
{
    StdoutWriter writer;
    BufferedPublisher publisher(TEST_BATCH_SIZE, TEST_PUBLISH_INTERVAL,
                                &writer);
    publisher.doPublish = false; // Disable the publishing thread for this test

    ASSERT_EQ(publisher.buffer.size(), 0);

    build::buildbox::CompilerExecutionData compilationData;
    *compilationData.mutable_command() = "test command";
    publisher.publish(compilationData.SerializeAsString());

    ASSERT_EQ(publisher.buffer.size(), 1);
}

TEST(BufferedPublisherTest, TestBatching)
{
    const int TEST_MESSAGE_COUNT =
        15; // Ensure TEST_BATCH_SIZE < TEST_MESSAGE_SIZE < 2 * TEST_BATCH_SIZE

    StdoutWriter writer;
    BufferedPublisher publisher(TEST_BATCH_SIZE, TEST_PUBLISH_INTERVAL,
                                &writer);
    publisher.doPublish = false; // Disable the publishing thread for this test

    ASSERT_EQ(publisher.buffer.size(), 0);

    // Publish a test message a number of times
    build::buildbox::CompilerExecutionData compilationData;
    *compilationData.mutable_command() = "test command";
    for (int i = 0; i < TEST_MESSAGE_COUNT; i++) {
        publisher.publish(compilationData.SerializeAsString());
    }

    // Now trigger the final publishing of a batch and make sure the right
    // number of messages were removed from the buffer
    ASSERT_EQ(publisher.buffer.size(), TEST_MESSAGE_COUNT);
    publisher.publishBatch();
    ASSERT_EQ(publisher.buffer.size(), TEST_MESSAGE_COUNT - TEST_BATCH_SIZE);

    // Trigger another batch publishing; this time without a full batch and
    // check that the buffer was fully drained
    publisher.publishBatch();
    ASSERT_EQ(publisher.buffer.size(), 0);
}

TEST(BufferedPublisherTest, TestBatchingCompositeWriter)
{
    std::string logFilepath = "rumbad.log";
    std::filesystem::remove(logFilepath);
    CompositeWriter writer;
    writer.addWriter(std::make_unique<StdoutWriter>());
    writer.addWriter(std::make_unique<FileWriter>(logFilepath));

    const int TEST_MESSAGE_COUNT =
        15; // Ensure TEST_BATCH_SIZE < TEST_MESSAGE_SIZE < 2 * TEST_BATCH_SIZE

    BufferedPublisher publisher(TEST_BATCH_SIZE, TEST_PUBLISH_INTERVAL,
                                &writer);
    publisher.doPublish = false; // Disable the publishing thread for this test

    ASSERT_EQ(publisher.buffer.size(), 0);

    // Publish a test message a number of times
    build::buildbox::CompilerExecutionData compilationData;
    *compilationData.mutable_command() = "test command";
    for (int i = 0; i < TEST_MESSAGE_COUNT; i++) {
        publisher.publish(compilationData.SerializeAsString());
    }

    // Now trigger the final publishing of a batch and make sure the right
    // number of messages were removed from the buffer
    ASSERT_EQ(publisher.buffer.size(), TEST_MESSAGE_COUNT);
    publisher.publishBatch();
    ASSERT_EQ(publisher.buffer.size(), TEST_MESSAGE_COUNT - TEST_BATCH_SIZE);

    // Trigger another batch publishing; this time without a full batch and
    // check that the buffer was fully drained
    publisher.publishBatch();
    ASSERT_EQ(publisher.buffer.size(), 0);

    // Inspect logfile
    int numberOfLogLines = 0;
    std::string line;
    std::ifstream f(logFilepath);
    while (getline(f, line)) {
        numberOfLogLines++;
    }
    f.close();
    ASSERT_EQ(numberOfLogLines, TEST_MESSAGE_COUNT);
    std::filesystem::remove(logFilepath);
}

TEST(BufferedPublisherTest, TestStdoutStats)
{
    // Ensure TEST_BATCH_SIZE > TEST_MESSAGE_SIZE
    const int TEST_MESSAGE_COUNT = 3;

    StdoutWriter writer;
    BufferedPublisher publisher(TEST_BATCH_SIZE, TEST_PUBLISH_INTERVAL,
                                &writer);
    publisher.doPublish = false; // Disable the publishing thread for this test
    StatsManager *statsManager = &publisher.statsManager;
    StatsData const *data = statsManager->readonlyStatsdata();

    ASSERT_EQ(publisher.buffer.size(), 0);

    // Setup a command message
    build::buildbox::CompilerExecutionData compilationData;
    *compilationData.mutable_command() = "test command";
    auto reccData = compilationData.mutable_recc_data();
    (*reccData->mutable_counter_metrics())[COUNTER_NAME_ACTION_CACHE_HIT] = 1;
    ASSERT_EQ(compilationData.has_recc_data(), true);

    /*
     *  Testing for: Compilation command cache hits + cache misses
     */
    // Given: creation of a cache hit message.
    ASSERT_EQ(reccData->counter_metrics().at(COUNTER_NAME_ACTION_CACHE_HIT),
              1);
    // When: multiple cache hit messages are added to the publisher queue
    for (int i = 0; i < TEST_MESSAGE_COUNT; i++) {
        publisher.publish(compilationData.SerializeAsString());
    }
    reccData->mutable_counter_metrics()->erase(COUNTER_NAME_ACTION_CACHE_HIT);

    // And When: same number of cache miss messages are added right after
    (*reccData->mutable_counter_metrics())[COUNTER_NAME_ACTION_CACHE_MISS] = 1;
    ASSERT_EQ(reccData->counter_metrics().at(COUNTER_NAME_ACTION_CACHE_MISS),
              1);
    for (int i = 0; i < TEST_MESSAGE_COUNT; i++) {
        publisher.publish(compilationData.SerializeAsString());
    }
    reccData->mutable_counter_metrics()->erase(COUNTER_NAME_ACTION_CACHE_MISS);

    // Then: the publish queue size is 0 upon triggering batch publish
    ASSERT_EQ(publisher.buffer.size(), 2 * TEST_MESSAGE_COUNT);
    publisher.publishBatch();
    ASSERT_EQ(publisher.buffer.size(), 0);

    // And: the cache hits and cache misses are available
    ASSERT_EQ(data->d_compile_command_counts, 2 * TEST_MESSAGE_COUNT);
    ASSERT_EQ(data->d_compile_command_cache_hit_counts, TEST_MESSAGE_COUNT);
    ASSERT_EQ(statsManager->compileCacheHitRate(),
              (float)(TEST_MESSAGE_COUNT) / (2 * TEST_MESSAGE_COUNT));

    /*
     *  Testing for: Linking command cache hit
     */
    // Given: a link command cache hit message is added to publisher queue
    (*reccData
          ->mutable_counter_metrics())[COUNTER_NAME_LINK_ACTION_CACHE_HIT] = 1;
    publisher.publish(compilationData.SerializeAsString());
    reccData->mutable_counter_metrics()->erase(
        COUNTER_NAME_LINK_ACTION_CACHE_HIT);

    // When: the queue content is published in batch
    publisher.publishBatch();

    // Then: the link cache hits is available
    ASSERT_EQ(data->d_linking_command_counts, 1);
    ASSERT_EQ(data->d_linking_command_cache_hit_counts, 1);
    ASSERT_EQ(statsManager->linkCacheHitRate(), 1.0);

    /*
     *  Testing for: Compilation command cache skip
     */
    // Given: A compilation cache skip message is added to publisher queue
    (*reccData->mutable_counter_metrics())[COUNTER_NAME_ACTION_CACHE_SKIP] = 1;
    publisher.publish(compilationData.SerializeAsString());
    reccData->mutable_counter_metrics()->erase(COUNTER_NAME_ACTION_CACHE_SKIP);

    // When: the queue content is published in batch
    publisher.publishBatch();

    // Then: the command cache skip is available
    ASSERT_EQ(data->d_compile_command_cache_skip_counts, 1);
}

} // namespace rumbad
