#pragma once

#include <cstdint>

#include <memory>

#include "arrow/buffer.h"

#include "arrow/io/concurrency.h"

#include "arrow/type_fwd.h"

namespace arrow {

namespace cuda {

class CudaContext;

class CudaIpcMemHandle;

class ARROW_EXPORT CudaBuffer : public Buffer {

public:

// XXX deprecate?

CudaBuffer(uint8_t* data, int64_t size, const std::shared_ptr<CudaContext>& context,

bool own_data = false, bool is_ipc = false);

CudaBuffer(uintptr_t address, int64_t size, const std::shared_ptr<CudaContext>& context,

bool own_data = false, bool is_ipc = false);

CudaBuffer(const std::shared_ptr<CudaBuffer>& parent, const int64_t offset,

const int64_t size);

~CudaBuffer();

/// \brief Convert back generic buffer into CudaBuffer

/// \param[in] buffer buffer to convert

/// \return CudaBuffer or Status

///

/// \note This function returns an error if the buffer isn't backed

/// by GPU memory

static Result<std::shared_ptr<CudaBuffer>> FromBuffer(std::shared_ptr<Buffer> buffer);

/// \brief Copy memory from GPU device to CPU host

/// \param[in] position start position inside buffer to copy bytes from

/// \param[in] nbytes number of bytes to copy

/// \param[out] out start address of the host memory area to copy to

/// \return Status

Status CopyToHost(const int64_t position, const int64_t nbytes, void* out) const;

/// \brief Copy memory to device at position

/// \param[in] position start position to copy bytes to

/// \param[in] data the host data to copy

/// \param[in] nbytes number of bytes to copy

/// \return Status

Status CopyFromHost(const int64_t position, const void* data, int64_t nbytes);

/// \brief Copy memory from device to device at position

/// \param[in] position start position inside buffer to copy bytes to

/// \param[in] data start address of the device memory area to copy from

/// \param[in] nbytes number of bytes to copy

/// \return Status

///

/// \note It is assumed that both source and destination device

/// memories have been allocated within the same context.

Status CopyFromDevice(const int64_t position, const void* data, int64_t nbytes);

/// \brief Copy memory from another device to device at position

/// \param[in] src_ctx context of the source device memory

/// \param[in] position start position inside buffer to copy bytes to

/// \param[in] data start address of the another device memory area to copy from

/// \param[in] nbytes number of bytes to copy

/// \return Status

Status CopyFromAnotherDevice(const std::shared_ptr<CudaContext>& src_ctx,

const int64_t position, const void* data, int64_t nbytes);

/// \brief Expose this device buffer as IPC memory which can be used in other processes

/// \return Handle or Status

///

/// \note After calling this function, this device memory will not be freed

/// when the CudaBuffer is destructed

virtual Result<std::shared_ptr<CudaIpcMemHandle>> ExportForIpc();

const std::shared_ptr<CudaContext>& context() const { return context_; }

protected:

std::shared_ptr<CudaContext> context_;

bool own_data_;

bool is_ipc_;

virtual Status Close();

};

class ARROW_EXPORT CudaHostBuffer : public MutableBuffer {

public:

using MutableBuffer::MutableBuffer;

~CudaHostBuffer();

/// \brief Return a device address the GPU can read this memory from.

Result<uintptr_t> GetDeviceAddress(const std::shared_ptr<CudaContext>& ctx);

};

class ARROW_EXPORT CudaIpcMemHandle {

public:

~CudaIpcMemHandle();

/// \brief Create CudaIpcMemHandle from opaque buffer (e.g. from another process)

/// \param[in] opaque_handle a CUipcMemHandle as a const void*

/// \return Handle or Status

static Result<std::shared_ptr<CudaIpcMemHandle>> FromBuffer(const void* opaque_handle);

/// \brief Write CudaIpcMemHandle to a Buffer

/// \param[in] pool a MemoryPool to allocate memory from

/// \return Buffer or Status

Result<std::shared_ptr<Buffer>> Serialize(

MemoryPool* pool = default_memory_pool()) const;

private:

explicit CudaIpcMemHandle(const void* handle);

CudaIpcMemHandle(int64_t memory_size, const void* cu_handle);

struct CudaIpcMemHandleImpl;

std::unique_ptr<CudaIpcMemHandleImpl> impl_;

const void* handle() const;

int64_t memory_size() const;

friend CudaBuffer;

friend CudaContext;

};

class ARROW_EXPORT CudaBufferReader

: public ::arrow::io::internal::RandomAccessFileConcurrencyWrapper<CudaBufferReader> {

public:

explicit CudaBufferReader(const std::shared_ptr<Buffer>& buffer);

bool closed() const override;

bool supports_zero_copy() const override;

std::shared_ptr<CudaBuffer> buffer() const { return buffer_; }

protected:

friend ::arrow::io::internal::RandomAccessFileConcurrencyWrapper<CudaBufferReader>;

Status DoClose();

Result<int64_t> DoRead(int64_t nbytes, void* buffer);

Result<std::shared_ptr<Buffer>> DoRead(int64_t nbytes);

Result<int64_t> DoReadAt(int64_t position, int64_t nbytes, void* out);

Result<std::shared_ptr<Buffer>> DoReadAt(int64_t position, int64_t nbytes);

Result<int64_t> DoTell() const;

Status DoSeek(int64_t position);

Result<int64_t> DoGetSize();

Status CheckClosed() const {

}

std::shared_ptr<CudaBuffer> buffer_;

std::shared_ptr<CudaContext> context_;

const uintptr_t address_;

int64_t size_;

int64_t position_;

bool is_open_;

};

class ARROW_EXPORT CudaBufferWriter : public io::WritableFile {

public:

explicit CudaBufferWriter(const std::shared_ptr<CudaBuffer>& buffer);

~CudaBufferWriter() override;

/// \brief Close writer and flush buffered bytes to GPU

Status Close() override;

bool closed() const override;

/// \brief Flush buffered bytes to GPU

Status Flush() override;

Status Seek(int64_t position) override;

Status Write(const void* data, int64_t nbytes) override;

Status WriteAt(int64_t position, const void* data, int64_t nbytes) override;

Result<int64_t> Tell() const override;

/// \brief Set CPU buffer size to limit calls to cudaMemcpy

/// \param[in] buffer_size the size of CPU buffer to allocate

/// \return Status

///

/// By default writes are unbuffered

Status SetBufferSize(const int64_t buffer_size);

/// \brief Returns size of host (CPU) buffer, 0 for unbuffered

int64_t buffer_size() const;

/// \brief Returns number of bytes buffered on host

int64_t num_bytes_buffered() const;

private:

class CudaBufferWriterImpl;

std::unique_ptr<CudaBufferWriterImpl> impl_;

};

Result<std::shared_ptr<CudaHostBuffer>> AllocateCudaHostBuffer(int device_number,

const int64_t size);

Result<uintptr_t> GetDeviceAddress(const uint8_t* cpu_data,

const std::shared_ptr<CudaContext>& ctx);

Result<uint8_t*> GetHostAddress(uintptr_t device_ptr);

} // namespace cuda

} // namespace arrow