/*
 * 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_grpcclient.h>
#include <buildboxcommon_grpcerror.h>
#include <buildboxcommon_logstreamreader.h>
#include <buildboxcommon_protos.h>

#include <memory>

#include <build/bazel/remote/logstream/v1/remote_logstream_mock.grpc.pb.h>
#include <google/bytestream/bytestream_mock.grpc.pb.h>
#include <grpcpp/test/mock_stream.h>
#include <gtest/gtest.h>

using namespace buildboxcommon;
using namespace testing;

class LogStreamReaderTestFixture : public testing::Test {
  protected:
    typedef grpc::testing::MockClientReader<google::bytestream::ReadResponse>
        MockClientReader;

    const std::string TESTING_RESOURCE_NAME = "dummy-resource-name";
    const int GRPC_RETRY_LIMIT = 3;
    const int GRPC_RETRY_DELAY = 1;

    LogStreamReaderTestFixture()
        : byteStreamClient(
              std::make_shared<google::bytestream::MockByteStreamStub>()),
          client(std::make_shared<GrpcClient>()),
          logStreamReader(TESTING_RESOURCE_NAME, client),
          mockClientReader(new MockClientReader()),
          logStreamClient(std::make_unique<build::bazel::remote::logstream::
                                               v1::MockLogStreamServiceStub>())

    {
        logStreamReader.init(byteStreamClient);
    }

    std::shared_ptr<google::bytestream::MockByteStreamStub> byteStreamClient;
    std::shared_ptr<GrpcClient> client;
    buildboxcommon::LogStreamReader logStreamReader;

    MockClientReader *mockClientReader;
    // Wrapping this in a smart pointer causes issues on destruction.
    // (It is freed by someone else.)

    std::unique_ptr<
        build::bazel::remote::logstream::v1::MockLogStreamServiceStub>
        logStreamClient;
};

TEST_F(LogStreamReaderTestFixture, TestSuccessfulRead)
{
    const std::string data = "Hello!!";
    std::string receivedData = "";

    ReadResponse response;
    response.set_data(data);

    ReadRequest request;
    EXPECT_CALL(*byteStreamClient, ReadRaw(_, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(mockClientReader)));

    EXPECT_CALL(*mockClientReader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(response), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*mockClientReader, Finish())
        .WillOnce(Return(grpc::Status::OK));

    auto handler = [&](const std::string &chunk) { receivedData += chunk; };

    auto shouldStop = [&]() { return false; };

    logStreamReader.read(handler, shouldStop);
    EXPECT_TRUE(receivedData == data);
    EXPECT_EQ(request.resource_name(), TESTING_RESOURCE_NAME);
};

TEST_F(LogStreamReaderTestFixture, TestChunkedRead)
{
    const std::string chunk1 = "Hello ";
    const std::string chunk2 = "world!";
    std::string receivedData = "";

    ReadResponse response1;
    response1.set_data(chunk1);
    ReadResponse response2;
    response2.set_data(chunk2);

    ReadRequest request;
    EXPECT_CALL(*byteStreamClient, ReadRaw(_, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(mockClientReader)));

    EXPECT_CALL(*mockClientReader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(response1), Return(true)))
        .WillOnce(DoAll(SetArgPointee<0>(response2), Return(true)))
        .WillOnce(Return(false));

    auto handler = [&](const std::string &chunk) { receivedData += chunk; };

    auto shouldStop = [&]() { return false; };

    logStreamReader.read(handler, shouldStop);
    EXPECT_TRUE(receivedData == chunk1 + chunk2);
    EXPECT_EQ(request.resource_name(), TESTING_RESOURCE_NAME);
};

TEST_F(LogStreamReaderTestFixture, TestInterruptedRead)
{
    const std::string chunk1 = "Hello ";
    const std::string chunk2 = "world!";
    std::string receivedData = "";

    ReadResponse response1;
    response1.set_data(chunk1);
    ReadResponse response2;
    response2.set_data(chunk2);

    ReadRequest request;
    EXPECT_CALL(*byteStreamClient, ReadRaw(_, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(mockClientReader)));

    EXPECT_CALL(*mockClientReader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(response1), Return(true)));

    auto handler = [&](const std::string &chunk) { receivedData += chunk; };

    auto shouldStop = [&]() { return true; };

    logStreamReader.read(handler, shouldStop);
    EXPECT_TRUE(receivedData == chunk1);
    EXPECT_EQ(request.resource_name(), TESTING_RESOURCE_NAME);
};
