/*
 * Copyright 2024 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 <buildboxcommon_logstreamreader.h>

#include <utility>

namespace buildboxcommon {

LogStreamReader::LogStreamReader(std::string resourceName,
                                 const ConnectionOptions &connectionOptions)
    : d_resourceName(std::move(resourceName)),
      d_grpcClient(std::make_shared<GrpcClient>())
{
    d_grpcClient->init(connectionOptions);
    d_byteStreamClient =
        google::bytestream::ByteStream::NewStub(d_grpcClient->channel());
}

LogStreamReader::LogStreamReader(std::string resourceName,
                                 const std::shared_ptr<GrpcClient> &client)
    : d_resourceName(std::move(resourceName)), d_grpcClient(client)
{
    if (d_grpcClient->channel() != nullptr) {
        d_byteStreamClient =
            google::bytestream::ByteStream::NewStub(d_grpcClient->channel());
    }
}

// Connect to another client
void LogStreamReader::init(
    std::shared_ptr<google::bytestream::ByteStream::StubInterface>
        bytestreamClient)
{
    d_byteStreamClient = std::move(bytestreamClient);
}

bool LogStreamReader::read(std::function<void(const std::string &)> onMessage,
                           std::function<bool()> shouldStop)
{
    BUILDBOX_LOG_TRACE("Reading LogStream " << d_resourceName);
    size_t currentReadOffset = 0;

    auto readLambda = [&](grpc::ClientContext &context) {
        ReadRequest request;
        request.set_resource_name(d_resourceName);
        request.set_read_offset(static_cast<int64_t>(currentReadOffset));
        auto reader = this->d_byteStreamClient->Read(&context, request);

        ReadResponse response;
        while (reader->Read(&response)) {
            onMessage(response.data());
            currentReadOffset += response.data().length();
            if (shouldStop()) {
                return grpc::Status::CANCELLED;
            }
        }

        grpc::Status readStatus = reader->Finish();
        return readStatus;
    };

    try {
        d_grpcClient->issueRequest(readLambda, "(LogStream) ByteStream.Read",
                                   nullptr);
        return true;
    }
    catch (GrpcError &) {
        return false;
    }
}

bool LogStreamReader::read()
{
    auto handler = [&](const std::string &data) { std::cout << data; };
    auto shouldStop = [&]() { return false; };
    return read(handler, shouldStop);
}

} // namespace buildboxcommon
