/*
 * 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.
 */

/* STL Headers */
#include <unistd.h>
#include <cstring>
#include <iostream>
#include <csignal>
#include <thread>
#include <chrono>
#include <ctime>
#include <getopt.h>

/* NvSIPL Header for INvSIPLNvMBuffer to Flip */
#include "NvSIPLClient.hpp" // Client

#ifndef CFILEWRITER_HPP
#define CFILEWRITER_HPP

#define MAX_NUM_SURFACES (3u)

using namespace std;
using namespace nvsipl;

/** FileWriter class */
class CFileWriter
{
 public:
    SIPLStatus Init(string sFilename)
    {
        m_pOutFile = nullptr;
        m_pOutFile = fopen(sFilename.c_str(), "wb");
        if (!m_pOutFile) {
            LOG_ERR("Failed to create output file\n");
            return NVSIPL_STATUS_ERROR;
        }
        return NVSIPL_STATUS_OK;
    }

    SIPLStatus WriteBufferToFile(INvSIPLClient::INvSIPLNvMBuffer *pNvmBuf)
    {
        NvMediaImage *imagePtr = pNvmBuf->GetImage();
        if (imagePtr == nullptr) {
            LOG_ERR("GetImage failed\n");
            return NVSIPL_STATUS_BAD_ARGUMENT;
        }

        NvMediaImageSurfaceMap surfaceMap;
        NvMediaStatus nvmStatus = NvMediaImageLock(imagePtr,
                                                   NVMEDIA_IMAGE_ACCESS_READ,
                                                   &surfaceMap);
        if (nvmStatus != NVMEDIA_STATUS_OK) {
            LOG_ERR("NvMediaImageLock failed\n");
            return NVSIPL_STATUS_ERROR;
        }
        NvMediaImageUnlock(imagePtr);

        float *xScalePtr = nullptr, *yScalePtr = nullptr;
        uint32_t *bytesPerPixelPtr = nullptr;
        uint32_t numSurfaces = -1u;
        SIPLStatus status = GetSurfParams(imagePtr->type,
                                          &xScalePtr,
                                          &yScalePtr,
                                          &bytesPerPixelPtr,
                                          &numSurfaces);
        if (status != NVSIPL_STATUS_OK) {
            LOG_ERR("GetSurfParams failed\n");
            return status;
        }

        uint32_t pBuffPitches[MAX_NUM_SURFACES] = { 0u };
        uint8_t *pBuff[MAX_NUM_SURFACES] = { nullptr };
        uint32_t size[MAX_NUM_SURFACES] = { 0u };
        uint32_t imageSize = 0u;
        for (uint32_t i = 0u; i < numSurfaces; i++) {
            size[i] = (surfaceMap.width * xScalePtr[i] * surfaceMap.height * yScalePtr[i] * bytesPerPixelPtr[i]);
            imageSize += size[i];
            pBuffPitches[i] = (uint32_t)((float)surfaceMap.width * xScalePtr[i]) * bytesPerPixelPtr[i];
        }

        // Add size for embedded data
        size[0] += imagePtr->embeddedDataTopSize;
        size[0] += imagePtr->embeddedDataBottomSize;
        imageSize += imagePtr->embeddedDataTopSize;
        imageSize += imagePtr->embeddedDataBottomSize;

        uint8_t *buff = new (std::nothrow) uint8_t[imageSize];
        if (buff == nullptr) {
            LOG_ERR("Out of memory\n");
            return NVSIPL_STATUS_OUT_OF_MEMORY;
        }
        std::fill(buff, buff + imageSize, 0x00);

        uint8_t *buffIter = buff;
        for (uint32_t i = 0u; i < numSurfaces; i++) {
            pBuff[i] = buffIter;
            buffIter += (uint32_t)(surfaceMap.height * yScalePtr[i] * pBuffPitches[i]);
        }

        nvmStatus = NvMediaImageLock(imagePtr, NVMEDIA_IMAGE_ACCESS_READ, &surfaceMap);
        if (nvmStatus != NVMEDIA_STATUS_OK) {
            LOG_ERR("NvMediaImageLock failed\n");
            return NVSIPL_STATUS_ERROR;
        }

        nvmStatus = NvMediaImageGetBits(imagePtr, nullptr, (void **)pBuff, pBuffPitches);
        NvMediaImageUnlock(imagePtr);
        if (nvmStatus != NVMEDIA_STATUS_OK) {
            LOG_ERR("NvMediaImageGetBits failed\n");
            return NVSIPL_STATUS_ERROR;
        }

        for (uint32_t i = 0u; i < numSurfaces; i++) {
            if (fwrite(pBuff[i], size[i], 1u, m_pOutFile) != 1u) {
                LOG_ERR("File write failed\n");
                delete [] buff;
                return NVSIPL_STATUS_ERROR;
            }
        }

        delete [] buff;

        static bool bDumpInfo = true;
        if (bDumpInfo) {
            cout << "Width: " << surfaceMap.width << " Height: " << surfaceMap.height << endl;
            cout << "Bpp:";
            for (uint32_t i = 0u; i < numSurfaces; i++) {
                cout << " " << bytesPerPixelPtr[i];
            }
            cout << endl;
            bDumpInfo = false;
        }
        return NVSIPL_STATUS_OK;
    }

    void Deinit(void)
    {
        if (m_pOutFile != nullptr) {
            fflush(m_pOutFile);
            fclose(m_pOutFile);
        }
    }

 private:
    FILE *m_pOutFile = nullptr;

    typedef struct {
        float heightFactor[6];
        float widthFactor[6];
        uint32_t numSurfaces;
    } ImgUtilSurfParams;

    ImgUtilSurfParams ImgSurfParamsTable_RAW  = {
        .heightFactor = {1, 0, 0, 0, 0, 0},
        .widthFactor = {1, 0, 0, 0, 0, 0},
        .numSurfaces = 1,
    };

    ImgUtilSurfParams ImgSurfParamsTable_YUV[3][4] = {
        { /* PLANAR */
            { /* 420 */
                .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
                .widthFactor = {1, 0.5, 0.5, 0, 0, 0},
                .numSurfaces = 3,
            },
            { /* 422 */
                .heightFactor = {1, 1, 1, 0, 0, 0},
                .widthFactor = {1, 0.5, 0.5, 0, 0, 0},
                .numSurfaces = 3,
            },
            { /* 444 */
                .heightFactor = {1, 1, 1, 0, 0, 0},
                .widthFactor = {1, 1, 1, 0, 0, 0},
                .numSurfaces = 3,
            },
            { /* 422R */
                .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
                .widthFactor = {1, 1, 1, 0, 0, 0},
                .numSurfaces = 3,
            },
        },
        { /* SEMI_PLANAR */
            { /* 420 */
                .heightFactor = {1, 0.5, 0, 0, 0, 0},
                .widthFactor = {1, 0.5, 0, 0, 0, 0},
                .numSurfaces = 2,
            },
            { /* 422 */
                .heightFactor = {1, 1, 0, 0, 0, 0},
                .widthFactor = {1, 0.5, 0, 0, 0, 0},
                .numSurfaces = 2,
            },
            { /* 444 */
                .heightFactor = {1, 1, 0.5, 0, 0, 0},
                .widthFactor = {1, 1, 0.5, 0, 0, 0},
                .numSurfaces = 2,
            },
            { /* 422R */
                .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
                .widthFactor = {1, 1, 0.5, 0, 0, 0},
                .numSurfaces = 2,
            },
        },
        { /* PACKED */
            { /* 420 */
                .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
                .widthFactor = {1, 0.5, 0.5, 0, 0, 0},
                .numSurfaces = 3,
            },
            { /* 422 */
                .heightFactor = {1, 1, 1, 0, 0, 0},
                .widthFactor = {1, 0.5, 0.5, 0, 0, 0},
                .numSurfaces = 3,
            },
            { /* 444 */
                .heightFactor = {1, 1, 1, 0, 0, 0},
                .widthFactor = {1, 1, 1, 0, 0, 0},
                .numSurfaces = 3,
            },
            { /* 422R */
                .heightFactor = {1, 0.5, 0.5, 0, 0, 0},
                .widthFactor = {1, 1, 1, 0, 0, 0},
                .numSurfaces = 3,
            },
        },
    };

    ImgUtilSurfParams ImgSurfParamsTable_Packed  = {
        .heightFactor = {1, 0, 0, 0, 0, 0},
        .widthFactor = {1, 0, 0, 0, 0, 0},
        .numSurfaces = 1,
    };

    uint32_t ImgBytesPerPixelTable_RAW[10][6] = {
        {1, 0, 0, 0, 0, 0}, /* 8 */
        {2, 0, 0, 0, 0, 0}, /* 10 */
        {2, 0, 0, 0, 0, 0}, /* 12 */
        {2, 0, 0, 0, 0, 0}, /* 14 */
        {2, 0, 0, 0, 0, 0}, /* 16 */
        {4, 0, 0, 0, 0, 0}, /* 32 */
        {4, 0, 0, 0, 0, 0}, /* 16_8_8 */
        {4, 0, 0, 0, 0, 0}, /* 10_8_8 */
        {4, 0, 0, 0, 0, 0}, /* 2_10_10_10 */
        {4, 0, 0, 0, 0, 0}, /* 20 */
    };

    uint32_t ImgBytesPerPixelTable_YUV[2][9][6] = {
        { /* PLANAR */
            {1, 1, 1, 0, 0, 0}, /* 8 */
            {2, 2, 2, 0, 0, 0}, /* 10 */
            {2, 2, 2, 0, 0, 0}, /* 12 */
            {2, 2, 2, 0, 0, 0}, /* 14 */
            {2, 2, 2, 0, 0, 0}, /* 16 */
            {4, 4, 4, 0, 0, 0}, /* 32 */
            {2, 1, 1, 0, 0, 0}, /* 16_8_8 */
            {2, 1, 1, 0, 0, 0}, /* 10_8_8 */
            {4, 0, 0, 0, 0, 0}, /* 2_10_10_10 */
        },
        { /* SEMI_PLANAR */
            {1, 2, 0, 0, 0, 0}, /* 8 */
            {2, 4, 0, 0, 0, 0}, /* 10 */
            {2, 4, 0, 0, 0, 0}, /* 12 */
            {2, 4, 0, 0, 0, 0}, /* 14 */
            {2, 4, 0, 0, 0, 0}, /* 16 */
            {4, 8, 0, 0, 0, 0}, /* 32 */
            {2, 2, 0, 0, 0, 0}, /* 16_8_8 */
            {2, 2, 0, 0, 0, 0}, /* 10_8_8 */
            {4, 0, 0, 0, 0, 0}, /* 2_10_10_10 */
        },
    };

    SIPLStatus GetBytesPerCompForPackedYUV(uint32_t surfBPC, uint32_t *bytesPerComp)
    {
        switch(surfBPC) {
            case NVM_SURF_ATTR_BITS_PER_COMPONENT_8:
            case NVM_SURF_ATTR_BITS_PER_COMPONENT_LAYOUT_2_10_10_10:
                *bytesPerComp = 1u;
                break;
            case NVM_SURF_ATTR_BITS_PER_COMPONENT_10:
            case NVM_SURF_ATTR_BITS_PER_COMPONENT_12:
            case NVM_SURF_ATTR_BITS_PER_COMPONENT_14:
            case NVM_SURF_ATTR_BITS_PER_COMPONENT_16:
                *bytesPerComp = 2u;
                break;
            case NVM_SURF_ATTR_BITS_PER_COMPONENT_20:
                *bytesPerComp = 3u;
                break;
            case NVM_SURF_ATTR_BITS_PER_COMPONENT_32:
                *bytesPerComp = 4u;
                break;
            default:
                return NVSIPL_STATUS_ERROR;
        }

        return NVSIPL_STATUS_OK;
    }

    SIPLStatus GetSurfParams(uint32_t surfaceType,
                             float **xScale,
                             float **yScale,
                             uint32_t **bytesPerPixel,
                             uint32_t *numSurfacesVal)
    {
        NvMediaStatus nvmStatus;
        static uint32_t yuvPackedTbl[6] = {1, 0, 0, 0, 0, 0};
        uint32_t surfType, surfMemoryType, surfSubSamplingType, surfBPC, surfCompOrder;
        uint32_t numSurfaces = 1u, numComps = 1u;
        uint32_t *bytesPerPixelPtr = nullptr;
        float *xScalePtr = nullptr, *yScalePtr = nullptr;

        NVM_SURF_FMT_DEFINE_ATTR(srcAttr);
        nvmStatus = NvMediaSurfaceFormatGetAttrs(surfaceType,
                                                 srcAttr,
                                                 NVM_SURF_FMT_ATTR_MAX);
        if (nvmStatus != NVMEDIA_STATUS_OK) {
            LOG_ERR("NvMediaSurfaceFormatGetAttrs failed\n");
            return NVSIPL_STATUS_ERROR;
        }

        surfType = srcAttr[NVM_SURF_ATTR_SURF_TYPE].value;
        surfMemoryType = srcAttr[NVM_SURF_ATTR_MEMORY].value;
        surfSubSamplingType = srcAttr[NVM_SURF_ATTR_SUB_SAMPLING_TYPE].value;
        surfBPC = srcAttr[NVM_SURF_ATTR_BITS_PER_COMPONENT].value;
        surfCompOrder = srcAttr[NVM_SURF_ATTR_COMPONENT_ORDER].value;

        switch (surfType) {
            case NVM_SURF_ATTR_SURF_TYPE_YUV:
                if ((surfSubSamplingType == NVM_SURF_ATTR_SUB_SAMPLING_TYPE_NONE)
                    && (surfMemoryType == NVM_SURF_ATTR_MEMORY_PACKED)) {

                    xScalePtr =  &ImgSurfParamsTable_Packed.widthFactor[0];
                    yScalePtr = &ImgSurfParamsTable_Packed.heightFactor[0];
                    numSurfaces = ImgSurfParamsTable_Packed.numSurfaces;

                    SIPLStatus status = GetBytesPerCompForPackedYUV(surfBPC, &yuvPackedTbl[0]);
                    if (status != NVSIPL_STATUS_OK) {
                        LOG_ERR("Invalid bits per component and packed YUV combination\n");
                        return status;
                    }

                    switch (surfCompOrder) {
                        case NVM_SURF_ATTR_COMPONENT_ORDER_VUYX:
                        case NVM_SURF_ATTR_COMPONENT_ORDER_XYUV:
                        case NVM_SURF_ATTR_COMPONENT_ORDER_XUYV:
                            numComps = 4u;
                            break;
                        case NVM_SURF_ATTR_COMPONENT_ORDER_UYVY:
                        case NVM_SURF_ATTR_COMPONENT_ORDER_VYUY:
                        case NVM_SURF_ATTR_COMPONENT_ORDER_YVYU:
                        case NVM_SURF_ATTR_COMPONENT_ORDER_YUYV:
                            numComps = 2u;
                            break;
                        case NVM_SURF_ATTR_COMPONENT_ORDER_LUMA:
                            numComps = 1u;
                            break;
                        default:
                            LOG_ERR("Invalid component order and packed YUV combination\n");
                            return NVSIPL_STATUS_ERROR;
                    }
                    yuvPackedTbl[0] = yuvPackedTbl[0] * numComps;
                    bytesPerPixelPtr = &yuvPackedTbl[0];
                } else {
                    xScalePtr = &ImgSurfParamsTable_YUV[0][surfSubSamplingType - NVM_SURF_ATTR_SUB_SAMPLING_TYPE_420].widthFactor[0];
                    yScalePtr = &ImgSurfParamsTable_YUV[0][surfSubSamplingType - NVM_SURF_ATTR_SUB_SAMPLING_TYPE_420].heightFactor[0];
                    numSurfaces = ImgSurfParamsTable_YUV[0][surfSubSamplingType - NVM_SURF_ATTR_SUB_SAMPLING_TYPE_420].numSurfaces;
                    bytesPerPixelPtr = &ImgBytesPerPixelTable_YUV[0][surfBPC - NVM_SURF_ATTR_BITS_PER_COMPONENT_8][0];
                }
                break;
            case NVM_SURF_ATTR_SURF_TYPE_RAW:
                bytesPerPixelPtr = &ImgBytesPerPixelTable_RAW[surfBPC - NVM_SURF_ATTR_BITS_PER_COMPONENT_8][0];
                xScalePtr = &ImgSurfParamsTable_RAW.widthFactor[0];
                yScalePtr = &ImgSurfParamsTable_RAW.heightFactor[0];
                numSurfaces =  ImgSurfParamsTable_RAW.numSurfaces;
                break;
            default:
                LOG_ERR("Unsupported pixel format:%u\n", surfType);
                return NVSIPL_STATUS_ERROR;
        }

        if (xScale) {
            *xScale = xScalePtr;
        }
        if (yScale) {
            *yScale = yScalePtr;
        }
        if (bytesPerPixel) {
            *bytesPerPixel = bytesPerPixelPtr;
        }
        if (numSurfacesVal) {
            *numSurfacesVal = numSurfaces;
        }

        return NVSIPL_STATUS_OK;
    }
};

#endif //CFileWriter_HPP
