#include <gtest/gtest.h>

#include <google/devtools/remoteworkers/v1test2/bots_mock.grpc.pb.h>
#include <google/protobuf/util/message_differencer.h>

#include <buildboxworker_botsessionutils.h>

#include <buildboxcommon_grpctestserver.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommonmetrics_durationmetricvalue.h>
#include <buildboxcommonmetrics_testingutils.h>
#include <buildboxworker_metricnames.h>
#include <buildboxworker_worker.h>

using namespace testing;
using namespace buildboxworker::proto;
using buildboxcommon::buildboxcommonmetrics::collectedByName;
using buildboxcommon::buildboxcommonmetrics::DurationMetricValue;
using buildboxworker::BotSessionUtils;

using google::devtools::remoteworkers::v1test2::MockBotsStub;

class BotSessionUtilsTests : public ::testing::Test {
  protected:
    buildboxcommon::GrpcTestServer testServer;
    BotStatus botStatus = BotStatus::OK;
    std::shared_ptr<Bots::StubInterface> stub;
    BotSession botSession;
    buildboxcommon::ConnectionOptions botsServerConnection;
    buildboxcommon::GrpcRetrier::GrpcStatusCodes defaultRetryStatusCodes = {
        grpc::StatusCode::UNAVAILABLE};
    std::atomic_bool cancelled{false};
    buildboxcommon::RequestMetadataGenerator requestMetadata;

    std::mutex sessionMutex;

    UpdateBotSessionRequest expectedUpdateBotSessionRequest;
    BotSession expectedBotSession;
    BotSessionUtilsTests() : stub(Bots::NewStub(testServer.channel()))
    {
        // Setup default UpdateBotSessionRequest
        expectedBotSession.set_status(botStatus);
        *expectedUpdateBotSessionRequest.mutable_bot_session() =
            expectedBotSession;
        requestMetadata.set_tool_details("buildbox-worker-tests", "0.0.1");
    }
    void validateRequestMetadata(buildboxcommon::GrpcTestServerContext &ctx)
    {
        build::bazel::remote::execution::v2::RequestMetadata clientMetadata =
            buildboxcommon::RequestMetadataGenerator::parse_request_metadata(
                ctx.d_ctx.client_metadata());
        EXPECT_EQ(clientMetadata.tool_details().tool_name(),
                  requestMetadata.getToolDetails().tool_name());
        EXPECT_EQ(clientMetadata.tool_details().tool_version(),
                  requestMetadata.getToolDetails().tool_version());
    }
};

/*
 * UpdateBotSession succeeds on the first try.
 */
TEST_F(BotSessionUtilsTests, UpdateBotSessionHappyPath)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/UpdateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedUpdateBotSessionRequest);
        ctx.writeAndFinish(botSession);
    });

    buildboxworker::RunnerPartialExecutionMetadata partialExecutionMetadata =
        {};
    std::unique_lock<std::mutex> lock(sessionMutex);
    const grpc::Status status = BotSessionUtils::updateBotSession(
        botStatus, stub, botsServerConnection, &botSession,
        defaultRetryStatusCodes, cancelled, requestMetadata,
        partialExecutionMetadata, &lock);
    serverHandler.join();
    EXPECT_TRUE(status.ok());
}

/*
 * UpdateBotSession succeeds after one retry.
 */
TEST_F(BotSessionUtilsTests, UpdateBotSessionEventualSuccess)
{
    botsServerConnection.d_retryLimit = "1";
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx1(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/UpdateBotSession");
        validateRequestMetadata(ctx1);
        ctx1.read(expectedUpdateBotSessionRequest);
        ctx1.finish(grpc::Status(grpc::UNAVAILABLE, "Failing for test"));

        buildboxcommon::GrpcTestServerContext ctx2(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/UpdateBotSession");
        validateRequestMetadata(ctx2);
        ctx2.read(expectedUpdateBotSessionRequest);
        ctx2.writeAndFinish(botSession);
    });

    buildboxworker::RunnerPartialExecutionMetadata partialExecutionMetadata =
        {};
    std::unique_lock<std::mutex> lock(sessionMutex);
    const grpc::Status status = BotSessionUtils::updateBotSession(
        botStatus, stub, botsServerConnection, &botSession,
        defaultRetryStatusCodes, cancelled, requestMetadata,
        partialExecutionMetadata, &lock);
    serverHandler.join();
    EXPECT_TRUE(status.ok());
}

/*
 * UpdateBotSession returns an the error status when failing
 */
TEST_F(BotSessionUtilsTests, UpdateBotSessionReturnsError)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/UpdateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedUpdateBotSessionRequest);
        ctx.finish(
            grpc::Status(grpc::INTERNAL, "Non-retriable error for test"));
    });
    const grpc::Status expectedError =
        grpc::Status(grpc::INTERNAL, "Non-retriable error for test");

    buildboxworker::RunnerPartialExecutionMetadata partialExecutionMetadata =
        {};
    std::unique_lock<std::mutex> lock(sessionMutex);
    const grpc::Status status = BotSessionUtils::updateBotSession(
        botStatus, stub, botsServerConnection, &botSession,
        defaultRetryStatusCodes, cancelled, requestMetadata,
        partialExecutionMetadata, &lock);
    serverHandler.join();
    EXPECT_EQ(status.error_code(), expectedError.error_code());
    EXPECT_EQ(status.error_message(), expectedError.error_message());
}

/*
 * UpdateBotSession is cancelled.
 */
TEST_F(BotSessionUtilsTests, UpdateBotSessionCancelled)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/UpdateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedUpdateBotSessionRequest);
        cancelled = true;
        std::this_thread::sleep_for(std::chrono::seconds(1));
        ctx.expectCancel();
    });

    buildboxworker::RunnerPartialExecutionMetadata partialExecutionMetadata =
        {};
    std::unique_lock<std::mutex> lock(sessionMutex);
    const grpc::Status status = BotSessionUtils::updateBotSession(
        botStatus, stub, botsServerConnection, &botSession,
        defaultRetryStatusCodes, cancelled, requestMetadata,
        partialExecutionMetadata, &lock);
    serverHandler.join();
    EXPECT_EQ(status.error_code(), grpc::StatusCode::CANCELLED);
}
/*
 * UpdateBotSession being cancelled before request sent.
 */
TEST_F(BotSessionUtilsTests, UpdateBotSessionCancelledAfterRequest)
{
    // No request being sent means no serverHandler, so instead
    // just inspect the error code/error message to verify
    cancelled = true;
    buildboxworker::RunnerPartialExecutionMetadata partialExecutionMetadata =
        {};
    std::unique_lock<std::mutex> lock(sessionMutex);
    const grpc::Status status = BotSessionUtils::updateBotSession(
        botStatus, stub, botsServerConnection, &botSession,
        defaultRetryStatusCodes, cancelled, requestMetadata,
        partialExecutionMetadata, &lock);
    EXPECT_EQ(status.error_code(), grpc::StatusCode::CANCELLED);
    EXPECT_EQ(status.error_message(),
              "Cancellation requested before sending request");
}
/*
 * Verify that metrics are collected for UpdateBotSession.
 */
TEST_F(BotSessionUtilsTests, UpdateBotSessionMetricsCollected)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/UpdateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedUpdateBotSessionRequest);
        ctx.writeAndFinish(botSession);
    });

    buildboxworker::RunnerPartialExecutionMetadata partialExecutionMetadata =
        {};
    std::unique_lock<std::mutex> lock(sessionMutex);
    BotSessionUtils::updateBotSession(botStatus, stub, botsServerConnection,
                                      &botSession, defaultRetryStatusCodes,
                                      cancelled, requestMetadata,
                                      partialExecutionMetadata, &lock);
    serverHandler.join();
    ASSERT_TRUE(collectedByName<DurationMetricValue>(
        buildboxworker::MetricNames::TIMER_NAME_UPDATE_BOTSESSION));
}

/*
 * Verify that executeOperationMetadata has no values if no metadata returned
 */
TEST_F(BotSessionUtilsTests, UpdateBotSessionEmptyExecuteOperationMetadata)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/UpdateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedUpdateBotSessionRequest);
        ctx.writeAndFinish(botSession);
    });

    BotSessionUtils::ExecuteOperationMetadataEntriesMultiMap multiMap;
    buildboxworker::RunnerPartialExecutionMetadata partialExecutionMetadata =
        {};
    std::unique_lock<std::mutex> lock(sessionMutex);
    BotSessionUtils::updateBotSession(
        botStatus, stub, botsServerConnection, &botSession,
        defaultRetryStatusCodes, cancelled, requestMetadata,
        partialExecutionMetadata, &lock, &multiMap);
    serverHandler.join();

    ASSERT_EQ(0, multiMap.size());
}
/*
 * Verify that ExecuteOperationMetadata attached to UpdateBotSession calls is
 * read
 */
TEST_F(BotSessionUtilsTests, UpdateBotSessionReadsExecuteOperationMetadata)
{
    buildboxcommon::Digest actionDigest;
    actionDigest.set_hash("hash");
    actionDigest.set_size_bytes(123);

    const auto streamResourceName = "stream/outputs/here";
    buildboxcommon::ExecuteOperationMetadata metadata;
    metadata.set_stdout_stream_name(streamResourceName);
    metadata.mutable_action_digest()->CopyFrom(actionDigest);

    const std::string serializedMetadata = metadata.SerializeAsString();

    const auto metadataName = "executeoperationmetadata-bin";

    std::thread serverHandler([this, serializedMetadata, metadataName]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/UpdateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedUpdateBotSessionRequest);
        ctx.d_ctx.AddTrailingMetadata(metadataName, serializedMetadata);
        ctx.writeAndFinish(botSession);
    });

    BotSessionUtils::ExecuteOperationMetadataEntriesMultiMap multiMap;
    buildboxworker::RunnerPartialExecutionMetadata partialExecutionMetadata =
        {};
    std::unique_lock<std::mutex> lock(sessionMutex);
    BotSessionUtils::updateBotSession(
        botStatus, stub, botsServerConnection, &botSession,
        defaultRetryStatusCodes, cancelled, requestMetadata,
        partialExecutionMetadata, &lock, &multiMap);
    serverHandler.join();

    ASSERT_EQ(multiMap.size(), 1);
    // Verify the multiMap has an entry for actionDigest whose value matches
    // the expected metadata
    if (auto element = multiMap.find(actionDigest);
        element != multiMap.end()) {
        ASSERT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            element->second, metadata));
    }
    else {
        ASSERT_FALSE("Metadata not returned by updateBotSession");
    }
}

class CreateBotSessionTests : public BotSessionUtilsTests {
  protected:
    CreateBotSessionRequest expectedCreateBotSessionRequest;

    CreateBotSessionTests()
    {
        botSession.set_bot_id("foo");
        botsServerConnection.d_instanceName = "dev";

        // Setup default expected CreateBotSessionRequest
        expectedBotSession.set_bot_id("foo");
        expectedBotSession.clear_status();
        expectedCreateBotSessionRequest.set_parent("dev");
        *expectedCreateBotSessionRequest.mutable_bot_session() =
            expectedBotSession;
    }
};

TEST_F(CreateBotSessionTests, CreateBotSessionHappyPath)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/CreateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedCreateBotSessionRequest);
        ctx.writeAndFinish(botSession);
    });

    const grpc::Status status = BotSessionUtils::createBotSession(
        stub, botsServerConnection, &botSession, defaultRetryStatusCodes,
        requestMetadata, cancelled);
    serverHandler.join();
    EXPECT_TRUE(status.ok());
}

/*
 * CreateBotSession succeeds after one retry.
 */
TEST_F(CreateBotSessionTests, CreateBotSessionEventualSuccess)
{
    botsServerConnection.d_retryLimit = "1";
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx1(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/CreateBotSession");
        validateRequestMetadata(ctx1);
        ctx1.read(expectedCreateBotSessionRequest);
        ctx1.finish(grpc::Status(grpc::UNAVAILABLE, "Failing for test"));

        buildboxcommon::GrpcTestServerContext ctx2(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/CreateBotSession");
        validateRequestMetadata(ctx2);
        ctx2.read(expectedCreateBotSessionRequest);
        ctx2.writeAndFinish(botSession);
    });
    const grpc::Status status = BotSessionUtils::createBotSession(
        stub, botsServerConnection, &botSession, defaultRetryStatusCodes,
        requestMetadata, cancelled);
    serverHandler.join();
    EXPECT_TRUE(status.ok());
}

TEST_F(CreateBotSessionTests, CreateBotSessionReturnsError)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/CreateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedCreateBotSessionRequest);
        ctx.finish(
            grpc::Status(grpc::INTERNAL, "Non-retriable error for test"));
    });
    const grpc::Status expectedError =
        grpc::Status(grpc::INTERNAL, "Non-retriable error for test");
    const grpc::Status status = BotSessionUtils::createBotSession(
        stub, botsServerConnection, &botSession, defaultRetryStatusCodes,
        requestMetadata, cancelled);
    serverHandler.join();
    EXPECT_FALSE(status.ok());
    EXPECT_EQ(status.error_code(), expectedError.error_code());
    EXPECT_EQ(status.error_message(), expectedError.error_message());
}

/*
 * CreateBotSession is cancelled.
 */
TEST_F(CreateBotSessionTests, CreateBotSessionCancelled)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/CreateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedCreateBotSessionRequest);
        cancelled = true;
        std::this_thread::sleep_for(std::chrono::seconds(1));
        ctx.expectCancel();
    });

    const grpc::Status status = BotSessionUtils::createBotSession(
        stub, botsServerConnection, &botSession, defaultRetryStatusCodes,
        requestMetadata, cancelled);
    serverHandler.join();
    EXPECT_EQ(status.error_code(), grpc::StatusCode::CANCELLED);
}

/*
 * UpdateBotSession being cancelled before request sent.
 */
TEST_F(CreateBotSessionTests, CreateBotSessionCancelledAfterRequest)
{
    // No request being sent means no serverHandler, so instead
    // just inspect the error code/error message to verify
    cancelled = true;
    const grpc::Status status = BotSessionUtils::createBotSession(
        stub, botsServerConnection, &botSession, defaultRetryStatusCodes,
        requestMetadata, cancelled);
    EXPECT_EQ(status.error_code(), grpc::StatusCode::CANCELLED);
    EXPECT_EQ(status.error_message(),
              "Cancellation requested before sending request");
}

TEST_F(CreateBotSessionTests, CreateBotSessionMetricsCollected)
{
    std::thread serverHandler([this]() {
        buildboxcommon::GrpcTestServerContext ctx(
            &testServer,
            "/google.devtools.remoteworkers.v1test2.Bots/CreateBotSession");
        validateRequestMetadata(ctx);
        ctx.read(expectedCreateBotSessionRequest);
        ctx.writeAndFinish(botSession);
    });
    BotSessionUtils::createBotSession(stub, botsServerConnection, &botSession,
                                      defaultRetryStatusCodes, requestMetadata,
                                      cancelled);
    serverHandler.join();
    ASSERT_TRUE(collectedByName<DurationMetricValue>(
        buildboxworker::MetricNames::TIMER_NAME_CREATE_BOTSESSION));
}

TEST(BotSessionUtilsTest, ExtractValidExecuteOperationMetadata)
{
    buildboxcommon::Digest actionDigest;
    actionDigest.set_hash("hash");
    actionDigest.set_size_bytes(123);

    const auto streamResourceName = "stream/outputs/here";
    buildboxcommon::ExecuteOperationMetadata metadata;
    metadata.set_stdout_stream_name(streamResourceName);
    metadata.mutable_action_digest()->CopyFrom(actionDigest);

    const std::string serializedMetadata = metadata.SerializeAsString();

    const auto metadataName = "executeoperationmetadata-bin";
    const std::multimap<grpc::string_ref, grpc::string_ref> metadataMultiMap =
        {{metadataName, serializedMetadata}, {"val", "123"}};

    const BotSessionUtils::ExecuteOperationMetadataEntriesMultiMap
        extractedMetadata =
            BotSessionUtils::extractExecuteOperationMetadata(metadataMultiMap);

    ASSERT_EQ(extractedMetadata.size(), 1);
    const auto extractedEntry = extractedMetadata.find(actionDigest)->second;
    ASSERT_EQ(extractedEntry.stdout_stream_name(),
              metadata.stdout_stream_name());
    ASSERT_EQ(extractedEntry.action_digest(), actionDigest);
}

TEST(BotSessionUtilsTest, ExtractEmptyExecuteOperationMetadata)
{
    const std::multimap<grpc::string_ref, grpc::string_ref> metadataMultiMap =
        {{"foo", "bar"}, {"a", "b"}};

    ASSERT_TRUE(
        BotSessionUtils::extractExecuteOperationMetadata(metadataMultiMap)
            .empty());
}
