// Copyright 2022-2023 Bloomberg Finance L.P
//
// 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_grpctestserver.h>
#include <buildboxcommon_localexecutionclient.h>

#include <gtest/gtest.h>

#include <signal.h>
#include <stdlib.h>
#include <system_error>
#include <thread>

using namespace buildboxcommon;
using namespace testing;

namespace {
const auto digestFunctionInitializer = []() {
    buildboxcommon::DigestGenerator::init();
    return 0;
}();
} // namespace

google::rpc::Status runAction(buildboxcommon::Action *action,
                              const Command &command,
                              ActionResult *actionResult, bool cancel = false)
{
    // Ignore SIGPIPE in case of using sockets + grpc without MSG_NOSIGNAL
    // support configured
    struct sigaction sa = {};
    sa.sa_handler = SIG_IGN;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = 0;
    if (sigaction(SIGPIPE, &sa, nullptr) < 0) {
        throw std::system_error(errno, std::system_category(),
                                "Unable to ignore SIGPIPE");
    }

    const char *runnerCommand = getenv("BUILDBOX_RUN");
    EXPECT_NE(runnerCommand, nullptr);

    const auto commandSerialized = command.SerializeAsString();
    const auto commandDigest = DigestGenerator::hash(commandSerialized);
    action->mutable_command_digest()->CopyFrom(commandDigest);

    buildboxcommon::GrpcTestServer testServer;

    ConnectionOptions casServer;
    casServer.setUrl(testServer.url());

    LocalExecutionClient execClient(casServer);
    execClient.setRunner(runnerCommand, {"--no-logs-capture"});
    execClient.init();

    auto operation = execClient.asyncExecuteAction(*action);
    if (!operation.done() && cancel) {
        EXPECT_TRUE(execClient.cancelOperation(operation.name()));
        operation = execClient.waitExecution(operation.name());
    }
    else if (!operation.done()) {
        std::thread serverHandler([&]() {
            buildboxcommon::GrpcTestServerContext ctxCap(
                &testServer, "/build.bazel.remote.execution.v2.Capabilities/"
                             "GetCapabilities");
            ctxCap.finish(grpc::Status(grpc::UNIMPLEMENTED, "Unimplemented"));

            ReadRequest expectedReadRequest;
            ReadResponse readResponse;
            expectedReadRequest.set_resource_name(
                "blobs/" + commandDigest.hash() + "/" +
                std::to_string(commandDigest.size_bytes()));
            readResponse.set_data(commandSerialized);
            buildboxcommon::GrpcTestServerContext ctx(
                &testServer, "/google.bytestream.ByteStream/Read");
            ctx.read(expectedReadRequest);
            ctx.writeAndFinish(readResponse);
        });

        try {
            operation = execClient.waitExecution(operation.name());
        }
        catch (...) {
            serverHandler.join();
            throw;
        }
        serverHandler.join();
    }

    EXPECT_TRUE(operation.done());

    buildboxcommon::ExecuteResponse executeResponse;
    EXPECT_TRUE(operation.response().UnpackTo(&executeResponse));
    *actionResult = executeResponse.result();
    return executeResponse.status();
}

Command createSleepCommand(int time)
{
    Command command;
    command.add_arguments("/usr/bin/env");
    command.add_arguments("sleep");
    command.add_arguments(std::to_string(time));
    return command;
}

TEST(LocalExecutionClientTests, RunSingleCommand)
{
    buildboxcommon::Action action;
    Command command = createSleepCommand(1);

    ActionResult actionResult;
    const auto status = runAction(&action, command, &actionResult);

    EXPECT_EQ(status.code(), grpc::StatusCode::OK);
    EXPECT_EQ(actionResult.exit_code(), 0);
}

TEST(LocalExecutionClientTests, RunSingleCommandTimeout)
{
    Action action;
    action.mutable_timeout()->CopyFrom(
        google::protobuf::util::TimeUtil::SecondsToDuration(1));
    Command command = createSleepCommand(5);

    ActionResult actionResult;
    const auto status = runAction(&action, command, &actionResult);

    EXPECT_EQ(status.code(), grpc::StatusCode::DEADLINE_EXCEEDED);
    EXPECT_NE(actionResult.exit_code(), 0);
}

TEST(LocalExecutionClientTests, RunCommandThenCancel)
{
    Action action;
    // long sleep, should be cancelled
    Command command = createSleepCommand(100);

    ActionResult actionResult;
    const auto status = runAction(&action, command, &actionResult, true);

    EXPECT_EQ(status.code(), grpc::StatusCode::CANCELLED);
    EXPECT_EQ(actionResult.exit_code(), 0);
}

TEST(LocalExecutionClientTests, RunNonexistantCommand)
{
    Action action;
    Command command;
    command.add_arguments("/my/imaginary/command");
    command.add_arguments("somearg");

    ActionResult actionResult;
    const auto status = runAction(&action, command, &actionResult);

    EXPECT_EQ(status.code(), grpc::StatusCode::NOT_FOUND);
    EXPECT_EQ(actionResult.exit_code(), 127);
}
