// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

#ifndef PLASMA_PROTOCOL_H
#define PLASMA_PROTOCOL_H

#include <memory>
#include <unordered_map>
#include <vector>

#include "arrow/status.h"
#include "plasma/plasma.h"
#include "plasma/plasma_generated.h"

namespace plasma {

using arrow::Status;

template <class T>
bool verify_flatbuffer(T* object, uint8_t* data, size_t size) {
  flatbuffers::Verifier verifier(data, size);
  return object->Verify(verifier);
}

/* Plasma receive message. */

Status PlasmaReceive(int sock, int64_t message_type, std::vector<uint8_t>* buffer);

/* Plasma Create message functions. */

Status SendCreateRequest(int sock, ObjectID object_id, int64_t data_size,
                         int64_t metadata_size, int device_num);

Status ReadCreateRequest(uint8_t* data, size_t size, ObjectID* object_id,
                         int64_t* data_size, int64_t* metadata_size, int* device_num);

Status SendCreateReply(int sock, ObjectID object_id, PlasmaObject* object, int error,
                       int64_t mmap_size);

Status ReadCreateReply(uint8_t* data, size_t size, ObjectID* object_id,
                       PlasmaObject* object, int* store_fd, int64_t* mmap_size);

Status SendAbortRequest(int sock, ObjectID object_id);

Status ReadAbortRequest(uint8_t* data, size_t size, ObjectID* object_id);

Status SendAbortReply(int sock, ObjectID object_id);

Status ReadAbortReply(uint8_t* data, size_t size, ObjectID* object_id);

/* Plasma Seal message functions. */

Status SendSealRequest(int sock, ObjectID object_id, unsigned char* digest);

Status ReadSealRequest(uint8_t* data, size_t size, ObjectID* object_id,
                       unsigned char* digest);

Status SendSealReply(int sock, ObjectID object_id, int error);

Status ReadSealReply(uint8_t* data, size_t size, ObjectID* object_id);

/* Plasma Get message functions. */

Status SendGetRequest(int sock, const ObjectID* object_ids, int64_t num_objects,
                      int64_t timeout_ms);

Status ReadGetRequest(uint8_t* data, size_t size, std::vector<ObjectID>& object_ids,
                      int64_t* timeout_ms);

Status SendGetReply(
    int sock, ObjectID object_ids[],
    std::unordered_map<ObjectID, PlasmaObject, UniqueIDHasher>& plasma_objects,
    int64_t num_objects, const std::vector<int>& store_fds,
    const std::vector<int64_t>& mmap_sizes);

Status ReadGetReply(uint8_t* data, size_t size, ObjectID object_ids[],
                    PlasmaObject plasma_objects[], int64_t num_objects,
                    std::vector<int>& store_fds, std::vector<int64_t>& mmap_sizes);

/* Plasma Release message functions. */

Status SendReleaseRequest(int sock, ObjectID object_id);

Status ReadReleaseRequest(uint8_t* data, size_t size, ObjectID* object_id);

Status SendReleaseReply(int sock, ObjectID object_id, int error);

Status ReadReleaseReply(uint8_t* data, size_t size, ObjectID* object_id);

/* Plasma Delete message functions. */

Status SendDeleteRequest(int sock, ObjectID object_id);

Status ReadDeleteRequest(uint8_t* data, size_t size, ObjectID* object_id);

Status SendDeleteReply(int sock, ObjectID object_id, int error);

Status ReadDeleteReply(uint8_t* data, size_t size, ObjectID* object_id);

/* Satus messages. */

Status SendStatusRequest(int sock, const ObjectID* object_ids, int64_t num_objects);

Status ReadStatusRequest(uint8_t* data, size_t size, ObjectID object_ids[],
                         int64_t num_objects);

Status SendStatusReply(int sock, ObjectID object_ids[], int object_status[],
                       int64_t num_objects);

int64_t ReadStatusReply_num_objects(uint8_t* data, size_t size);

Status ReadStatusReply(uint8_t* data, size_t size, ObjectID object_ids[],
                       int object_status[], int64_t num_objects);

/* Plasma Constains message functions. */

Status SendContainsRequest(int sock, ObjectID object_id);

Status ReadContainsRequest(uint8_t* data, size_t size, ObjectID* object_id);

Status SendContainsReply(int sock, ObjectID object_id, bool has_object);

Status ReadContainsReply(uint8_t* data, size_t size, ObjectID* object_id,
                         bool* has_object);

/* Plasma Connect message functions. */

Status SendConnectRequest(int sock);

Status ReadConnectRequest(uint8_t* data, size_t size);

Status SendConnectReply(int sock, int64_t memory_capacity);

Status ReadConnectReply(uint8_t* data, size_t size, int64_t* memory_capacity);

/* Plasma Evict message functions (no reply so far). */

Status SendEvictRequest(int sock, int64_t num_bytes);

Status ReadEvictRequest(uint8_t* data, size_t size, int64_t* num_bytes);

Status SendEvictReply(int sock, int64_t num_bytes);

Status ReadEvictReply(uint8_t* data, size_t size, int64_t& num_bytes);

/* Plasma Fetch Remote message functions. */

Status SendFetchRequest(int sock, const ObjectID* object_ids, int64_t num_objects);

Status ReadFetchRequest(uint8_t* data, size_t size, std::vector<ObjectID>& object_ids);

/* Plasma Wait message functions. */

Status SendWaitRequest(int sock, ObjectRequest object_requests[], int64_t num_requests,
                       int num_ready_objects, int64_t timeout_ms);

Status ReadWaitRequest(uint8_t* data, size_t size, ObjectRequestMap& object_requests,
                       int64_t* timeout_ms, int* num_ready_objects);

Status SendWaitReply(int sock, const ObjectRequestMap& object_requests,
                     int num_ready_objects);

Status ReadWaitReply(uint8_t* data, size_t size, ObjectRequest object_requests[],
                     int* num_ready_objects);

/* Plasma Subscribe message functions. */

Status SendSubscribeRequest(int sock);

/* Data messages. */

Status SendDataRequest(int sock, ObjectID object_id, const char* address, int port);

Status ReadDataRequest(uint8_t* data, size_t size, ObjectID* object_id, char** address,
                       int* port);

Status SendDataReply(int sock, ObjectID object_id, int64_t object_size,
                     int64_t metadata_size);

Status ReadDataReply(uint8_t* data, size_t size, ObjectID* object_id,
                     int64_t* object_size, int64_t* metadata_size);

}  // namespace plasma

#endif /* PLASMA_PROTOCOL */
