/*
 * Copyright (c) 2018-2020, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA CORPORATION and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA CORPORATION is strictly prohibited.
 */

#include "CUtils.hpp"
#include <cstring>

using namespace std;

// Log utils
CLogger& CLogger::GetInstance()
{
    static CLogger instance;
    return instance;
}

void CLogger::SetLogLevel(LogLevel level)
{
    m_level = (level > LEVEL_DBG) ? LEVEL_DBG : level;
}

CLogger::LogLevel CLogger::GetLogLevel(void)
{
    return m_level;
}

void CLogger::SetLogStyle(LogStyle style)
{
    m_style = (style > LOG_STYLE_FUNCTION_LINE) ? LOG_STYLE_FUNCTION_LINE
                                                : style;
}

void CLogger::LogLevelMessageVa(LogLevel level, const char *functionName,
                                       uint32_t lineNumber, const char *format,
                                                                    va_list ap)
{
    char str[256] = {'\0',};

    if (level > m_level) {
        return;
    }

    strcpy(str, "nvsipl_camera: ");
    switch (level) {
        case LEVEL_NONE:
            break;
        case LEVEL_ERR:
            strcat(str, "ERROR: ");
            break;
        case LEVEL_WARN:
            strcat(str, "WARNING: ");
            break;
        case LEVEL_INFO:
            break;
        case LEVEL_DBG:
            // Empty
            break;
    }

    vsnprintf(str + strlen(str), sizeof(str) - strlen(str), format, ap);

    if (m_style == LOG_STYLE_NORMAL) {
        if (strlen(str) != 0 && str[strlen(str) - 1] == '\n') {
            strcat(str, "\n");
        }
    } else if (m_style == LOG_STYLE_FUNCTION_LINE) {
        if (strlen(str) != 0 && str[strlen(str) - 1] == '\n') {
            str[strlen(str) - 1] = 0;
        }
        snprintf(str + strlen(str), sizeof(str) - strlen(str), " at %s():%d\n",
                                                     functionName, lineNumber);
    }

    cout << str;
}

void CLogger::LogLevelMessage(LogLevel level, const char *functionName,
                               uint32_t lineNumber, const char *format, ...)
{
    va_list ap;
    va_start(ap, format);
    LogLevelMessageVa(level, functionName, lineNumber, format, ap);
    va_end(ap);
}

void CLogger::LogLevelMessage(LogLevel level, std::string functionName,
                               uint32_t lineNumber, std::string format, ...)
{
    va_list ap;
    va_start(ap, format);
    LogLevelMessageVa(level, functionName.c_str(), lineNumber,
                                                       format.c_str(), ap);
    va_end(ap);
}

void CLogger::LogMessageVa(const char *format, va_list ap)
{
    char str[128] = {'\0',};
    vsnprintf(str, sizeof(str), format, ap);
    cout << str;
}

void CLogger::LogMessage(const char *format, ...)
{
    va_list ap;
    va_start(ap, format);
    LogMessageVa(format, ap);
    va_end(ap);
}

void CLogger::LogMessage(std::string format, ...)
{
    va_list ap;
    va_start(ap, format);
    LogMessageVa(format.c_str(), ap);
    va_end(ap);
}

// Image utils
SIPLStatus CUtils::IsRawImage(NvMediaImage *nvmImg, bool &bIsRaw)
{
    bIsRaw = false;
    // Check if nvmImg is RAW
    NVM_SURF_FMT_DEFINE_ATTR(imgAttr);
    auto nvmStatus = NvMediaSurfaceFormatGetAttrs(nvmImg->type, imgAttr, NVM_SURF_FMT_ATTR_MAX);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        cout << "NvMediaSurfaceFormatGetAttrs failed\n";
        return NVSIPL_STATUS_ERROR;
    }

    if (imgAttr[NVM_SURF_ATTR_SURF_TYPE].value == NVM_SURF_ATTR_SURF_TYPE_RAW) {
        bIsRaw = true;
    }

    return NVSIPL_STATUS_OK;
}

NvMediaImage* CUtils::CreateRgbaImage(NvMediaDevice *nvmDevice, NvMediaImage* imgRaw)
{
    NvMediaSurfAllocAttr allocAttrs[8];
    // Allocate RGBA images
    allocAttrs[0].type = NVM_SURF_ATTR_WIDTH;
    allocAttrs[0].value = imgRaw->width/2;
    allocAttrs[1].type = NVM_SURF_ATTR_HEIGHT;
    allocAttrs[1].value = imgRaw->height/2;
    allocAttrs[2].type = NVM_SURF_ATTR_CPU_ACCESS;
    allocAttrs[2].value = NVM_SURF_ATTR_CPU_ACCESS_UNCACHED;
    auto numallocAttrs = 3;

    NVM_SURF_FMT_DEFINE_ATTR(surfFormatAttrs);
    NVM_SURF_FMT_SET_ATTR_RGBA(surfFormatAttrs,RGBA,UINT,8,PL);
    auto rgbaSurfType = NvMediaSurfaceFormatGetType(surfFormatAttrs, NVM_SURF_FMT_ATTR_MAX);
    auto rgbaImage = NvMediaImageCreateNew(nvmDevice, rgbaSurfType, allocAttrs, numallocAttrs, 0);
    if (rgbaImage == NULL) {
        cout << "NvMediaImageCreateNew failed\n";
        return NULL;
    }
    return rgbaImage;
}

SIPLStatus CUtils::ValidateImage(NvMediaImage *img,
                                 uint32_t expectedSurfFormat,
                                 uint32_t &pixelOrder,
                                 uint32_t &bytesPerPixel)
{
    NVM_SURF_FMT_DEFINE_ATTR(imgAttr);
    auto nvmStatus = NvMediaSurfaceFormatGetAttrs(img->type, imgAttr, NVM_SURF_FMT_ATTR_MAX);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        cout << "NvMediaSurfaceFormatGetAttrs failed\n";
        return NVSIPL_STATUS_ERROR;
    }

    if (imgAttr[NVM_SURF_ATTR_SURF_TYPE].value != expectedSurfFormat) {
        cout << "Surface format does not match expected value\n";
        return NVSIPL_STATUS_ERROR;
    }

    switch (imgAttr[NVM_SURF_ATTR_SURF_TYPE].value) {
        case NVM_SURF_ATTR_SURF_TYPE_RAW:
            if (imgAttr[NVM_SURF_ATTR_DATA_TYPE].value != NVM_SURF_ATTR_DATA_TYPE_UINT) {
                cout << "Unsupported raw input data type\n";
                return NVSIPL_STATUS_ERROR;
            }
            if ((imgAttr[NVM_SURF_ATTR_BITS_PER_COMPONENT].value != NVM_SURF_ATTR_BITS_PER_COMPONENT_10)
                and (imgAttr[NVM_SURF_ATTR_BITS_PER_COMPONENT].value != NVM_SURF_ATTR_BITS_PER_COMPONENT_12)
                and (imgAttr[NVM_SURF_ATTR_BITS_PER_COMPONENT].value != NVM_SURF_ATTR_BITS_PER_COMPONENT_16)) {
                cout << "Unsupported raw input bits per component\n";
                return NVSIPL_STATUS_ERROR;
            }
            bytesPerPixel = 2u;
            break;
        case NVM_SURF_ATTR_SURF_TYPE_RGBA:
            bytesPerPixel = 4u;
            break;
        default:
            cout << "Unsupported surface type\n";
            break;
    }
    pixelOrder = imgAttr[NVM_SURF_ATTR_COMPONENT_ORDER].value;

    return NVSIPL_STATUS_OK;
}

uint8_t* CUtils::CreateImageBuffer(NvMediaImage *img, uint32_t surfFormat)
{
    uint32_t pixelOrder = 0u, bytesPerPixel = 0u;
    auto status = ValidateImage(img, surfFormat, pixelOrder, bytesPerPixel);
    if (status != NVSIPL_STATUS_OK) {
        cout << "ValidateImage failed\n";
        return nullptr;
    }

    // TODO: Try using surface map pointer directly (needs image to be allocated for CPU access)
    auto imgSize = (img->height) * (img->width) * bytesPerPixel;
    if (surfFormat == NVM_SURF_ATTR_SURF_TYPE_RAW) {
        imgSize += img->embeddedDataTopSize;
        imgSize += img->embeddedDataBottomSize;
    }
    uint8_t *pImgBuf = new (std::nothrow) uint8_t[imgSize];
    if (pImgBuf == nullptr) {
        cout << "Failed to allocate memory for image buffer\n";
        return nullptr;
    }
    std::fill(pImgBuf, pImgBuf + imgSize, 0x00);

    return pImgBuf;
}

uint8_t* CUtils::CreateRawBuffer(NvMediaImage *imgSrc)
{
    return CUtils::CreateImageBuffer(imgSrc, NVM_SURF_ATTR_SURF_TYPE_RAW);
}

uint8_t* CUtils::CreateRgbaBuffer(NvMediaImage *imgDst)
{
    return CUtils::CreateImageBuffer(imgDst, NVM_SURF_ATTR_SURF_TYPE_RGBA);
}

/** Function to convert raw NvMediaImage to RGBA NvMediaImage */
SIPLStatus CUtils::ConvertRawToRgba(NvMediaImage *imgSrc,
                                    uint8_t *pSrcBuf,
                                    NvMediaImage *imgDst,
                                    uint8_t *pDstBuf)
{
    uint32_t srcPixelOrder = 0u, srcBytesPerPixel = 0u, dstPixelOrder = 0u, dstBytesPerPixel = 0u;
    auto status = CUtils::ValidateImage(imgSrc,
                                        NVM_SURF_ATTR_SURF_TYPE_RAW,
                                        srcPixelOrder,
                                        srcBytesPerPixel);
    if (status != NVSIPL_STATUS_OK) {
        cout << "ValidateImage failed for source image\n";
        return NVSIPL_STATUS_ERROR;
    }
    if (pSrcBuf == nullptr) {
        cout << "Invalid source image buffer\n";
        return NVSIPL_STATUS_BAD_ARGUMENT;
    }
    status = CUtils::ValidateImage(imgDst,
                                   NVM_SURF_ATTR_SURF_TYPE_RGBA,
                                   dstPixelOrder,
                                   dstBytesPerPixel);
    if (status != NVSIPL_STATUS_OK) {
        cout << "ValidateImage failed for destination image\n";
        return NVSIPL_STATUS_ERROR;
    }
    if (pDstBuf == nullptr) {
        cout << "Invalid destination image buffer\n";
        return NVSIPL_STATUS_BAD_ARGUMENT;
    }
    auto srcPitch = imgSrc->width * srcBytesPerPixel;
    auto dstPitch = imgDst->width * dstBytesPerPixel;
    uint8_t *pSrcBufCpy = pSrcBuf;
    uint8_t *pDstBufCpy = pDstBuf;

    // Read source image
    NvMediaImageSurfaceMap srcSurfMap{};
    auto nvmStatus = NvMediaImageLock(imgSrc, NVMEDIA_IMAGE_ACCESS_READ, &srcSurfMap);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        cout << "NvMediaImageLock failed for source image\n";
        return NVSIPL_STATUS_ERROR;
    }

    nvmStatus = NvMediaImageGetBits(imgSrc, NULL, (void **)&pSrcBufCpy, &srcPitch);
    NvMediaImageUnlock(imgSrc);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        cout << "NvMediaImageGetBits failed\n";
        return NVSIPL_STATUS_ERROR;
    }

    // Do poor man's demosaic
    // Get offsets for each raw component within 2x2 block
    auto xR = 0;
    auto yR = 0;
    auto xG1 = 0;
    auto yG1 = 0;
    auto xG2 = 0;
    auto yG2 = 0;
    auto xB = 0;
    auto yB = 0;

    switch (srcPixelOrder) {
        case NVM_SURF_ATTR_COMPONENT_ORDER_RGGB:
        case NVM_SURF_ATTR_COMPONENT_ORDER_RCCB:
        case NVM_SURF_ATTR_COMPONENT_ORDER_RCCC:
            xR = 0; yR = 0;
            xG1 = 1; yG1 = 0;
            xG2 = 0; yG2 = 1;
            xB = 1; yB = 1;
            break;
        case NVM_SURF_ATTR_COMPONENT_ORDER_GRBG:
        case NVM_SURF_ATTR_COMPONENT_ORDER_CRBC:
        case NVM_SURF_ATTR_COMPONENT_ORDER_CRCC:
            xG1 = 0; yG1 = 0;
            xR = 1; yR = 0;
            xB = 0; yB = 1;
            xG2 = 1; yG2 = 1;
            break;
        case NVM_SURF_ATTR_COMPONENT_ORDER_GBRG:
        case NVM_SURF_ATTR_COMPONENT_ORDER_CBRC:
        case NVM_SURF_ATTR_COMPONENT_ORDER_CCRC:
            xG1 = 0; yG1 = 0;
            xB = 1; yB = 0;
            xR = 0; yR = 1;
            xG2 = 1; yG2 = 1;
            break;
        case NVM_SURF_ATTR_COMPONENT_ORDER_BGGR:
        case NVM_SURF_ATTR_COMPONENT_ORDER_BCCR:
        case NVM_SURF_ATTR_COMPONENT_ORDER_CCCR:
        case NVM_SURF_ATTR_COMPONENT_ORDER_CCCC:
        default:
            xB = 0; yB = 0;
            xG1 = 1; yG1 = 0;
            xG2 = 0; yG2 = 1;
            xR = 1; yR = 1;
            break;
    }

    // Skip embedded lines
    auto embTopLines = imgSrc->embeddedDataTopSize / srcPitch;
    auto embBotLines = imgSrc->embeddedDataBottomSize / srcPitch;

    for (auto y = embTopLines; y < (imgSrc->height - embBotLines); y += 2) {
        for (auto x = 0u; x < imgSrc->width; x += 2) {
            // R
            *pDstBuf++ = pSrcBuf[srcPitch*(y + yR) + 2*(x + xR) + 1];
            // G (average of G1 and G2)
            auto g1 = pSrcBuf[srcPitch*(y + yG1) + 2*(x + xG1) + 1];
            auto g2 = pSrcBuf[srcPitch*(y + yG2) + 2*(x + xG2) + 1];
            *pDstBuf++ = (g1 + g2)/2;
            // B
            *pDstBuf++ = pSrcBuf[srcPitch*(y + yB) + 2*(x + xB) + 1];
            // A
            *pDstBuf++ = 0xFF;
        }
    }

    // Write to destination image
    NvMediaImageSurfaceMap dstSurfMap{};
    nvmStatus = NvMediaImageLock(imgDst, NVMEDIA_IMAGE_ACCESS_WRITE, &dstSurfMap);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        cout << "NvMediaImageLock failed for destination image\n";
        return NVSIPL_STATUS_ERROR;
    }

    nvmStatus = NvMediaImagePutBits(imgDst, NULL, (void **)&pDstBufCpy, &dstPitch);
    NvMediaImageUnlock(imgDst);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        cout << "NvMediaImagePutBits failed\n";
        return NVSIPL_STATUS_ERROR;
    }

    return NVSIPL_STATUS_OK;
}

/* Loads NITO file for given camera module.
 The function assumes the .nito files to be named same as camera module name.
 */
SIPLStatus LoadNITOFile(std::string folderPath,
                        std::string moduleName,
                        std::vector<std::uint8_t>& nito)
{
    // Set up blob file
    string nitoFilePath = (folderPath != "") ? folderPath : "/opt/nvidia/nvmedia/nit/";

    string nitoFile = nitoFilePath + moduleName + ".nito";

    string moduleNameLower{};
    for (auto& c : moduleName) {
        moduleNameLower.push_back(std::tolower(c));
    }
    string nitoFileLower = nitoFilePath + moduleNameLower + ".nito";
    string nitoFileDefault = nitoFilePath + "default.nito";

    // Open NITO file
    auto fp = fopen(nitoFile.c_str(), "rb");
    if (fp == NULL) {
        LOG_INFO("File \"%s\" not found\n", nitoFile.c_str());
        // Try lower case name
        fp = fopen(nitoFileLower.c_str(), "rb");
        if (fp == NULL) {
            LOG_INFO("File \"%s\" not found\n", nitoFileLower.c_str());
            LOG_ERR("Unable to open NITO file for module \"%s\", image quality is not supported!\n", moduleName.c_str());
            return NVSIPL_STATUS_BAD_ARGUMENT;
        } else {
            LOG_MSG("nvsipl_camera: Opened NITO file for module \"%s\"\n", moduleName.c_str());
        }
    } else {
        LOG_MSG("nvsipl_camera: Opened NITO file for module \"%s\"\n", moduleName.c_str());
    }

    // Check file size
    fseek(fp, 0, SEEK_END);
    auto fsize = ftell(fp);
    rewind(fp);

    if (fsize <= 0U) {
        LOG_ERR("NITO file for module \"%s\" is of invalid size\n", moduleName.c_str());
        fclose(fp);
        return NVSIPL_STATUS_BAD_ARGUMENT;
    }

    /* allocate blob memory */
    nito.resize(fsize);

    /* load nito */
    auto result = (long int) fread(nito.data(), 1, fsize, fp);
    if (result != fsize) {
        LOG_ERR("Fail to read data from NITO file for module \"%s\", image quality is not supported!\n", moduleName.c_str());
        nito.resize(0);
        fclose(fp);
        return NVSIPL_STATUS_BAD_ARGUMENT;
    }
    /* close file */
    fclose(fp);

    LOG_INFO("data from NITO file loaded for module \"%s\"\n", moduleName.c_str());

    return NVSIPL_STATUS_OK;
}
