#pragma once

#include <cuda_gl_interop.h>
#include <cuda_helper.hpp>

class Interop {
public:
    void init() {
        cudaStreamCreate(&stream);
    }

    void destroy() {
        cudaStreamDestroy(stream);
        unregisterImage();
    }

    void map() {
        cudaArray_t array;
        CUDA::checkCudaError(cudaGraphicsMapResources(1, &graphic_resource, stream));
        CUDA::checkCudaError(cudaGraphicsSubResourceGetMappedArray(&array, graphic_resource, 0, 0));
        cudaResourceDesc surfRes = cudaResourceDesc();
        surfRes.resType = cudaResourceTypeArray;
        surfRes.res.array.array = array;
        CUDA::checkCudaError(cudaCreateSurfaceObject(&surface, &surfRes));
        mapped = true;
    }

    void unmap() {
        if (mapped) {
            CUDA::checkCudaError(cudaGraphicsUnmapResources(1, &graphic_resource, stream));
            CUDA::checkCudaError(cudaDestroySurfaceObject(surface));
        }
        mapped = false;
    }

    void registerImage(unsigned int buffer, int target) {
        CUDA::checkCudaError(cudaGraphicsGLRegisterImage(&graphic_resource, buffer, target,
                                                         cudaGraphicsRegisterFlagsSurfaceLoadStore));
    }

    void unregisterImage() {
        if (mapped) {
            unmap();
        }
        if (graphic_resource != nullptr) {
            CUDA::checkCudaError(cudaGraphicsUnregisterResource(graphic_resource));
        }
        graphic_resource = nullptr;
        surface = 0;
    }

    [[nodiscard]] cudaSurfaceObject_t getImage() const {
        return surface;
    }

private:
    bool mapped = false;
    cudaGraphicsResource* graphic_resource = nullptr;
    cudaSurfaceObject_t surface = 0;
    cudaStream_t stream = nullptr;
};
