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

#ifndef MEOW_LLIMAGEJ2CNV_H
#define MEOW_LLIMAGEJ2CNV_H

#include "llimagej2c.h"

#include <vector>

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

typedef std::vector<std::vector<U8>> BitStreamData;

#define CHECK_CUDA(call) {                                                                                  \
    cudaError_t _e = (call);                                                                                \
    if (_e != cudaSuccess) {                                                                                \
        LL_ERRS() << "CUDA Runtime failure: '#" << _e << "' at " << __FILE__ << ":" << __LINE__ << LL_ENDL; \
        return false;                                                                                       \
    }                                                                                                       \
}

#define CHECK_NVJ2K(call) {                                                                                 \
    nvjpeg2kStatus_t _e = (call);                                                                           \
    if (_e != NVJPEG2K_STATUS_SUCCESS) {                                                                    \
        LL_ERRS() << "nvJPEG2000 failure: '#" << _e << "' at " << __FILE__ << ":" << __LINE__ << LL_ENDL;   \
        return false;                                                                                       \
    }                                                                                                       \
}

#define CHECK_NVJ2K_NORETURN(call) {                                                                                 \
    nvjpeg2kStatus_t _e = (call);                                                                           \
    if (_e != NVJPEG2K_STATUS_SUCCESS) {                                                                    \
        LL_ERRS() << "nvJPEG2000 failure: '#" << _e << "' at " << __FILE__ << ":" << __LINE__ << LL_ENDL;   \
    }                                                                                                       \
}

class NvImage {

private:
    std::vector<void*> m_pixelDataDevice;
    std::vector<void*> m_pixelDataHost;
    std::vector<size_t> m_pitchBytesDevice;
    std::vector<size_t> m_pitchBytesHost;

    std::vector<size_t> m_pixelDataSize;
    U32 m_capacity;

    nvjpeg2kImage_t m_imageDevice;
    nvjpeg2kImage_t m_imageHost;

    nvjpeg2kImageInfo_t m_info;
    std::vector<nvjpeg2kImageComponentInfo_t> m_compInfo;
    nvjpeg2kColorSpace_t m_colorSpace;

    void clear();

public:
    nvjpeg2kColorSpace_t getColorSpace();
    nvjpeg2kImage_t& getImageHost();
    nvjpeg2kImage_t& getImageDevice();
    nvjpeg2kImageInfo_t& getImageInfo();
    nvjpeg2kImageComponentInfo_t* getComponentInfo();

    NvImage();

    bool init(nvjpeg2kImageInfo_t& imageInfo, nvjpeg2kImageComponentInfo_t* imageCompInfo, nvjpeg2kColorSpace_t colorSpace);
    bool cleanup();
    bool copyToDevice();
};

class LLNVJ2KEncoder {

public:
    ~LLNVJ2KEncoder();

    bool createEncoder();

    bool encode(const LLImageRaw &rawImageIn, LLImageJ2C &compressedImageOut);

private:
    void cleanup();

    bool setImage(const LLImageRaw& raw);

    nvjpeg2kEncoder_t m_encoderHandle = nullptr;
    nvjpeg2kEncodeState_t m_encoderState = nullptr;
    nvjpeg2kEncodeParams_t m_encoderParams = nullptr;

    NvImage m_nvImage;

    bool m_encoderCreated = false;
};

class LLImageJ2CNV : public LLImageJ2CImpl {

public:
    LLImageJ2CNV();
    virtual ~LLImageJ2CNV();
protected:
    // Find out the image size and number of channels.
    // Return value:
    // true: image size and number of channels was determined
    // false: error on decode
    virtual bool getMetadata(LLImageJ2C &base);
    // Decode the raw image optionally aborting (to continue later) after
    // decode_time seconds.  Decode at most max_channel_count and start
    // decoding channel first_channel.
    // Return value:
    // true: decoding complete (even if it failed)
    // false: time expired while decoding
    virtual bool decodeImpl(LLImageJ2C &base, LLImageRaw &raw_image, F32 decode_time, S32 first_channel, S32 max_channel_count);
    virtual bool encodeImpl(LLImageJ2C &base, const LLImageRaw &raw_image, const char* comment_text, F32 encode_time=0.0,
                                bool reversible = false);
    virtual bool initDecode(LLImageJ2C &base, LLImageRaw &raw_image, int discard_level = -1, int* region = NULL);
    virtual bool initEncode(LLImageJ2C &base, LLImageRaw &raw_image, int blocks_size = -1, int precincts_size = -1, int levels = 0);
    virtual std::string getEngineInfo() const;

};

#endif // MEOW_LLIMAGEJ2CNV_H