/**
 * @file llimagej2cnv.cpp
 * @brief This is an implementation of JPEG2000 encode/decode using nvJPEG2000.
 */

#include "linden_common.h"
#include "llimagej2cnv.h"

#include <cuda_runtime_api.h>
#include <nvjpeg2k.h>
#include <nvjpeg2k_version.h>

// Factory function: see declaration in llimagej2c.cpp
LLImageJ2CImpl* fallbackCreateLLImageJ2CImpl() {
    return new LLImageJ2CNV();
}

#pragma region NVIMAGE

void NvImage::clear() {
    m_imageDevice.num_components = 0;
    m_imageDevice.pixel_data = nullptr;
    m_imageDevice.pitch_in_bytes = nullptr;
    m_imageDevice.pixel_type = NVJPEG2K_UINT8;

    m_imageHost.num_components = 0;
    m_imageHost.pixel_data = nullptr;
    m_imageHost.pitch_in_bytes = nullptr;
    m_imageHost.pixel_type = NVJPEG2K_UINT8;

    m_capacity = 0;
}

nvjpeg2kColorSpace_t NvImage::getColorSpace() {
    return m_colorSpace;
}

nvjpeg2kImage_t& NvImage::getImageHost() {
    return m_imageHost;
}

nvjpeg2kImage_t& NvImage::getImageDevice() {
    return m_imageDevice;
}

nvjpeg2kImageInfo_t& NvImage::getImageInfo() {
    return m_info;
}

nvjpeg2kImageComponentInfo_t* NvImage::getComponentInfo() {
    return m_compInfo.data();
}

NvImage::NvImage() {
    clear();
}

bool NvImage::init(nvjpeg2kImageInfo_t& imageInfo, nvjpeg2kImageComponentInfo_t* compInfo, nvjpeg2kColorSpace_t colorSpace) {
    memcpy(&m_info, &imageInfo, sizeof(imageInfo));
    m_compInfo.resize(m_info.num_components);
    m_colorSpace = colorSpace;

    for (U32 c = 0; c < m_info.num_components; c++) {
        memcpy(&m_compInfo[c], &compInfo[c], sizeof(m_compInfo[c]));
    }

    if (m_info.num_components > m_capacity) {
        m_pixelDataDevice.resize(m_info.num_components, nullptr);
        m_pitchBytesDevice.resize(m_info.num_components, 0);
        m_pixelDataHost.resize(m_info.num_components, nullptr);
        m_pitchBytesHost.resize(m_info.num_components, 0);
        m_pixelDataSize.resize(m_info.num_components, 0);
        m_capacity = m_info.num_components;
    }

    m_imageDevice.pixel_data = m_pixelDataDevice.data();
    m_imageDevice.pitch_in_bytes = m_pitchBytesDevice.data();
    m_imageHost.pixel_data = m_pixelDataHost.data();
    m_imageHost.pitch_in_bytes = m_pitchBytesHost.data();

    if (m_compInfo[0].precision <= 8) {
        m_imageDevice.pixel_type = NVJPEG2K_UINT8;
        m_imageHost.pixel_type = NVJPEG2K_UINT8;
    } else if (m_compInfo[0].precision > 8 && m_compInfo[0].precision <= 16) {
        m_imageDevice.pixel_type = NVJPEG2K_UINT16;
        m_imageHost.pixel_type = NVJPEG2K_UINT16;
    } else {
        LL_ERRS() << "nvJPEG2000 failed to initialise image: precision too large" << LL_ENDL;
        return false;
    }

    m_imageDevice.num_components = m_info.num_components;
    m_imageHost.num_components = m_info.num_components;

    U8 bytesPerElement = (m_imageDevice.pixel_type == NVJPEG2K_UINT16) ? 2 : 1;

    for (U32 c = 0; c < m_info.num_components; c++) {
        m_imageDevice.pitch_in_bytes[c] =
        m_imageHost.pitch_in_bytes[c] = m_compInfo[c].component_width * bytesPerElement;
        size_t compSize = m_compInfo[c].component_height * m_imageDevice.pitch_in_bytes[c];

        if (compSize > m_pixelDataSize[c]) {
            if (m_imageDevice.pixel_data[c]) {
                CHECK_CUDA(cudaFree(m_imageDevice.pixel_data[c]));
            }

            if (m_imageHost.pixel_data[c]) {
                free(m_imageHost.pixel_data[c]);
            }

            m_pixelDataSize[c] = compSize;
            CHECK_CUDA(cudaMalloc(&m_imageDevice.pixel_data[c], compSize));
            m_imageHost.pixel_data[c] = malloc(compSize);
        }
    }

    return true;
}

bool NvImage::cleanup() {
    for (U32 c = 0; c < m_capacity; c++) {
        if (m_imageDevice.pixel_data[c]) {
            CHECK_CUDA(cudaFree(m_pixelDataDevice[c]));
            m_pixelDataDevice[c] = nullptr;
        }

        if (m_imageHost.pixel_data[c]) {
            free(m_pixelDataHost[c]);
            m_pixelDataHost[c] = nullptr;
        }

        m_pixelDataSize[c] = 0;
    }

    return true;
}

bool NvImage::copyToDevice() {
    U8 bytesPerElement = m_imageDevice.pixel_type == NVJPEG2K_UINT16 ? 2 : 1;

    for (U32 c = 0; c < m_imageDevice.num_components; c++) {
        CHECK_CUDA(cudaMemcpy2D(
            m_imageDevice.pixel_data[c],
            m_imageDevice.pitch_in_bytes[c],
            m_imageHost.pixel_data[c],
            m_imageHost.pitch_in_bytes[c],
            m_compInfo[c].component_width * bytesPerElement,
            m_compInfo[c].component_height, cudaMemcpyHostToDevice));

        CHECK_CUDA(cudaDeviceSynchronize());
    }

    return true;
}

#pragma endregion NVIMAGE

#pragma region ENCODER

bool LLNVJ2KEncoder::createEncoder() {
    cleanup();

    CHECK_NVJ2K(nvjpeg2kEncoderCreateSimple(&m_encoderHandle));
    CHECK_NVJ2K(nvjpeg2kEncodeStateCreate(m_encoderHandle, &m_encoderState));
    CHECK_NVJ2K(nvjpeg2kEncodeParamsCreate(&m_encoderParams));

    m_encoderCreated = true;
}

void LLNVJ2KEncoder::cleanup() {
    if (m_encoderParams != nullptr) {
        CHECK_NVJ2K_NORETURN(nvjpeg2kEncodeParamsDestroy(m_encoderParams));
    }
    m_encoderParams = nullptr;

    if (m_encoderState != nullptr) {
        CHECK_NVJ2K_NORETURN(nvjpeg2kEncodeStateDestroy(m_encoderState));
    }
    m_encoderState = nullptr;

    if (m_encoderHandle != nullptr) {
        CHECK_NVJ2K_NORETURN(nvjpeg2kEncoderDestroy(m_encoderHandle));
    }
    m_encoderHandle = nullptr;

    m_encoderCreated = false;
}

LLNVJ2KEncoder::~LLNVJ2KEncoder() {
    cleanup();
}

bool LLNVJ2KEncoder::encode(const LLImageRaw &rawImageIn, LLImageJ2C &compressedImageOut) {
    LLImageDataSharedLock lockIn(&rawImageIn);
    LLImageDataLock lockOut(&compressedImageOut);

    if (!m_encoderCreated) {
        if (!createEncoder()) {
            return false;
        }
    }

    if (!setImage(rawImageIn)) {
        return false;
    }

    nvjpeg2kEncodeConfig_t encodeCfg;

    encodeCfg.stream_type = NVJPEG2K_STREAM_J2K;
    encodeCfg.color_space = m_nvImage.getColorSpace();
    encodeCfg.rsiz = 0;
    encodeCfg.image_width = m_nvImage.getImageInfo().image_width;
    encodeCfg.image_height = m_nvImage.getImageInfo().image_height;
    encodeCfg.prog_order = NVJPEG2K_RLCP;
}

bool LLNVJ2KEncoder::setImage(const LLImageRaw& raw) {
    S8 numcomps = raw.getComponents();
    U16 width = raw.getWidth();
    U16 height = raw.getHeight();

    nvjpeg2kColorSpace_t colorSpace = NVJPEG2K_COLORSPACE_SRGB;

    nvjpeg2kImageInfo_t imageInfo;
    std::vector<nvjpeg2kImageComponentInfo_t> compInfo(numcomps);

    imageInfo.num_components = numcomps;
    imageInfo.image_width = width;
    imageInfo.image_height = height;

    for (auto &comp : compInfo) {
        comp.component_width = width;
        comp.component_height = height;
        comp.precision = 8;
        comp.sgn = 0;
    }

    if (!m_nvImage.init(imageInfo, compInfo.data(), colorSpace)) {
        return false;
    }

    const U8 *pSrcData = raw.getData();
    auto &dstImage = m_nvImage.getImageHost();

    // ????????????????
    S32 i = 0;
    for (S32 y = height - 1; y >= 0; y--) {
        for (S32 x = 0; x < width; x++) {
            const U8 *pixel = pSrcData + (y * width + x) * numcomps;
            for (S32 c = 0; c < numcomps; c++) {
                // Absolute fuckery ahead, viewer discretion advised
                U8 *pixBuf = static_cast<U8*>(dstImage.pixel_data[c]);
                pixBuf[i] = *pixel;
                pixel++;
            }
            i++;
        }
    }
}

#pragma endregion ENCODER

std::string LLImageJ2CNV::getEngineInfo() const {
    cudaDeviceProp props;
    int dev = 0;
    cudaGetDevice(&dev);
    cudaGetDeviceProperties(&props, dev);
    return llformat("nvJPEG2000, Version: %i.%i.%i.%i, GPU: %s, CC: %i.%i",
        NVJPEG2K_VER_MAJOR, NVJPEG2K_VER_MINOR, NVJPEG2K_VER_PATCH, NVJPEG2K_VER_BUILD,
        props.name, props.major, props.minor);
}

LLImageJ2CNV::LLImageJ2CNV() : LLImageJ2CImpl() {}

LLImageJ2CNV::~LLImageJ2CNV() {}

bool LLImageJ2CNV::initEncode(LLImageJ2C &base, LLImageRaw &raw_image, int blocks_size = -1, int precincts_size = -1, int levels = 0) {
    // TODO: impl
    return false;
}

bool LLImageJ2CNV::initDecode(LLImageJ2C &base, LLImageRaw &raw_image, int discard_level = -1, int* region = NULL) {
    // TODO: impl
    return false;
}

bool LLImageJ2CNV::encodeImpl(LLImageJ2C &base, const LLImageRaw &raw_image, const char* comment_text, F32 encode_time=0.0, bool reversible = false) {
    if (raw_image.isBufferInvalid()) {
        base.setLastError("Invalid input, no buffer");
        return false;
    }

    LLNVJ2KEncoder encoder;
    if (!encoder.createEncoder()) {
        LL_WARNS() << "Failed to create nvJPEG2000 encoder." << LL_ENDL;
        return false;
    }
    bool encoded = encoder.encode(raw_image, base);
    if (!encoded) {
        LL_WARNS() << "nvJPEG2000 encoding was unsuccessful, returning false." << LL_ENDL;
    }
    return encoded;
}

bool LLImageJ2CNV::decodeImpl(LLImageJ2C &base, LLImageRaw &raw_image, F32 decode_time, S32 first_channel, S32 max_channel_count) {
    // TODO: impl
    return false;
}

bool LLImageJ2CNV::getMetadata(LLImageJ2C &base) {
    // TODO: impl
    return false;
}