/*
 * 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 <bytestreamservicer.h>

#include <buildboxcommon_logging.h>

#include <sstream>

LogStreamBytestreamServicer::LogStreamBytestreamServicer(
    const DataReceivedCallback &callback)
    : d_dataReceivedCallback(callback)
{
}

grpc::Status LogStreamBytestreamServicer::Read(
    grpc::ServerContext *, const google::bytestream::ReadRequest *request,
    grpc::ServerWriter<google::bytestream::ReadResponse> *writer)
{
    return grpc::Status(grpc::StatusCode::UNIMPLEMENTED, "Not implemented.");
}

grpc::Status LogStreamBytestreamServicer::Write(
    grpc::ServerContext *,
    grpc::ServerReader<google::bytestream::WriteRequest> *request,
    google::bytestream::WriteResponse *response)
{

    google::bytestream::WriteRequest request_message;
    if (!request->Read(&request_message)) {
        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                            "No request received.");
    }

    const std::string &resource_name = request_message.resource_name();
    BUILDBOX_LOG_INFO("`ByteStream.Write()` called");

    bool commit_write = false;
    google::protobuf::int64 bytes_written = 0;

    while (true) {
        BUILDBOX_LOG_INFO("`WriteRequest` received "
                          << printableWriteRequest(request_message));

        const std::string new_request_resource_name =
            request_message.resource_name();
        // If the `resource_name` was set on subsequent requests, it has to
        // match the initial resource name. However, `resource_name` is
        // optional and can be omitted in subsequent requests.
        // ref:
        // https://github.com/googleapis/googleapis/blob/6da3d64919c006ef40cad2026f1e39084253afe2/google/bytestream/bytestream.proto#L130-133
        if (!new_request_resource_name.empty() &&
            new_request_resource_name != resource_name) {
            return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                                "`resource_name` changed between requests.");
        }
        else if (commit_write) {
            // `finish_write` must be set only in the last request.
            return grpc::Status(
                grpc::StatusCode::INVALID_ARGUMENT,
                "Extra request sent after setting `finish_write`.");
        }

        if (request_message.write_offset() != bytes_written) {
            // (Data can only be appended at the end.)
            grpc::Status(grpc::StatusCode::OUT_OF_RANGE,
                         "`write_offset` is not valid");
        }

        // All arguments are valid, invoking the callback:
        d_dataReceivedCallback(resource_name, request_message.data(),
                               commit_write);

        commit_write = request_message.finish_write();
        bytes_written += static_cast<google::protobuf::int64>(
            request_message.data().size());

        // Read the next request message
        if (!request->Read(&request_message)) {
            break;
        }
    }

    // Finished receiving requests.

    if (!commit_write) {
        return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                            "`finish_write` was not set in the last request.");
    }

    response->set_committed_size(bytes_written);

    BUILDBOX_LOG_INFO("`ByteStream.Write()` returned OK");
    return grpc::Status::OK;
}

std::string LogStreamBytestreamServicer::printableWriteRequest(
    const google::bytestream::WriteRequest &request)
{
    std::ostringstream os;
    os << "[resource_name=" << request.resource_name() << ", "
       << "data.size()=" << request.data().size() << ", "
       << "write_offset=" << request.write_offset() << ", "
       << "finish_write=" << request.finish_write() << "]";

    return os.str();
}

grpc::Status LogStreamBytestreamServicer::QueryWriteStatus(
    grpc::ServerContext *,
    const google::bytestream::QueryWriteStatusRequest *request,
    google::bytestream::QueryWriteStatusResponse *response)
{
    BUILDBOX_LOG_INFO("`QueryWriteStatus()` request for resource name ["
                      << request->resource_name() << "].");

    // We don't support resuming, so always return 0.
    response->set_committed_size(0);
    response->set_complete(false);
    return grpc::Status::OK;
}
