// 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.

#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/stream_executor/device_memory.h"
#include "tensorflow/stream_executor/event.h"
#include "tensorflow/stream_executor/stream.h"

#ifdef GOOGLE_CUDA
#include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"
#include "tensorflow/core/platform/stream_executor.h"
#endif

#include "arrow/adapters/tensorflow/convert.h"
#include "arrow/api.h"
#include "arrow/io/api.h"
#include "arrow/util/logging.h"

// These headers do not include Python.h
#include "arrow/python/deserialize.h"
#include "arrow/python/serialize.h"

#include "plasma/client.h"

namespace tf = tensorflow;

using ArrowStatus = arrow::Status;
using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;

using Event = perftools::gputools::Event;
using Stream = perftools::gputools::Stream;

// NOTE(zongheng): for some reason using unique_ptr or shared_ptr results in
// CUDA_ERROR_DEINITIALIZED on program exit.  I suspect this is because the
// static object's dtor gets called *after* TensorFlow's own CUDA cleanup.
// Instead, we use a raw pointer here and manually clean up in the Ops' dtors.
static Stream* d2h_stream = nullptr;
static tf::mutex d2h_stream_mu;

// TODO(zongheng): CPU kernels' std::memcpy might be able to be sped up by
// parallelization.

int64_t get_byte_width(const arrow::DataType& dtype) {
  return arrow::internal::checked_cast<const arrow::FixedWidthType&>(dtype)
      .bit_width() / CHAR_BIT;
}

// Put:  tf.Tensor -> plasma.
template <typename Device>
class TensorToPlasmaOp : public tf::AsyncOpKernel {
 public:
  explicit TensorToPlasmaOp(tf::OpKernelConstruction* context) : tf::AsyncOpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("plasma_store_socket_name",
                                             &plasma_store_socket_name_));
    tf::mutex_lock lock(mu_);
    if (!connected_) {
      VLOG(1) << "Connecting to Plasma...";
      ARROW_CHECK_OK(client_.Connect(plasma_store_socket_name_));
      VLOG(1) << "Connected!";
      connected_ = true;
    }
  }

  ~TensorToPlasmaOp() override {
    {
      tf::mutex_lock lock(mu_);
      ARROW_CHECK_OK(client_.Disconnect());
      connected_ = false;
    }
    {
      tf::mutex_lock lock(d2h_stream_mu);
      if (d2h_stream != nullptr) {
        delete d2h_stream;
      }
    }
  }

  void ComputeAsync(tf::OpKernelContext* context, DoneCallback done) override {
    const int num_inputs = context->num_inputs();
    OP_REQUIRES_ASYNC(
        context, num_inputs >= 2,
        tf::errors::InvalidArgument("Input should have at least 1 tensor and 1 object_id"),
        done);
    const int num_tensors = num_inputs - 1;

    // Check that all tensors have the same dtype
    tf::DataType tf_dtype = context->input(0).dtype();
    for (int i = 1; i < num_inputs - 1; i++) {
      if (tf_dtype != context->input(i).dtype()) {
        ARROW_CHECK_OK(arrow::Status(arrow::StatusCode::TypeError,
                                     "All input tensors must have the same data type"));
      }
    }

    std::shared_ptr<arrow::DataType> arrow_dtype;
    ARROW_CHECK_OK(arrow::adapters::tensorflow::GetArrowType(tf_dtype, &arrow_dtype));
    int64_t byte_width = get_byte_width(*arrow_dtype);

    std::vector<size_t> offsets;
    offsets.reserve(num_tensors + 1);
    offsets.push_back(0);
    int64_t total_bytes = 0;
    for (int i = 0; i < num_tensors; ++i) {
      const size_t s = context->input(i).TotalBytes();
      CHECK_EQ(s, context->input(i).NumElements() * byte_width);
      CHECK_GT(s, 0);
      total_bytes += s;
      offsets.push_back(total_bytes);
    }

    const tf::Tensor& plasma_object_id = context->input(num_inputs - 1);
    CHECK_EQ(plasma_object_id.NumElements(), 1);
    const std::string& plasma_object_id_str = plasma_object_id.flat<std::string>()(0);
    VLOG(1) << "plasma_object_id_str: '" << plasma_object_id_str << "'";
    const plasma::ObjectID object_id =
        plasma::ObjectID::from_binary(plasma_object_id_str);

    std::vector<int64_t> shape = {total_bytes / byte_width};

    arrow::io::MockOutputStream mock;
    ARROW_CHECK_OK(arrow::py::WriteNdarrayHeader(arrow_dtype, shape, 0, &mock));
    int64_t header_size = mock.GetExtentBytesWritten();

    std::shared_ptr<Buffer> data_buffer;
    {
      tf::mutex_lock lock(mu_);
      ARROW_CHECK_OK(client_.Create(object_id, header_size + total_bytes,
                                    /*metadata=*/nullptr, 0, &data_buffer));
    }

    int64_t offset;
    arrow::io::FixedSizeBufferWriter buf(data_buffer);
    ARROW_CHECK_OK(arrow::py::WriteNdarrayHeader(arrow_dtype, shape, total_bytes, &buf));
    ARROW_CHECK_OK(buf.Tell(&offset));

    uint8_t* data = reinterpret_cast<uint8_t*>(data_buffer->mutable_data() + offset);

    auto wrapped_callback = [this, context, done, data_buffer, data, object_id]() {
      {
        tf::mutex_lock lock(mu_);
        ARROW_CHECK_OK(client_.Seal(object_id));
        ARROW_CHECK_OK(client_.Release(object_id));
#ifdef GOOGLE_CUDA
        auto orig_stream = context->op_device_context()->stream();
        auto stream_executor = orig_stream->parent();
        CHECK(stream_executor->HostMemoryUnregister(static_cast<void*>(data)));
#endif
      }
      context->SetStatus(tensorflow::Status::OK());
      done();
    };

    if (std::is_same<Device, CPUDevice>::value) {
      for (int i = 0; i < num_tensors; ++i) {
        const auto& input_tensor = context->input(i);
        std::memcpy(static_cast<void*>(data + offsets[i]),
                    input_tensor.tensor_data().data(),
                    static_cast<tf::uint64>(offsets[i + 1] - offsets[i]));
      }
      wrapped_callback();
    } else {
#ifdef GOOGLE_CUDA
      auto orig_stream = context->op_device_context()->stream();
      OP_REQUIRES_ASYNC(context, orig_stream != nullptr,
                        tf::errors::Internal("No GPU stream available."), done);
      auto stream_executor = orig_stream->parent();

      // NOTE(zongheng): this is critical of getting good performance out of D2H
      // async memcpy.  Under the hood it performs cuMemHostRegister(), see:
      // http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1gf0a9fe11544326dabd743b7aa6b54223
      CHECK(stream_executor->HostMemoryRegister(static_cast<void*>(data),
                                                static_cast<tf::uint64>(total_bytes)));

      {
        tf::mutex_lock l(d2h_stream_mu);
        if (d2h_stream == nullptr) {
          d2h_stream = new Stream(stream_executor);
          CHECK(d2h_stream->Init().ok());
        }
      }

      // Needed to make sure the input buffers have been computed.
      // NOTE(ekl): this is unnecessary when the op is behind a NCCL allreduce already
      CHECK(d2h_stream->ThenWaitFor(orig_stream).ok());

      for (int i = 0; i < num_tensors; ++i) {
        const auto& input_tensor = context->input(i);
        auto input_buffer = const_cast<char*>(input_tensor.tensor_data().data());
        perftools::gputools::DeviceMemoryBase wrapped_src(
            static_cast<void*>(input_buffer));
        const bool success =
            d2h_stream
                ->ThenMemcpy(static_cast<void*>(data + offsets[i]), wrapped_src,
                             static_cast<tf::uint64>(offsets[i + 1] - offsets[i]))
                .ok();
        OP_REQUIRES_ASYNC(context, success,
                          tf::errors::Internal("D2H memcpy failed to be enqueued."), done);
      }
      context->device()->tensorflow_gpu_device_info()->event_mgr->ThenExecute(
          d2h_stream, std::move(wrapped_callback));
#endif
    }
  }

 private:
  std::string plasma_store_socket_name_;

  tf::mutex mu_;
  bool connected_ = false;
  plasma::PlasmaClient client_ GUARDED_BY(mu_);
};

static Stream* h2d_stream = nullptr;
static tf::mutex h2d_stream_mu;

// Get:  plasma -> tf.Tensor.
template <typename Device>
class PlasmaToTensorOp : public tf::AsyncOpKernel {
 public:
  explicit PlasmaToTensorOp(tf::OpKernelConstruction* context) : tf::AsyncOpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("plasma_store_socket_name",
                                             &plasma_store_socket_name_));
    tf::mutex_lock lock(mu_);
    if (!connected_) {
      VLOG(1) << "Connecting to Plasma...";
      ARROW_CHECK_OK(client_.Connect(plasma_store_socket_name_));
      VLOG(1) << "Connected!";
      connected_ = true;
    }
  }

  ~PlasmaToTensorOp() override {
    {
      tf::mutex_lock lock(mu_);
      ARROW_CHECK_OK(client_.Disconnect());
      connected_ = false;
    }
    {
      tf::mutex_lock lock(h2d_stream_mu);
      if (h2d_stream != nullptr) {
        delete h2d_stream;
      }
    }
  }

  void ComputeAsync(tf::OpKernelContext* context, DoneCallback done) override {
    const tf::Tensor& plasma_object_id = context->input(0);
    CHECK_EQ(plasma_object_id.NumElements(), 1);
    const std::string& plasma_object_id_str = plasma_object_id.flat<std::string>()(0);

    VLOG(1) << "plasma_object_id_str: '" << plasma_object_id_str << "'";
    const plasma::ObjectID object_id =
        plasma::ObjectID::from_binary(plasma_object_id_str);

    plasma::ObjectBuffer object_buffer;
    {
      tf::mutex_lock lock(mu_);
      // NOTE(zongheng): this is a blocking call.  We might want to (1) make
      // Plasma asynchronous, (2) launch a thread / event here ourselves, or
      // something like that...
      ARROW_CHECK_OK(client_.Get(&object_id, /*num_objects=*/1,
                                 /*timeout_ms=*/-1, &object_buffer));
    }

    std::shared_ptr<arrow::Tensor> ndarray;
    ARROW_CHECK_OK(arrow::py::NdarrayFromBuffer(object_buffer.data, &ndarray));

    int64_t byte_width = get_byte_width(*ndarray->type());
    const int64_t size_in_bytes = ndarray->data()->size();

    tf::TensorShape shape({static_cast<int64_t>(size_in_bytes / byte_width)});

    const float* plasma_data = reinterpret_cast<const float*>(ndarray->raw_data());

    tf::Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK_ASYNC(context, context->allocate_output(0, shape, &output_tensor),
                         done);

    auto wrapped_callback = [this, context, done, plasma_data, object_id]() {
      {
        tf::mutex_lock lock(mu_);
        ARROW_CHECK_OK(client_.Release(object_id));
#ifdef GOOGLE_CUDA
        auto orig_stream = context->op_device_context()->stream();
        auto stream_executor = orig_stream->parent();
        CHECK(stream_executor->HostMemoryUnregister(
            const_cast<void*>(static_cast<const void*>(plasma_data))));
#endif
      }
      done();
    };

    if (std::is_same<Device, CPUDevice>::value) {
      std::memcpy(
          reinterpret_cast<void*>(const_cast<char*>(output_tensor->tensor_data().data())),
          plasma_data, size_in_bytes);
      wrapped_callback();
    } else {
#ifdef GOOGLE_CUDA
      auto orig_stream = context->op_device_context()->stream();
      OP_REQUIRES_ASYNC(context, orig_stream != nullptr,
                        tf::errors::Internal("No GPU stream available."), done);
      auto stream_executor = orig_stream->parent();

      {
        tf::mutex_lock l(h2d_stream_mu);
        if (h2d_stream == nullptr) {
          h2d_stream = new Stream(stream_executor);
          CHECK(h2d_stream->Init().ok());
        }
      }

      // Important.  See note in T2P op.
      CHECK(stream_executor->HostMemoryRegister(
          const_cast<void*>(static_cast<const void*>(plasma_data)),
          static_cast<tf::uint64>(size_in_bytes)));

      perftools::gputools::DeviceMemoryBase wrapped_dst(
          reinterpret_cast<void*>(const_cast<char*>(output_tensor->tensor_data().data())));
      const bool success =
          h2d_stream
              ->ThenMemcpy(&wrapped_dst, static_cast<const void*>(plasma_data),
                           static_cast<tf::uint64>(size_in_bytes))
              .ok();
      OP_REQUIRES_ASYNC(context, success,
                        tf::errors::Internal("H2D memcpy failed to be enqueued."), done);

      // Without this sync the main compute stream might proceed to use the
      // Tensor buffer, but its contents might still be in-flight from our
      // h2d_stream.
      CHECK(orig_stream->ThenWaitFor(h2d_stream).ok());

      context->device()->tensorflow_gpu_device_info()->event_mgr->ThenExecute(
          h2d_stream, std::move(wrapped_callback));
#endif
    }
  }

 private:
  std::string plasma_store_socket_name_;

  tf::mutex mu_;
  bool connected_ = false;
  plasma::PlasmaClient client_ GUARDED_BY(mu_);
};

REGISTER_OP("TensorToPlasma")
    .Input("input_tensor: dtypes")
    .Input("plasma_object_id: string")
    .Attr("dtypes: list(type)")
    .Attr("plasma_store_socket_name: string");

REGISTER_KERNEL_BUILDER(Name("TensorToPlasma").Device(tf::DEVICE_CPU),
                        TensorToPlasmaOp<CPUDevice>);
#ifdef GOOGLE_CUDA
REGISTER_KERNEL_BUILDER(Name("TensorToPlasma").Device(tf::DEVICE_GPU),
                        TensorToPlasmaOp<GPUDevice>);
#endif

REGISTER_OP("PlasmaToTensor")
    .Input("plasma_object_id: string")
    .Output("tensor: dtype")
    .Attr("dtype: type")
    .Attr("plasma_store_socket_name: string");

REGISTER_KERNEL_BUILDER(Name("PlasmaToTensor").Device(tf::DEVICE_CPU),
                        PlasmaToTensorOp<CPUDevice>);
#ifdef GOOGLE_CUDA
REGISTER_KERNEL_BUILDER(Name("PlasmaToTensor").Device(tf::DEVICE_GPU),
                        PlasmaToTensorOp<GPUDevice>);
#endif