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

#include <grpcpp/support/status.h>
#include <memory>
#include <sstream>
#include <utility>

namespace buildboxcommon {

LogStreamWriter::LogStreamWriter(std::string resourceName,
                                 const ConnectionOptions &connectionOptions)
    : d_resourceName(std::move(resourceName)),
      d_grpcClient(std::make_shared<GrpcClient>()), d_writeOffset(0),
      d_writeCommitted(false), d_resourceReady(false)
{
    d_grpcClient->init(connectionOptions);
    d_byteStreamClient =
        google::bytestream::ByteStream::NewStub(d_grpcClient->channel());
}

LogStreamWriter::LogStreamWriter(std::string resourceName,
                                 const std::shared_ptr<GrpcClient> &client)
    : d_resourceName(std::move(resourceName)), d_grpcClient(client),
      d_writeOffset(0), d_writeCommitted(false), d_resourceReady(false)
{
    if (client->channel() != nullptr) {
        d_byteStreamClient =
            google::bytestream::ByteStream::NewStub(d_grpcClient->channel());
    }
}

void LogStreamWriter::init(
    std::shared_ptr<ByteStream::StubInterface> bytestreamClient)
{
    d_byteStreamClient = std::move(bytestreamClient);
}

void LogStreamWriter::reset()
{
    d_writeOffset = 0;
    d_writeCommitted = false;
    d_resourceReady = false;
    d_bytestreamWriter = nullptr;
    d_clientContext = nullptr;
}

bool LogStreamWriter::write(const std::string &data)
{
    if (d_writeCommitted) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Attempted to `write()` after `commit()`.")
    }

    auto writeLambda = [&data, this](grpc::ClientContext &) {
        if (!d_resourceReady) {
            // query the current write status to know if resource exists and
            // current offset
            auto queryStatus = queryStreamWriteStatus();

            if (d_resourceReady) {
                BUILDBOX_LOG_DEBUG(
                    "`QueryWriteStatus()` returned successfully. We can "
                    "now start writing to the stream.");
            }
            else {
                BUILDBOX_LOG_WARNING(
                    "`QueryWriteStatus()` failed. Aborting the "
                    "call to `ByteStream.Write()`");
                return queryStatus;
            }
        }

        WriteRequest request;
        request.set_resource_name(d_resourceName);
        request.set_write_offset(d_writeOffset);
        request.set_data(data);
        request.set_finish_write(false);
        auto &writer = bytestreamWriter();
        if (!writer->Write(request)) {
            auto status = writer->Finish();
            BUILDBOX_LOG_ERROR("Failed to write logstream chunk, code="
                               << status.error_code()
                               << " message=" << status.error_message());
            // reset the writer for the next retry
            reset();
            return status;
        }
        d_writeOffset +=
            static_cast<google::protobuf::int64>(request.data().length());
        return grpc::Status::OK;
    };

    try {
        d_grpcClient->issueRequest(writeLambda, "(LogStream) ByteStream.Write",
                                   nullptr);
        return true;
    }
    catch (GrpcError &) {
        return false;
    }
}

bool LogStreamWriter::commit()
{
    if (d_writeCommitted) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Attempted to `commit()` an already commited write.")
    }

    auto commitWriteLambda = [this](grpc::ClientContext &) {
        if (!d_resourceReady) {
            auto queryStatus = queryStreamWriteStatus();
            if (!d_resourceReady) {
                BUILDBOX_LOG_WARNING(
                    "Failed to query write status during commit, code="
                    << queryStatus.error_code()
                    << " message=" << queryStatus.error_message());
                return queryStatus;
            }
        }

        WriteRequest request;
        request.set_resource_name(d_resourceName);
        request.set_write_offset(d_writeOffset);
        request.set_finish_write(true);
        auto &writer = bytestreamWriter();
        if (!writer->Write(request)) {
            auto status = writer->Finish();
            BUILDBOX_LOG_ERROR("Failed to commit logstream , code="
                               << status.error_code()
                               << " message=" << status.error_message());
            // reset the writer for the next retry
            reset();
            return status;
        }

        bytestreamWriter()->WritesDone();
        auto status = bytestreamWriter()->Finish();
        if (status.ok()) {
            const auto bytesWritten = d_writeOffset;
            if (d_writeResponse.committed_size() != bytesWritten) {
                std::stringstream errorMessage;
                errorMessage << "Server reported uncommitted data: "
                             << d_writeResponse.committed_size() << " of "
                             << bytesWritten << " bytes";
                BUILDBOX_LOG_DEBUG(errorMessage.str());
                return grpc::Status(grpc::StatusCode::DATA_LOSS,
                                    errorMessage.str());
            }
        }
        d_writeCommitted = true;
        return status;
    };

    try {
        d_grpcClient->issueRequest(
            commitWriteLambda,
            "(LogStream) ByteStream.Write(set_finish_write=True)", nullptr);
        return true;
    }
    catch (GrpcError &) {
        return false;
    }
}

grpc::Status LogStreamWriter::queryStreamWriteStatus()
{
    QueryWriteStatusRequest request;
    request.set_resource_name(d_resourceName);

    auto queryWriteStatusLambda = [&request,
                                   this](grpc::ClientContext &context) {
        QueryWriteStatusResponse response;

        auto status =
            d_byteStreamClient->QueryWriteStatus(&context, request, &response);
        d_writeOffset = response.committed_size();
        d_resourceReady = status.ok();

        return status;
    };

    try {
        d_grpcClient->issueRequest(queryWriteStatusLambda,
                                   "(LogStream)QueryWriteStatus()", nullptr);
        return grpc::Status::OK;
    }
    catch (GrpcError &err) {
        const auto &status = err.status;
        BUILDBOX_LOG_ERROR("Failed to query LogStream write status, code="
                           << status.error_code()
                           << " message=" << status.error_message());
        return err.status;
    }
}

LogStreamWriter::ByteStreamClientWriter &LogStreamWriter::bytestreamWriter()
{
    if (d_clientContext == nullptr) {
        d_clientContext = std::make_unique<grpc::ClientContext>();
    }
    if (d_bytestreamWriter == nullptr) {
        d_bytestreamWriter =
            d_byteStreamClient->Write(d_clientContext.get(), &d_writeResponse);
    }

    return d_bytestreamWriter;
}

LogStream LogStreamWriter::createLogStream(
    const std::string &parent,
    const buildboxcommon::ConnectionOptions &connectionOptions)
{
    const auto channel = connectionOptions.createChannel();
    std::unique_ptr<LogStreamService::StubInterface> logStreamClient =
        LogStreamService::NewStub(channel);

    const int retryLimit = std::stoi(connectionOptions.d_retryLimit);
    const int retryDelay = std::stoi(connectionOptions.d_retryDelay);

    return createLogStream(parent, retryLimit, retryDelay,
                           logStreamClient.get());
}

LogStream LogStreamWriter::createLogStream(
    const std::string &parent, const int retryLimit, const int retryDelay,
    LogStreamService::StubInterface *logstreamClient)
{
    if (logstreamClient == nullptr) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::invalid_argument,
                                       "logstreamClient argument is nullptr");
    }

    grpc::ClientContext context;

    CreateLogStreamRequest request;
    request.set_parent(parent);

    LogStream createdLogStream;
    auto createLogStreamLambda = [&context, &request, &logstreamClient,
                                  &createdLogStream](grpc::ClientContext &) {
        return logstreamClient->CreateLogStream(&context, request,
                                                &createdLogStream);
    };

    GrpcRetrier retrier(retryLimit, std::chrono::milliseconds(retryDelay),
                        createLogStreamLambda, "CreateLogStream()");

    if (retrier.issueRequest() && retrier.status().ok()) {
        return createdLogStream;
    }

    const auto errorMessage = "CreateLogStream() failed: [" +
                              std::to_string(retrier.status().error_code()) +
                              ": " + retrier.status().error_message() + "]";
    throw GrpcError(errorMessage, retrier.status());
}

} // namespace buildboxcommon
