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

#include <chrono>
#include <grpcpp/support/status.h>
#include <memory>

#include <build/bazel/remote/logstream/v1/remote_logstream_mock.grpc.pb.h>
#include <gmock/gmock.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;

struct LogStreamWriterTestFixture : public testing::Test {

    typedef grpc::testing::MockClientWriter<google::bytestream::WriteRequest>
        MockClientWriter;

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

    LogStreamWriterTestFixture()
        : byteStreamClient(
              std::make_shared<google::bytestream::MockByteStreamStub>()),
          client(std::make_shared<GrpcClient>()),
          logStreamWriter(TESTING_RESOURCE_NAME, client),
          logStreamClient(std::make_unique<build::bazel::remote::logstream::
                                               v1::MockLogStreamServiceStub>())

    {
        logStreamWriter.init(byteStreamClient);
        client->setRetryLimit(GRPC_RETRY_LIMIT);
        client->setRequestTimeout(std::chrono::seconds(GRPC_REQUEST_TIMEOUT));
    }

    std::shared_ptr<google::bytestream::MockByteStreamStub> byteStreamClient;
    std::shared_ptr<GrpcClient> client;
    buildboxcommon::LogStreamWriter logStreamWriter;

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

TEST_F(LogStreamWriterTestFixture, TestSuccessfulWrite)
{
    auto mockClientWriter = std::make_unique<MockClientWriter>();

    ASSERT_TRUE(client != nullptr);
    const std::string data = "Hello!!";

    // Initial `QueryWriteStatus()` request on first call to `write()`:
    EXPECT_CALL(*byteStreamClient, QueryWriteStatus(_, _, _))
        .WillOnce(Return(grpc::Status::OK));

    WriteResponse response;
    response.set_committed_size(
        static_cast<google::protobuf::int64>(data.size()));

    WriteRequest request;
    EXPECT_CALL(*mockClientWriter, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request), Return(true)));
    EXPECT_CALL(*byteStreamClient, WriteRaw(_, _))
        .WillOnce(DoAll(SetArgPointee<1>(response),
                        Return(mockClientWriter.release())));

    EXPECT_TRUE(logStreamWriter.write(data));

    EXPECT_FALSE(request.finish_write());
    EXPECT_EQ(request.resource_name(), TESTING_RESOURCE_NAME);
}

TEST_F(LogStreamWriterTestFixture, TestWriteFailsWithUncommitedData)
{
    auto mockClientWriter = std::make_unique<MockClientWriter>();

    // Initial `QueryWriteStatus()` request on first call to `write()`:
    EXPECT_CALL(*byteStreamClient, QueryWriteStatus(_, _, _))
        .WillOnce(Return(grpc::Status::OK));

    WriteResponse response;
    response.set_committed_size(0);

    WriteRequest request;
    EXPECT_CALL(*byteStreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(response), Return(mockClientWriter.get())));

    EXPECT_CALL(*mockClientWriter, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request), Return(true)));

    EXPECT_TRUE(logStreamWriter.write("ABCD"));

    EXPECT_CALL(*mockClientWriter, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*mockClientWriter, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*mockClientWriter, Finish())
        .WillOnce(Return(grpc::Status::OK));

    auto _ = mockClientWriter.release();
    (void)_;
    EXPECT_FALSE(logStreamWriter.commit());
}

TEST_F(LogStreamWriterTestFixture, TestMultipleWritesAndCommit)
{
    auto mockClientWriter = std::make_unique<MockClientWriter>();

    const std::string data1 = "This is the first part...";
    const std::string data2 = "Second part.";

    // Initial `QueryWriteStatus()` request on first call to `write()`:
    EXPECT_CALL(*byteStreamClient, QueryWriteStatus(_, _, _))
        .WillOnce(Return(grpc::Status::OK));

    WriteResponse write_response;
    write_response.set_committed_size(
        static_cast<google::protobuf::int64>(data1.size()) +
        static_cast<google::protobuf::int64>(data2.size()));

    WriteRequest request1;
    EXPECT_CALL(*byteStreamClient, WriteRaw(_, _))
        .WillOnce(DoAll(SetArgPointee<1>(write_response),
                        Return(mockClientWriter.get())));

    EXPECT_CALL(*mockClientWriter, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request1), Return(true)));

    EXPECT_NO_THROW(logStreamWriter.write(data1));
    EXPECT_FALSE(request1.finish_write());
    EXPECT_EQ(request1.data(), data1);
    EXPECT_EQ(request1.write_offset(), 0);

    WriteRequest request2;
    EXPECT_CALL(*mockClientWriter, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request2), Return(true)));

    EXPECT_NO_THROW(logStreamWriter.write(data2));
    EXPECT_FALSE(request2.finish_write());
    EXPECT_EQ(request2.data(), data2);
    EXPECT_EQ(request2.write_offset(), data1.size());

    // Calling `commit()`:
    WriteRequest commit_request;
    EXPECT_CALL(*mockClientWriter, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&commit_request), Return(true)));
    EXPECT_CALL(*mockClientWriter, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*mockClientWriter, Finish())
        .WillOnce(Return(grpc::Status::OK));
    auto _ = mockClientWriter.release();
    (void)_;
    EXPECT_TRUE(logStreamWriter.commit());
    EXPECT_TRUE(commit_request.finish_write());
    EXPECT_EQ(commit_request.write_offset(), data1.size() + data2.size());
}

TEST_F(LogStreamWriterTestFixture, TestMultipleWritesRetries)
{
    auto mockClientWriter1 = std::make_unique<MockClientWriter>();

    const std::string data1 = "123";
    const std::string data2 = "456";

    QueryWriteStatusResponse queryResponse0, queryResponse1, queryResponse2;
    queryResponse1.set_committed_size(0);
    queryResponse2.set_committed_size(3);
    EXPECT_CALL(*byteStreamClient, QueryWriteStatus(_, _, _))
        .WillOnce(
            DoAll(SetArgPointee<2>(queryResponse1), Return(grpc::Status::OK)))
        .WillRepeatedly(
            DoAll(SetArgPointee<2>(queryResponse2), Return(grpc::Status::OK)));

    WriteResponse response;
    response.set_committed_size(
        static_cast<int64_t>(data1.size() + data2.size()));

    // The second writer used after retry
    auto mockClientWriter2 = std::make_unique<MockClientWriter>();

    EXPECT_CALL(*byteStreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(response), Return(mockClientWriter1.get())))
        .WillOnce(DoAll(SetArgPointee<1>(response),
                        Return(mockClientWriter2.get())));

    EXPECT_CALL(*mockClientWriter1, Write(_, _))
        .WillOnce(Return(true))
        .WillOnce(Return(false));
    EXPECT_CALL(*mockClientWriter1, Finish())
        .WillOnce(Return(
            grpc::Status(grpc::StatusCode::UNAVAILABLE, "unavailable")));

    EXPECT_CALL(*mockClientWriter2, Write(_, _)).WillRepeatedly(Return(true));
    EXPECT_CALL(*mockClientWriter2, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*mockClientWriter2, Finish())
        .WillOnce(Return(grpc::Status::OK));

    EXPECT_TRUE(logStreamWriter.write(data1));
    EXPECT_TRUE(logStreamWriter.write(data2));

    auto _1 = mockClientWriter1.release();
    (void)_1;
    auto _2 = mockClientWriter2.release();
    (void)_2;
    EXPECT_TRUE(logStreamWriter.commit());
}

TEST_F(LogStreamWriterTestFixture, TestFinishWrite)
{
    auto mockClientWriter = std::make_unique<MockClientWriter>();

    WriteRequest request;

    EXPECT_CALL(*mockClientWriter, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request), Return(true)));
    EXPECT_CALL(*mockClientWriter, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*mockClientWriter, Finish())
        .WillOnce(Return(grpc::Status::OK));

    EXPECT_CALL(*byteStreamClient, WriteRaw(_, _))
        .WillOnce(Return(mockClientWriter.get()));

    EXPECT_TRUE(logStreamWriter.commit());
    EXPECT_TRUE(request.finish_write());
    EXPECT_EQ(request.resource_name(), TESTING_RESOURCE_NAME);

    auto _ = mockClientWriter.release();
    (void)_;
}

TEST_F(LogStreamWriterTestFixture, TestOperationsAfterCommitThrow)
{
    auto mockClientWriter = std::make_unique<MockClientWriter>();

    EXPECT_CALL(*mockClientWriter, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*mockClientWriter, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*mockClientWriter, Finish())
        .WillOnce(Return(grpc::Status::OK));

    EXPECT_CALL(*byteStreamClient, WriteRaw(_, _))
        .WillOnce(Return(mockClientWriter.get()));

    EXPECT_TRUE(logStreamWriter.commit());
    EXPECT_THROW(logStreamWriter.write("More data"), std::runtime_error);
    auto _ = mockClientWriter.release();
    (void)_;
    EXPECT_THROW(logStreamWriter.commit(), std::runtime_error);
}

TEST_F(LogStreamWriterTestFixture, TestQueryWriteStatusReturnsNotFound)
{
    const std::string data = "Hello!!";
    // The `QueryWriteStatus()` request before performing a
    // `ByteStream.Write()` returns `NOT_FOUND`. This means we cannot write to
    // the stream.
    EXPECT_CALL(*byteStreamClient, QueryWriteStatus(_, _, _))
        .WillOnce(Return(grpc::Status(grpc::StatusCode::NOT_FOUND, "")));

    EXPECT_FALSE(logStreamWriter.write(data));
}

TEST_F(LogStreamWriterTestFixture, SuccessfulCreateLogStream)
{
    const std::string parent = "parent";
    const std::string readName = parent + "/foo";
    const std::string writeName = parent + "/foo/WRITE";

    LogStream response;
    response.set_name(readName);
    response.set_write_resource_name(writeName);

    CreateLogStreamRequest request;
    EXPECT_CALL(*logStreamClient, CreateLogStream(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), SetArgPointee<2>(response),
                        Return(grpc::Status::OK)));

    const LogStream returnedLogStream = LogStreamWriter::createLogStream(
        parent, GRPC_RETRY_LIMIT, GRPC_RETRY_DELAY, logStreamClient.get());

    // The request contains the parent value we specified:
    EXPECT_EQ(request.parent(), parent);

    // And the returned LogStream matches the one sent by the server:
    EXPECT_EQ(returnedLogStream.name(), response.name());
    EXPECT_EQ(returnedLogStream.write_resource_name(),
              response.write_resource_name());
}

TEST_F(LogStreamWriterTestFixture, CreateLogStreamReturnsError)
{
    const grpc::Status errorStatus = grpc::Status(
        grpc::StatusCode::UNAVAILABLE, "LogStream server is taking a nap.");

    EXPECT_CALL(*logStreamClient, CreateLogStream(_, _, _))
        .WillRepeatedly(Return(errorStatus));

    ASSERT_THROW(LogStreamWriter::createLogStream("parent", GRPC_RETRY_LIMIT,
                                                  GRPC_RETRY_DELAY,
                                                  logStreamClient.get()),
                 GrpcError);
}
