/*
 * 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 <thread>
#include <queue>
#include <mutex>
#include <atomic>
#include <condition_variable>

#include "nvmedia_2d.h" // NvMedia2D Header for Blit
#include "nvmedia_idp.h" // NvMediaIDP Header  for Display

#include "NvSIPLClient.hpp" // NvSIPL Client header for input buffers

#ifndef CCOMPOSITE_HPP
#define CCOMPOSITE_HPP

using namespace nvsipl;

class CComposite
{
  protected:
    typedef enum {
        QueueType_Mailbox,
        QueueType_Fifo
    } QueueType;

  public:
    // Destructor
    virtual ~CComposite()
    {
        if (m_bRunning) {
            this->Stop();
        }
    }

    // Initializes compositor
    virtual SIPLStatus Init(uint32_t uDispID,
                            uint32_t uWinID,
                            uint32_t uWinDepth,
                            NvMediaRect *pRect,
                            NvSciBufModule bufModule,
                            NvSciSyncModule syncModule);

    // Registers a source and returns id assigned to it.
    virtual SIPLStatus RegisterSource(uint32_t groupIndex,
                                      uint32_t modIndex,
                                      uint32_t outIndex,
                                      bool isRaw,
                                      uint32_t &id,
                                      NvSciStreamBlock *consumer,
                                      NvSciStreamBlock **upstream,
                                      NvSciStreamBlock *queue,
                                      QueueType queueType=QueueType_Fifo);

    SIPLStatus Post(uint32_t id, INvSIPLClient::INvSIPLNvMBuffer *pBuffer);

    SIPLStatus SetActiveGroup(uint32_t groupIndex);

    void PrintDisplayableGroups() const;

    SIPLStatus Start();

    SIPLStatus Stop();

    virtual SIPLStatus Deinit();

  protected:
    virtual void ThreadFunc(void);

    SIPLStatus ClearOutputBuffers();

    virtual bool CheckInputQueues(void);

    SIPLStatus InitDisplay(uint32_t uDispID, uint32_t uWinID, uint32_t uWinDepth);

  protected:
    struct DestroyNvMediaDevice
    {
        void operator ()(NvMediaDevice *p) const
        {
            NvMediaDeviceDestroy(p);
        }
    };

    struct DestroyNvMedia2DDevice
    {
        void operator ()(NvMedia2D *p) const
        {
            NvMedia2DDestroy(p);
        }
    };

    struct DestroyNvMediaImage
    {
        void operator ()(NvMediaImage *p) const
        {
            NvMediaImageDestroy(p);
        }
    };

    struct DestroyNvMediaIDPDevice
    {
        void operator ()(NvMediaIDP *p) const
        {
            NvMediaIDPDestroy(p);
        }
    };

    std::unique_ptr<NvMediaDevice, DestroyNvMediaDevice> m_pDevice;
    std::unique_ptr<NvMediaIDP, DestroyNvMediaIDPDevice> m_pNvmIDP;

    // Display stuff
    bool m_bRectSet = false;
    NvMediaRect m_outputRect;
    NvMediaIDPDeviceParams m_oDispInfo;

    // Input stuff
    static const uint32_t NUM_OF_GROUPS = 4;
    static const uint32_t NUM_OF_COLS = 3;
    static const uint32_t NUM_OF_ROWS = 4;
    static const uint32_t NUM_OF_INPUTS = NUM_OF_GROUPS *NUM_OF_COLS*NUM_OF_ROWS;

    // Create 2D context for each row (camera module) to get around the limit on number of surfaces
    // that can be used with single context of 2D.
    std::unique_ptr<NvMedia2D, DestroyNvMedia2DDevice> m_p2DDevices[NUM_OF_ROWS] {nullptr};

    NvMediaRect m_oInputRects[NUM_OF_INPUTS];
    std::unique_ptr<NvMediaImage, DestroyNvMediaImage> m_scratchRGBAImages[NUM_OF_INPUTS] {nullptr};
    std::unique_ptr<uint8_t[]> m_rawImageBufs[NUM_OF_INPUTS] {nullptr};
    std::unique_ptr<uint8_t[]> m_rgbaImageBufs[NUM_OF_INPUTS] {nullptr};

    // Output stuff
    uint32_t m_outputWidth = 0u;
    uint32_t m_outputHeight = 0u;
    static const uint32_t NUM_OUT_BUFFERS = 2;
    std::unique_ptr<NvMediaImage, DestroyNvMediaImage> m_outputImages[NUM_OUT_BUFFERS] {nullptr};
    std::queue<NvMediaImage*> m_outputQueue;
    std::unique_ptr<NvMediaImage, DestroyNvMediaImage> m_pBlankFrame {nullptr};

    // Group (of sources that need to be composited)
    std::vector<uint32_t> m_GroupInfos[NUM_OF_GROUPS]; // Array of vectors containing IDs of sources in each group)
    uint32_t m_uGroupIndex = -1U; // Currently active Group
    std::atomic<uint32_t> m_uNewGroupIndex; // ID of new active Group

    // Thread stuff
    std::unique_ptr<std::thread> m_pthread {nullptr};
    std::atomic<bool> m_bRunning; // Flag indicating if composite is running.

  private:
    std::mutex m_vInputQueueMutex[NUM_OF_INPUTS];
    std::queue<INvSIPLClient::INvSIPLNvMBuffer*> m_vInputQueue[NUM_OF_INPUTS];
};

#endif // CCOMPOSITE
