// Copyright (c) 2022, Edgeware AB. All rights reserved.

#pragma once

#include <cstddef>
#include <cstdint>

#include <cuda_runtime.h>

class DeviceMemory {
public:
    DeviceMemory() = default;

    explicit DeviceMemory(size_t numberOfBytes);

    explicit DeviceMemory(size_t numberOfBytes, cudaStream_t cudaStream);

    explicit DeviceMemory(void* deviceMemory) noexcept;

    bool allocateMemory(size_t numberOfBytes);

    bool allocateMemoryAsync(size_t numberOfBytes, cudaStream_t cudaStream);

    bool reallocateMemory(size_t numberOfBytes);

    bool allocateAndResetMemory(size_t numberOfBytes);

    bool allocateAndResetMemoryAsync(size_t numberOfBytes, cudaStream_t cudaStream);

    bool freeMemory();

    bool freeMemoryAsync(cudaStream_t cudaStream);

    void setFreeingCudaStream(cudaStream_t cudaStream);

    ~DeviceMemory();

    template <typename T = uint8_t> [[nodiscard]] T* getDevicePointer() const {
        return reinterpret_cast<T*>(dMemory);
    }

    [[nodiscard]] size_t getSize() const;

    DeviceMemory(DeviceMemory&& other) noexcept;
    DeviceMemory& operator=(DeviceMemory&& other) noexcept;

    void swap(DeviceMemory& other) noexcept;

    DeviceMemory(DeviceMemory const&) = delete;
    DeviceMemory operator=(DeviceMemory const&) = delete;

private:
    uint8_t* dMemory = nullptr;
    size_t mAllocatedBytes = 0;
    cudaStream_t mCudaStream = nullptr; // Stream to use for Async free
};