/*
 * 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 <mutex>
#include <cmath>

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

#include "CComposite.hpp"
#include "CNvSIPLMasterNvSci.hpp"

#ifdef NVMEDIA_QNX
#include "nvtime.h"
#endif /* NVMEDIA_QNX */

#ifndef CNVSIPLCONSUMER_HPP
#define CNVSIPLCONSUMER_HPP

using namespace std;
using namespace nvsipl;

/** NvSIPL consumer class.
 * NvSIPL consumer consumes the output buffers of NvSIPL.
 * The output buffer can be either an NvMediaImage or CUeglFrame.
 * NvSIPL consumer must be derived from INvSIPLClient::INvMCallback
 * and implement callbacks to handle/process the output buffers.
 */
class CNvSIPLConsumer : public INvSIPLClient::INvMCallback
{
 public:
    typedef struct {
        std::mutex profDataMut;
        uint64_t uFrameCount;
        uint64_t uPrevFrameCount;
        uint64_t uFirstLatencyUs;
        uint64_t uTotLatencyUs;
        uint64_t uTotSqJitterUs;
    } ProfilingData;

    SIPLStatus Init(CComposite *pComposite=nullptr,
                    CNvSIPLMasterNvSci *pMasterNvSci=nullptr,
                    uint32_t uID=-1,
                    uint32_t uSensor=-1,
                    INvSIPLClient::ConsumerDesc::OutputType outputType=INvSIPLClient::ConsumerDesc::OutputType::ICP,
                    std::string sFilenamePrefix="",
                    uint32_t uNumSkipFrames=0u,
                    uint64_t uNumWriteFrames=-1u)
    {
        if ((pComposite != nullptr) && (pMasterNvSci != nullptr)) {
            LOG_ERR("CNvSIPLConsumer expects only one of pComposite and pMasterNvSci\n");
            return NVSIPL_STATUS_BAD_ARGUMENT;
        }

        m_pComposite = pComposite;
        m_pMasterNvSci = pMasterNvSci;
        m_uID = uID;
        m_uSensor = uSensor;
        m_outputType = outputType;
        m_sFilenamePrefix = sFilenamePrefix;
        m_uNumSkipFrames = uNumSkipFrames;
        m_uNumWriteFrames = uNumWriteFrames;

        m_profData.profDataMut.lock();
        m_profData.uFrameCount = 0u;
        m_profData.uPrevFrameCount = 0u;
        m_profData.uFirstLatencyUs = 0u;
        m_profData.uTotLatencyUs = 0u;
        m_profData.uTotSqJitterUs = 0u;
        m_profData.profDataMut.unlock();

        // Create file writer if necessary
        if (m_sFilenamePrefix != "") {
            bool bRawOut = m_outputType == INvSIPLClient::ConsumerDesc::OutputType::ICP;
            string sFileExt = bRawOut ? ".raw" : ".yuv";
            string sFilename = m_sFilenamePrefix + "_cam_" + std::to_string(m_uSensor)
                                   + "_out_" + std::to_string((uint32_t)m_outputType) + sFileExt;
            m_pFileWriter.reset(new CFileWriter);
            auto status = m_pFileWriter->Init(sFilename);
            if (status != NVSIPL_STATUS_OK) {
                LOG_ERR("Failed to initialize file writer\n");
                return status;
            }
        }

        return NVSIPL_STATUS_OK;
    }

    /** Deinitializes NvSIPL NvMedia consumer. */
    void Deinit(void)
    {
        if (m_pFileWriter) {
            m_pFileWriter->Deinit();
        }

        return;
    }

    /** Buffer callback.
     * Implements the callback to process the buffer from NvSIPL.
     * @param[in] NvSIPL Buffer. Points to an instance of INvSIPLNvMBuffer
     * depending on the configuration of the consumer.
     * @return NvSIPLStatus
     *  NVSIPL_STATUS_OK => Operation completed successfully.
     */
    SIPLStatus OnFrameAvailable(INvSIPLClient::INvSIPLBuffer* pBuffer) override
    {
        static uint64_t uPrevLatencyUs = 0u;
        auto pNvmBuffer = (INvSIPLClient::INvSIPLNvMBuffer*) pBuffer;
        uint64_t uFrameCounter = 0u;

        const auto& md = pNvmBuffer->GetImageData();
        NvMediaImage *image = pNvmBuffer->GetImage();
        uint64_t frameTimeUs = image->captureGlobalTimeStamp;
        struct timespec receivedTimespec;
#ifdef NVMEDIA_QNX
        nvtime_gettime_mono(&receivedTimespec);
#else /* NVMEDIA_QNX */
        clock_gettime(CLOCK_MONOTONIC, &receivedTimespec);
#endif /* NVMEDIA_QNX */
        uint64_t receivedTimeUs = ((uint64_t)receivedTimespec.tv_sec * 1000000)
                                + ((uint64_t)receivedTimespec.tv_nsec / 1000);
        uint64_t uLatencyUs = receivedTimeUs - frameTimeUs;

        m_profData.profDataMut.lock();
        uFrameCounter = m_profData.uFrameCount;
        if (m_profData.uFirstLatencyUs == 0u) {
            m_profData.uFirstLatencyUs = uLatencyUs;
        }
        m_profData.uTotLatencyUs += uLatencyUs;
        if (uPrevLatencyUs != 0u) {
            // Use root mean square
            m_profData.uTotSqJitterUs += pow((int64_t)(uLatencyUs - uPrevLatencyUs), 2);
        }
        m_profData.uFrameCount++;
        m_profData.profDataMut.unlock();
        uPrevLatencyUs = uLatencyUs;

        // Send to file writer
        if (m_pFileWriter && (uFrameCounter >= m_uNumSkipFrames)) {
            if ((m_uNumWriteFrames == -1u) || (uFrameCounter < (m_uNumSkipFrames + m_uNumWriteFrames))) {
                auto status = m_pFileWriter->WriteBufferToFile(pNvmBuffer);
                if (status != NVSIPL_STATUS_OK) {
                    LOG_ERR("WriteBufferToFile failed\n");
                    return status;
                }
            }
        }

        if (m_bShowMetadata) {
            cout << "Camera ID: " << m_uSensor << endl;
            cout << " Frame Counter: " << (md.frameSeqNumInfo.frameSeqNumValid ?
                                           md.frameSeqNumInfo.frameSequenceNumber : uFrameCounter)
                                       << endl;
            cout << " TSC: " << md.frameCaptureTSC << endl;
            if (md.badPixelStatsValid) {
                cout << " Bad pixel stats:" << endl;
                cout << "     highInWin: " << md.badPixelStats.highInWin << endl;
                cout << "     lowInWin: " << md.badPixelStats.lowInWin << endl;
                cout << "     highMagInWin: " << md.badPixelStats.highMagInWin << endl;
                cout << "     lowMagInWin: " << md.badPixelStats.lowMagInWin << endl;
                cout << "     highOutWin: " << md.badPixelStats.highOutWin << endl;
                cout << "     lowOutWin: " << md.badPixelStats.lowOutWin << endl;
                cout << "     highMagOutWin: " << md.badPixelStats.highMagOutWin << endl;
                cout << "     lowMagOutWin: " << md.badPixelStats.lowMagOutWin << endl;
            }
        }

        if (m_pComposite != nullptr) {
            auto status = m_pComposite->Post(m_uID, pNvmBuffer);
            if (status != NVSIPL_STATUS_OK) {
                LOG_ERR("Compositor post failed\n");
                return status;
            }
        }
        if (m_pMasterNvSci != nullptr) {
            auto status = m_pMasterNvSci->Post(m_uSensor, m_outputType, pNvmBuffer);
            if (status != NVSIPL_STATUS_OK) {
                LOG_ERR("Master post failed\n");
                return status;
            }
        }

        return NVSIPL_STATUS_OK;
    }

    void EnableMetadataLogging(void) {
        m_bShowMetadata = true;
    }

    ProfilingData m_profData;
    uint32_t m_uSensor = -1;
    INvSIPLClient::ConsumerDesc::OutputType m_outputType;

    virtual ~CNvSIPLConsumer() = default;

 private:
    uint32_t m_uID = -1;
    unique_ptr<CFileWriter> m_pFileWriter = nullptr;
    CComposite *m_pComposite = nullptr;
    CNvSIPLMasterNvSci *m_pMasterNvSci = nullptr;
    string m_sFilenamePrefix = "";
    uint32_t m_uNumSkipFrames = 0u;
    uint64_t m_uNumWriteFrames = -1u;
    bool m_bShowMetadata = false;
};

#endif //CNVSIPLCONSUMER_HPP
