/*
 * Copyright 2023 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 <buildboxcasd_executionproxyinstance.h>
#include <buildboxcasd_server.h>
#include <buildboxcommon_grpctestserver.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_remoteexecutionclient.h>

#include <build/bazel/remote/execution/v2/remote_execution_mock.grpc.pb.h>
#include <google/longrunning/operations_mock.grpc.pb.h>
#include <memory>

using namespace buildboxcasd;
using namespace buildboxcommon;
using namespace google::longrunning;

class MockWriter final : public grpc::ServerWriterInterface<Operation> {
  public:
    MockWriter() {};

    bool Write(const Operation &msg, grpc::WriteOptions options) override
    {
        write_calls += 1;
        operations.push_back(msg);
        return true;
    };

    void SendInitialMetadata() override {};

    int write_calls = 0;
    std::vector<Operation> operations;
};

/**
 * Fixture to support testing the Execution proxy
 *
 * This fixture sets up a test gRPC server, and several expected requests
 * and responses which will be used by tests to verify the proxy behaviour.
 *
 * The tests implement a thread which provides the mock handler for the
 * test gRPC server, allowing it to respond to requests sent during the test
 * in a predictable manner.
 */
class ExecutionInstanceTestFixture : public ::testing::Test {
  protected:
    GrpcTestServer testServer;

    std::shared_ptr<ExecutionProxyInstance> instance;

    const std::string instance_name = "";
    const std::string remote_instance_name = "exec-dev";

    Digest actionDigest;
    ExecuteRequest expectedExecuteRequest;
    ExecuteRequest initialExecuteRequest;
    ExecuteResponse executeResponse;
    WaitExecutionRequest expectedWaitExecutionRequest;
    GetOperationRequest expectedGetOperationRequest;
    GetOperationRequest getOperationRequest;
    CancelOperationRequest expectedCancelOperationRequest;
    CancelOperationRequest cancelOperationRequest;
    ListOperationsRequest expectedListOperationsRequest;
    ListOperationsRequest listOperationsRequest;

    Digest stdErrDigest;

    Operation operation;
    Operation expected_operation;
    ListOperationsResponse expectedListOperationsResponse;
    ListOperationsResponse listOperationsResponse;

    ExecutionInstanceTestFixture()
    {

        ConnectionOptions exec_endpoint;
        exec_endpoint.setInstanceName(remote_instance_name);
        exec_endpoint.setUrl(testServer.url());

        instance = std::make_shared<ExecutionProxyInstance>(exec_endpoint,
                                                            instance_name);

        // Construct the expected requests we expect the proxy to send.
        actionDigest.set_hash("Action digest hash here");
        *expectedExecuteRequest.mutable_action_digest() = actionDigest;

        // Populate the example Operation message we'll respond with
        operation.set_name(remote_instance_name + "/operation-name");
        operation.set_done(true);
        operation.mutable_response()->PackFrom(executeResponse);

        expected_operation.CopyFrom(operation);
        expected_operation.set_name(instance_name + "#" + operation.name());

        expectedGetOperationRequest.set_name(operation.name());
        getOperationRequest.set_name(instance_name + "#" + operation.name());

        expectedCancelOperationRequest.set_name(operation.name());
        cancelOperationRequest.set_name(instance_name + "#" +
                                        operation.name());

        expectedListOperationsRequest.set_name(remote_instance_name);
        listOperationsRequest.set_name(instance_name + "#" +
                                       expectedListOperationsRequest.name());

        Operation *ptr = listOperationsResponse.add_operations();
        *ptr = operation;
        Operation *expected_ptr =
            expectedListOperationsResponse.add_operations();
        *expected_ptr = expected_operation;

        expectedExecuteRequest.set_instance_name(remote_instance_name);
        *expectedWaitExecutionRequest.mutable_name() = operation.name();
    }
};

TEST_F(ExecutionInstanceTestFixture, ExecuteRequestProxy)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer, "/build.bazel.remote.execution.v2.Execution/Execute");
        ctx.read(expectedExecuteRequest);
        ctx.writeAndFinish(operation);
    });

    // Send an Execute request and ensure that the operation is sent to the
    // writer and the status is returned.
    try {
        grpc::ServerContext context;
        MockWriter writer = MockWriter();
        grpc::Status status =
            instance->Execute(&context, expectedExecuteRequest, &writer);

        EXPECT_TRUE(writer.write_calls == 1);
        EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            writer.operations[0], expected_operation));
        EXPECT_TRUE(status.ok());
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture, ExecuteRequestProxyFailedPrecondition)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer, "/build.bazel.remote.execution.v2.Execution/Execute");
        grpc::Status s = grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
                                      "Command not in storage");
        ctx.finish(s);
    });

    // Send an Execute request and ensure that the operation is sent to the
    // writer and the status is returned.
    try {
        grpc::ServerContext context;
        MockWriter writer = MockWriter();
        grpc::Status status =
            instance->Execute(&context, expectedExecuteRequest, &writer);

        EXPECT_TRUE(status.error_code() ==
                    grpc::StatusCode::FAILED_PRECONDITION);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture, WaitExecutionRequestProxy)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer,
            "/build.bazel.remote.execution.v2.Execution/WaitExecution");
        ctx.read(expectedWaitExecutionRequest);
        ctx.writeAndFinish(operation);
    });

    // Send a WaitExecution request and ensure that the operation is sent to
    // the writer and the status is returned.
    try {
        grpc::ServerContext context;
        MockWriter writer = MockWriter();
        grpc::Status status = instance->WaitExecution(
            &context, expectedWaitExecutionRequest, &writer);

        EXPECT_TRUE(writer.write_calls == 1);
        EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            writer.operations[0], expected_operation));
        EXPECT_TRUE(status.ok());
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture, WaitExecutionRequestProxyInvalidArgument)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer,
            "/build.bazel.remote.execution.v2.Execution/WaitExecution");
        grpc::Status s = grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                                      "Operation name does not exist");
        ctx.finish(s);
    });

    // Send an WaitExecution request and ensure that the operation is sent to
    // the writer and the status is returned.
    try {
        grpc::ServerContext context;
        MockWriter writer = MockWriter();
        grpc::Status status = instance->WaitExecution(
            &context, expectedWaitExecutionRequest, &writer);

        EXPECT_TRUE(status.error_code() == grpc::StatusCode::INVALID_ARGUMENT);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture, GetOperationRequestProxy)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/GetOperation");
        // Read the GetOperationRequest and check that it matches the request
        // we expect to be sent to the remote
        ctx.read(expectedGetOperationRequest);
        ctx.writeAndFinish(operation);
    });

    try {
        Operation actual;
        // Handle a GetOperationRequest, including the local instance name in
        // the name field.
        grpc::Status status =
            instance->GetOperation(getOperationRequest, &actual);

        EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            expected_operation, actual));
        EXPECT_TRUE(actual.name() == getOperationRequest.name());
        EXPECT_TRUE(status.ok());
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture, GetOperationRequestProxyMissingOperation)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/GetOperation");
        ctx.read(expectedGetOperationRequest);
        grpc::Status s = grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                                      "Operation not found");
        ctx.finish(s);
    });

    // Send a GetOperation request and ensure that a non-OK status is
    // returned correctly.
    try {
        Operation actual;
        grpc::Status status =
            instance->GetOperation(expectedGetOperationRequest, &actual);

        EXPECT_TRUE(status.error_code() == grpc::StatusCode::INVALID_ARGUMENT);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture, CancelOperationRequestProxy)
{
    bool cancelled = false;
    std::thread serverHandler([this, &cancelled]() {
        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/CancelOperation");
        ctx.read(expectedCancelOperationRequest);
        cancelled = true;
        ctx.writeAndFinish(google::protobuf::Empty());
    });

    // Send a CancelOperation request and ensure that the server handler was
    // called as expected.
    try {
        google::protobuf::Empty empty;
        grpc::Status status =
            instance->CancelOperation(cancelOperationRequest, &empty);

        EXPECT_TRUE(status.ok());
        EXPECT_TRUE(cancelled);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture,
       CancelOperationRequestProxyMissingOperation)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/CancelOperation");
        ctx.read(expectedCancelOperationRequest);
        grpc::Status s = grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                                      "Operation not found");
        ctx.finish(s);
    });

    // Send a CancelOperation request and ensure that the server handler was
    // called as expected.
    try {
        google::protobuf::Empty empty;
        grpc::Status status =
            instance->CancelOperation(cancelOperationRequest, &empty);

        EXPECT_TRUE(status.error_code() == grpc::StatusCode::INVALID_ARGUMENT);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture, ListOperationsRequestProxy)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/ListOperations");
        ctx.read(expectedListOperationsRequest);
        ctx.writeAndFinish(listOperationsResponse);
    });

    // Send a ListOperations request and ensure that the server handler was
    // called as expected.
    try {
        ListOperationsResponse response;
        grpc::Status status =
            instance->ListOperations(listOperationsRequest, &response);

        EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            response, expectedListOperationsResponse));
        EXPECT_TRUE(status.ok());
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionInstanceTestFixture, ListOperationsRequestProxyBadQuery)
{
    std::thread serverHandler([this]() {
        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/ListOperations");
        ctx.read(expectedListOperationsRequest);
        grpc::Status s =
            grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Bad query");
        ctx.finish(s);
    });

    // Send a ListOperations request and ensure that the non-OK status code
    // is propagated to the client correctly
    try {
        ListOperationsResponse response;
        grpc::Status status =
            instance->ListOperations(listOperationsRequest, &response);

        EXPECT_TRUE(status.error_code() == grpc::StatusCode::INVALID_ARGUMENT);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}
