#include "augment_renderer.h"
#include "cuda_helper.hpp"

#include <cuda_gl_interop.h>

#include <glad/gl.h>

struct augment_renderer::impl {
    const cv::cuda::GpuMat *bg_img = nullptr;
    GLuint bg_tex = 0, bg_pbo = 0;
    cudaGraphicsResource *bg_res = nullptr;

    texture_renderer *tex_renderer = nullptr;
    const render_config *config = nullptr;

    ~impl() {
        cudaGraphicsUnregisterResource(bg_res);
    }

    bool initialize() {
        // generate and allocate pixel buffer
        glGenBuffers(1, &bg_pbo);
        glBindBuffer(GL_PIXEL_UNPACK_BUFFER, bg_pbo);
        glBufferStorage(GL_PIXEL_UNPACK_BUFFER, rgb_image_size, nullptr, GL_DYNAMIC_STORAGE_BIT);
        glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);

        // generate and allocate texture
        glGenTextures(1, &bg_tex);
        glBindTexture(GL_TEXTURE_2D, bg_tex);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
        glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
        glTexStorage2D(GL_TEXTURE_2D, 1, GL_RGBA8, image_width, image_height);

        // register background pbo
        CUDA_API_CHECK(cudaGraphicsGLRegisterBuffer(&bg_res, bg_pbo,
                                                    cudaGraphicsRegisterFlagsWriteDiscard));

        return true;
    }

    bool render_background() {
        if (bg_img == nullptr) return true;

        // upload background image to pbo
        void *pbo_ptr;
        size_t pbo_size;
        CUDA_API_CHECK(cudaGraphicsMapResources(1, &bg_res));
        CUDA_API_CHECK(cudaGraphicsResourceGetMappedPointer(&pbo_ptr, &pbo_size, bg_res));
        assert(pbo_size == rgb_image_size);
        CUDA_API_CHECK(cudaMemcpy2D(pbo_ptr, rgb_image_pitch, bg_img->cudaPtr(),
                                    bg_img->step, rgb_image_pitch, image_height, cudaMemcpyDeviceToDevice));
        CUDA_API_CHECK(cudaGraphicsUnmapResources(1, &bg_res));

        // unpack pbo to texture
        glBindBuffer(GL_PIXEL_UNPACK_BUFFER, bg_pbo);
        glBindTexture(GL_TEXTURE_2D, bg_tex);
        glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, image_width, image_height, GL_BGR, GL_UNSIGNED_BYTE, nullptr);
        glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);

        // render texture
        assert(tex_renderer != nullptr);
        assert(config != nullptr);
        texture_renderer::render_config tex_config{
                .tex = bg_tex,
                .x = config->x,
                .y = config->y,
                .width = config->width,
                .height = config->height
        };
        tex_renderer->render(&tex_config);

        return true;
    }

    bool render() {
        CALL_CHECK(render_background());
        return true;
    }

};

augment_renderer::augment_renderer()
        : pimpl(std::make_unique<impl>()) {}

augment_renderer::~augment_renderer() = default;

bool augment_renderer::initialize(texture_renderer *renderer) {
    pimpl->tex_renderer = renderer;
    return pimpl->initialize();
}

void augment_renderer::set_background(const cv::cuda::GpuMat *background) {
    assert(background->rows == image_height);
    assert(background->cols == image_width);
    pimpl->bg_img = background;
}

bool augment_renderer::render(const render_config &config) {
    pimpl->config = &config;
    return pimpl->render();
}