#pragma once
#include <cuda_runtime.h>
#include "../common.h"
#include "error.h"

namespace utils::cuda {
    class Resource {
    public:
        virtual ~Resource() {}

        virtual void *getBuffer() const = 0;

        virtual size_t size() const = 0;
    };

    class BufferResource : public Resource {
    public:
        BufferResource(void *buffer = nullptr, size_t size = 0)
            : _p_buffer(buffer), _ownBuffer(false), _size(size) {}
        BufferResource(size_t size) : _ownBuffer(true), _size(size) {
            void *p_buffer;
            THROW_IF_FAILED(cudaMalloc(&p_buffer, size));
            _p_buffer = std::shared_ptr<void>(p_buffer);
        }

        virtual ~BufferResource() {
            if (!_ownBuffer || _p_buffer.use_count() > 1)
                return;
            try {
                THROW_IF_FAILED(cudaFree(_p_buffer.get()));
            } catch (std::exception &ex) {
                common::Logger::instance.warning("Exception raised in destructor: %s", ex.what());
            }
        }

        virtual void *getBuffer() const { return _p_buffer.get(); }

        virtual size_t size() const { return _size; }

    private:
        std::shared_ptr<void> _p_buffer;
        bool _ownBuffer;
        size_t _size;
    };

    class GraphicsResource : public Resource {
    public:
        cudaGraphicsResource_t getHandler() { return _res; }

        virtual ~GraphicsResource() {
            if (_res == nullptr)
                return;
            try {
                THROW_IF_FAILED(cudaGraphicsUnregisterResource(_res));
            } catch (std::exception &ex) {
                common::Logger::instance.warning("Exception raised in destructor: %s", ex.what());
            }
            _res = nullptr;
        }

        virtual size_t size() const { return _size; }

    protected:
        cudaGraphicsResource_t _res;
        size_t _size;

        GraphicsResource() : _res(nullptr), _size(0) {}
    };

} // namespace utils::cuda