#include "xla/pjrt/pjrt_c_api_client.h"
#include "tensorflow/core/framework/tensor.h"

namespace tensorflow {

StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor);

absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
                                    xla::PjRtCApiClient* c_api_client,
                                    Tensor* tensor);

absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(

  tensorflow::Tensor tensor(DT_FLOAT, {1});
  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
  PJRT_Buffer* c_buffer = CreateCBuffer();

  TF_EXPECT_OK(SetPjRtCBufferToTensor(
      c_buffer, down_cast<xla::PjRtCApiClient*>(pjrt_client.get()), &tensor));
}

TEST(TensorPjRtBufferUtilTest, SetPjRtCBufferToTensorSuccess) {
  auto allocator = std::make_unique<AsyncValueAllocator>();

  if (c_api_buffer == nullptr) {
    return absl::InternalError(
        "The PjRtBuffer in the tensor is not type PjRtCApiBuffer.");
  }
  return c_api_buffer->c_buffer();
}

absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
                                    xla::PjRtCApiClient* c_api_client,
                                    Tensor* tensor) {

      tensorflow::GetPjRtCApiClient(tensorflow::DeviceType(device_type));
  if (!pjrt_c_api_client.ok()) {
    status->status = pjrt_c_api_client.status();
    return;
  }
  auto set_buffer_status =
      SetPjRtCBufferToTensor(c_buffer, *pjrt_c_api_client, &tensor);
  status->status = set_buffer_status;