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

#include <build/buildbox/local_execution.pb.h>

#include <buildboxcommon_logging.h>

#include <rumbad_bufferedpublisher.h>

namespace rumbad {

BufferedPublisher::BufferedPublisher(int batchSize,
                                     std::chrono::seconds publishInterval,
                                     Writer *writer)
    : batchSize(batchSize), publishInterval(publishInterval), writer(writer)
{

    // Start a thread to regularly write a batch of messages to the
    // backend database. This thread is joined when the BufferedPublisher
    // is destroyed.
    this->publishThread = std::thread([this] { periodicallyPublish(); });
}

BufferedPublisher::~BufferedPublisher()
{
    this->doPublish = false;
    this->waitCondition.notify_all();
    try {
        this->publishThread.join();
    }
    catch (std::system_error const &ex) {
        BUILDBOX_LOG_ERROR("Error joining publisher thread: " << ex.what());
    }

    // Attempt to make sure all the buffered messages get properly published.
    try {
        while (this->buffer.size() > 0) {
            this->publishBatch();
        }
    }
    catch (...) {
        BUILDBOX_LOG_ERROR("Error publishing buffered messages");
    }
}

void BufferedPublisher::periodicallyPublish()
{
    std::unique_lock<std::mutex> lock(this->waitLock);
    while (this->doPublish) {
        this->waitCondition.wait_for(lock, this->publishInterval);
        while (this->buffer.size() > 0) {
            try {
                this->publishBatch();
            }
            catch (...) {
                BUILDBOX_LOG_ERROR(
                    "Error publishing a batch of buffered messages");
            }
        }
    }
}

void BufferedPublisher::publish(const std::string &msg)
{
    SerializedMessage message;
    message.timestamp = std::chrono::system_clock::now();
    message.serializedData = msg;
    this->buffer.push(message);
}

void BufferedPublisher::publishBatch()
{
    int messageCount = 0;
    std::vector<Message> batch;
    while (messageCount < this->batchSize && this->buffer.size() > 0) {
        Message message;
        SerializedMessage messageData;
        bool popped = false;
        std::tie(messageData, popped) = this->buffer.pop();
        if (!popped) {
            // The buffer ran out of things, stop trying to get more
            break;
        }

        message.timestamp = messageData.timestamp;
        message.data.ParseFromString(messageData.serializedData);
        batch.push_back(message);
        messageCount++;
    }
    // If there was nothing to publish, break out early
    if (messageCount == 0) {
        return;
    }

    int attempts = 0;
    while (attempts < this->maxRetryAttempts) {
        try {
            this->writer->write(batch);
            break;
        }
        catch (std::exception &ex) {
            if (attempts++ < this->maxRetryAttempts) {
                BUILDBOX_LOG_DEBUG(
                    "Exception when writing batch, retrying: " << ex.what());
                std::this_thread::sleep_for(
                    std::chrono::seconds(1 << attempts));
            }
            else {
                BUILDBOX_LOG_WARNING(
                    "Exceeded retry limit when writing batch: " << ex.what());
                throw;
            }
        }
    }

    // Update caching stats
    this->statsManager.updateStatsInBatch(batch);
    if (this->statsManager.saveFilepathIsSet()) {
        this->statsManager.saveStatsToFile();
    }
}

} // namespace rumbad
