/* Copyright 2022 The JAX Authors

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

#ifndef JAXLIB_PY_ARRAY_H_
#define JAXLIB_PY_ARRAY_H_

#include <Python.h>

#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string_view>
#include <utility>
#include <vector>

// placeholder for index annotation headers
#include "absl/base/thread_annotations.h"
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "nanobind/nanobind.h"
#include "jaxlib/free_threading.h"
#include "jaxlib/nb_class_ptr.h"
#include "jaxlib/py_client.h"
#include "jaxlib/py_user_context.h"
#include "jaxlib/traceback.h"
#include "xla/future.h"
#include "xla/pjrt/exceptions.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/pjrt_layout.h"
#include "xla/python/ifrt/array.h"
#include "xla/python/ifrt/device_list.h"
#include "xla/python/ifrt/rtti.h"
#include "xla/python/ifrt/shape.h"
#include "xla/python/ifrt/user_context.h"
#include "xla/python/nb_numpy.h"
#include "xla/python/pjrt_ifrt/pjrt_array.h"
#include "xla/shape.h"
#include "xla/tsl/concurrency/future.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/util.h"

namespace jax {

class PyArray;
struct PyHostValue;

// Private to PyArray, but you cannot forward declare member classes.
struct PyArray_Storage {
  PyArray_Storage(nanobind::object aval, bool weak_type, xla::nb_dtype dtype,
                  std::vector<int64_t> shape, nanobind::object sharding,
                  bool committed, nb_class_ptr<PyClient> py_client,
                  xla::ifrt::ArrayRef ifrt_array, xla::Future<> result_status);

  ~PyArray_Storage();
  nanobind::handle AsHandle();

  const nb_class_ptr<PyClient> py_client;

  // Doubly-linked list of all PyArrays known to the client. Protected by the
  // GIL or by a sharded mutex in the client. Since multiple PyArrays may share
  // the same PjRtBuffer, there may be duplicate PjRtBuffers in this list.
  PyArray_Storage* next = nullptr;
  PyArray_Storage* prev = nullptr;

  // To which shard of the client does this array belong?
  const uint8_t thread_id_bucket;

  const bool committed = false;
  const bool weak_type = false;
  mutable ft_mutex mu;

  const nanobind::object aval;
  const xla::nb_dtype dtype;
  const std::vector<int64_t> shape;

  const nanobind::object sharding;
  nanobind::object npy_value ABSL_GUARDED_BY(mu) = nanobind::none();

  xla::ifrt::ArrayRef ifrt_array ABSL_GUARDED_BY(mu);
  nanobind::object fully_replicated_array ABSL_GUARDED_BY(mu) =
      nanobind::none();

  // optional field, used only in python
  nanobind::object py_arrays ABSL_GUARDED_BY(mu) = nanobind::none();
  tsl::Future<std::shared_ptr<PyHostValue>> host_value ABSL_GUARDED_BY(mu);
  std::optional<std::vector<int64_t>> dynamic_shape ABSL_GUARDED_BY(mu) =
      std::nullopt;
  // Only set if this Array was generated by a computation that has effects.
  // This is the result status of the XLA computation that generated this
  // array.
  xla::Future<> result_status ABSL_GUARDED_BY(mu);
};

// The C++ implementation of jax.Array. A few key methods and data members are
// implemented in C++ for performance, while most of the functionalities are
// still implemented in python.
class PyArray : public nanobind::object {
 public:
  NB_OBJECT(PyArray, nanobind::object, "Array", PyArray::IsPyArray);
  PyArray() = default;

  // "__init__" methods. Only used in python
  static void PyInit(PyArray self, nanobind::object aval,
                     nanobind::object sharding,
                     absl::Span<const PyArray> py_arrays, bool committed,
                     bool skip_checks);

  // Only used in C++. Requires that the array cannot have consistency issues
  // (e.g., `sharding` devices different than `ifrt_array` devices). Arrays
  // created by users or with potentially inconsistent metadata should be
  // checked (e.g., via `MakeFromIfrtArrayAndSharding`).
  PyArray(nanobind::object aval, bool weak_type, xla::nb_dtype dtype,
          std::vector<int64_t> shape, nanobind::object sharding,
          nb_class_ptr<PyClient> py_client, xla::ifrt::ArrayRef ifrt_array,
          bool committed, xla::Future<> result_status = xla::Future<>());

  static PyArray MakeFromSingleDeviceArray(
      nb_class_ptr<PyClient> py_client, xla::ifrt::ArrayRef ifrt_array,
      bool weak_type, bool committed,
      xla::Future<> result_status = xla::Future<>());

  static PyArray MakeFromIfrtArrayAndSharding(nb_class_ptr<PyClient> py_client,
                                              xla::ifrt::ArrayRef ifrt_array,
                                              nanobind::object sharding,
                                              bool weak_type, bool committed,
                                              bool skip_checks);

  // Registers Array and related types in module m.
  static absl::Status Register(nanobind::module_& m);

  static PyArray borrow(PyObject* ptr) {
    return nanobind::borrow<PyArray>(ptr);
  }

  using Storage = PyArray_Storage;

  const nanobind::object& aval() const { return GetStorage().aval; }

  bool weak_type() const { return GetStorage().weak_type; }

  const xla::nb_dtype& dtype() const { return GetStorage().dtype; }
  absl::Span<const int64_t> shape() const { return GetStorage().shape; }

  const nanobind::object& sharding() const { return GetStorage().sharding; }
  absl::StatusOr<std::vector<int64_t>> dynamic_shape() const;

  absl::StatusOr<std::shared_ptr<const xla::PjRtLayout>> layout() {
    ABSL_ASSIGN_OR_RETURN(xla::ifrt::ArrayRef ifrt_array_ptr, GetIfrtArray());
    ABSL_ASSIGN_OR_RETURN(std::shared_ptr<const xla::PjRtLayout> layout,
                          ifrt_array_ptr->pjrt_layout());
    if (layout == nullptr) {
      ABSL_ASSIGN_OR_RETURN(
          xla::ifrt::Shape shard_shape,
          ifrt_array_ptr->sharding().GetShardShape(ifrt_array_ptr->shape()));
      ABSL_ASSIGN_OR_RETURN(
          layout, ifrt_array_ptr->client()->GetDefaultPjRtLayout(
                      ifrt_array_ptr->dtype(), shard_shape.dims(),
                      ifrt_array_ptr->sharding().devices()->devices().front(),
                      ifrt_array_ptr->sharding().memory_kind()));
    }
    return layout;
  }

  bool committed() const { return GetStorage().committed; }

  nanobind::object npy_value() const {
    const Storage& storage = GetStorage();
    ft_lock_guard lock(storage.mu);
    return storage.npy_value;
  }
  void set_npy_value(nanobind::object v) {
    Storage& storage = GetStorage();
    nanobind::object old_npy_value;
    ft_lock_guard lock(storage.mu);
    old_npy_value = std::move(storage.npy_value);
    storage.npy_value = std::move(v);
  }

  const nb_class_ptr<PyClient>& py_client() const {
    return GetStorage().py_client;
  }

  std::optional<Traceback> traceback() const {
    xla::ifrt::ArrayRef ifrt_array_ptr = ifrt_array_ref();
    if (ifrt_array_ptr == nullptr) {
      return std::nullopt;
    }
    return GetTraceback(ifrt_array_ptr->user_context().get());
  }

  // Returns xla::InvalidArgument if the buffer has been deleted.
  // See `Future` for the semantics of `IsReady` and `IsKnownReady`.
  absl::StatusOr<bool> IsReady() {
    ABSL_ASSIGN_OR_RETURN(xla::ifrt::ArrayRef ifrt_array_ptr, GetIfrtArray());
    if (ifrt_array_ptr->IsDeleted()) {
      return xla::InvalidArgument("Array has been deleted.");
    }
    xla::ifrt::UserContextScope user_context_scope(
        jax::PyUserContext::Create());
    return ifrt_array_ptr->GetReadyFuture().IsReady();
  }

  xla::Future<> result_status() const {
    const Storage& storage = GetStorage();
    ft_lock_guard lock(storage.mu);
    return storage.result_status;
  }

  xla::ifrt::ArrayRef ifrt_array_ref() const {
    const Storage& storage = GetStorage();
    ft_lock_guard lock(storage.mu);
    return storage.ifrt_array;
  }

  // Returns the underlying IFRT array, or an InvalidArgument error if the array
  // has been deleted or donated. Unlike `ifrt_array_ref()`, which returns a
  // null pointer when deleted, this method returns a non-null `ArrayRef` or a
  // formatted error status.
  absl::StatusOr<xla::ifrt::ArrayRef> GetIfrtArray(
      std::string_view api = "") const;

  int num_addressable_shards() const {
    xla::ifrt::ArrayRef ifrt_array_ptr = ifrt_array_ref();
    if (ifrt_array_ptr == nullptr) {
      return 0;
    }
    auto* arr = xla::ifrt::dyn_cast_or_null<xla::ifrt::PjRtCompatibleArray>(
        ifrt_array_ptr.get());
    if (arr == nullptr) {
      // TODO(hyeontaek): Add num_addressable_shards to ifrt.
      return num_shards();
    }
    return arr->pjrt_buffers().size();
  }

  nanobind::object py_arrays_cached();

  nanobind::object arrays();
  absl::StatusOr<PyArray> FullyReplicatedShard();

  int num_shards() const {
    xla::ifrt::ArrayRef ifrt_array_ptr = ifrt_array_ref();
    if (ifrt_array_ptr == nullptr) {
      return 0;
    }
    return ifrt_array_ptr->sharding().devices()->size();
  }

  static nanobind::handle type() {
    DCHECK(type_);
    return nanobind::handle(type_);
  }

  static bool IsPyArray(nanobind::handle arg) {
    return arg.type().is(PyArray::type());
  }

  absl::Status BlockUntilReady() const;

  absl::Status BlockUntilResultStatusIsReady();

  absl::StatusOr<size_t> GetOnDeviceSizeInBytes();
  absl::StatusOr<std::pair<nanobind::object, bool>>
  SingleDeviceArrayToNumpyArrayDidCopy();
  // Returns the numpy array and a boolean indicating whether the array was
  // copied from the devices.
  absl::StatusOr<std::pair<nanobind::object, bool>> ToNumpyArrayDidCopy();
  absl::StatusOr<nanobind::object> SingleDeviceArrayToNumpyArray();
  absl::Status CopySingleDeviceArrayToHostAsync();
  nanobind::dict CudaArrayInterface();
  absl::StatusOr<std::uintptr_t> UnsafeBufferPointer();

  absl::Status Delete();

  bool IsDeleted() const;

  PyArray Clone() const;

  // Returns the state needed to copy the array to host, or std::nullopt if
  // the array has no addressable devices. It is safe for the returned object to
  // outlive `this`.
  struct CopyToHostState;
  absl::StatusOr<std::optional<CopyToHostState>> GetCopyToHostState(
      xla::ifrt::Client* client);

  static absl::StatusOr<std::vector<PyArray>> BatchedCopyToDeviceWithSharding(
      absl::Span<const PyArray> py_arrays,
      absl::Span<const xla::ifrt::DeviceListRef> dst_device_lists,
      absl::Span<const nanobind::object> dst_shardings,
      absl::Span<const xla::ifrt::ArrayCopySemantics> array_copy_semantics);

  static absl::StatusOr<PyArray> BatchedDevicePut(
      nanobind::object aval, nanobind::object sharding,
      std::vector<nanobind::object> xs,
      absl::Span<const PyDevice* const> dst_devices, bool committed,
      bool force_copy,
      xla::PjRtClient::HostBufferSemantics host_buffer_semantics,
      bool jax_enable_x64);

  static absl::StatusOr<PyArray> ReorderShards(
      PyArray x, nanobind::object dst_sharding,
      xla::ifrt::ArrayCopySemantics array_copy_semantics);

  static absl::Status BatchedBlockUntilReady(
      std::vector<nanobind::object> objs);

  // Copies a batch of multi-device PyArrays to the host.
  static absl::Status BatchedCopyToHostAsync(nanobind::sequence py_arrays);

  absl::Status ReplaceWithAlias(PyArray o);

 private:
  absl::StatusOr<tsl::Future<>> CopyToHostAsync();

  absl::StatusOr<PyArray> AssertUnsharded(std::string_view api);

  static nanobind::object CheckAndRearrange(absl::Span<const PyArray> py_arrays,
                                            nanobind::object sharding,
                                            nanobind::object aval);

  static absl::Status BatchedCopyToHostAsyncHelper(
      xla::ifrt::Client* client, std::vector<CopyToHostState> array_states);

  Storage& GetStorage();
  const Storage& GetStorage() const;

  inline static PyObject* type_ = nullptr;
};

class PyArrayResultHandler {
 public:
  PyArrayResultHandler(nanobind::object aval, nanobind::object sharding,
                       bool committed,
                       std::vector<nanobind::callable> wrappers = {});

  nanobind::object Call(absl::Span<const PyArray> py_arrays) const;
  nanobind::object Call(PyArray py_array) const;

  nanobind::object Call(nb_class_ptr<PyClient> py_client,
                        xla::ifrt::ArrayRef ifrt_array,
                        xla::Future<> result_status = xla::Future<>()) const;

  const std::vector<nanobind::callable>& wrappers() const { return wrappers_; }

  nanobind::object aval() const { return aval_; }
  nanobind::object sharding() const { return sharding_; }
  bool committed() const { return committed_; }

 private:
  nanobind::object aval_;
  nanobind::object sharding_;
  bool weak_type_;
  bool committed_;

  xla::nb_dtype dtype_;
  std::vector<int64_t> shape_;
  std::vector<nanobind::callable> wrappers_;
};

absl::StatusOr<nanobind::object> CudaArrayInterfaceToBuffer(
    const nanobind::dict& cai, nb_class_ptr<PyClient> cuda_client,
    std::optional<int> device_id);

}  // namespace jax

#endif  // JAXLIB_PY_ARRAY_H_