// 
// Notice Regarding Standards.  AMD does not provide a license or sublicense to
// any Intellectual Property Rights relating to any standards, including but not
// limited to any audio and/or video codec technologies such as MPEG-2, MPEG-4;
// AVC/H.264; HEVC/H.265; AAC decode/FFMPEG; AAC encode/FFMPEG; VC-1; and MP3
// (collectively, the "Media Technologies"). For clarity, you will pay any
// royalties due for such third party technologies, which may include the Media
// Technologies that are owed as a result of AMD providing the Software to you.
// 
// MIT license 
// 
//
// Copyright (c) 2018 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//

// this sample demonstrates managing the DPB externally with the reference frame injection
// and the reconstructed frame dump features. The regular output will be written to an
// H.264 / H.265 elementary stream and the reconstructed frames will be dumped as NV12 frames
// based on the static variable options

#include <stdio.h>
#include <fstream>
#include <codecvt>
#include <locale>

#ifdef _WIN32
#include <tchar.h>
#include <d3d9.h>
#include <d3d11.h>
#endif

#include "public/common/AMFFactory.h"
#include "public/common/Thread.h"
#include "public/common/AMFSTL.h"
#include "public/common/TraceAdapter.h"

#include "public/include/components/VideoEncoderVCE.h"
#include "public/include/components/VideoEncoderHEVC.h"

#include "public/samples/CPPSamples/common/SurfaceGenerator.h"
#include "public/samples/CPPSamples/common/SurfaceUtils.h"
#include "public/samples/CPPSamples/common/MiscHelpers.h"
#include "public/samples/CPPSamples/common/PollingThread.h"

#define AMF_FACILITY L"SimpleFrameInjection"

static const wchar_t* pCodec = AMFVideoEncoderVCE_AVC;
//static const wchar_t *pCodec = AMFVideoEncoder_HEVC;
static const bool isH264 = (amf_wstring(pCodec) == amf_wstring(AMFVideoEncoderVCE_AVC)) ? true : false;

//static const amf_int64 picMgmtMode = (isH264) ? AMF_VIDEO_ENCODER_PICTURE_TRANSFER_MODE_OFF : AMF_VIDEO_ENCODER_HEVC_PICTURE_TRANSFER_MODE_OFF; // disable ref/recon pic transfer
static const amf_int64 picMgmtMode = (isH264) ? amf_int64(AMF_VIDEO_ENCODER_PICTURE_TRANSFER_MODE_ON) : amf_int64(AMF_VIDEO_ENCODER_HEVC_PICTURE_TRANSFER_MODE_ON);     // enable ref/recon pic transfer

static const bool injectReconDump = true; // inject the reconstructed picture dump from the encoder
//static const bool injectReconDump = false; // inject the frame generated by this sample

static const wchar_t* fileNameOut = isH264 ? L"./output.h264" : L"./output.h265";
static const wchar_t* fileNameOutRecon = L"./recon_dump.nv12";    // reconstructed picture dump from encoder

//#define ENABLE_EFC // color conversion inside encoder component. Will use GFX or HW if available

static bool bEnable4K = false;
static bool bMaximumSpeed = true;

#ifdef _WIN32
//static amf::AMF_MEMORY_TYPE memoryTypeIn  = amf::AMF_MEMORY_DX9;
static amf::AMF_MEMORY_TYPE memoryTypeIn = amf::AMF_MEMORY_DX11;
#else
static amf::AMF_MEMORY_TYPE memoryTypeIn  = amf::AMF_MEMORY_VULKAN;
#endif
#if defined ENABLE_EFC
static amf::AMF_SURFACE_FORMAT formatIn = amf::AMF_SURFACE_RGBA;
//static amf::AMF_SURFACE_FORMAT formatIn = amf::AMF_SURFACE_BGRA;
//static amf::AMF_SURFACE_FORMAT formatIn = amf::AMF_SURFACE_R10G10B10A2;
//static amf::AMF_SURFACE_FORMAT formatIn = amf::AMF_SURFACE_RGBA_F16;
#else
static amf::AMF_SURFACE_FORMAT formatIn   = amf::AMF_SURFACE_NV12;
#endif
static amf_int32 widthIn                  = (bEnable4K ? 1920 * 2 : 1920);
static amf_int32 heightIn                 = (bEnable4K ? 1080 * 2 : 1080);
static amf_int32 frameRateIn              = 30;
static amf_int64 bitRateIn                = 5000000L; // in bits, 5MBit
amf_int32 rectSize                        = 50;
static amf_int32 frameCount               = 200;


#ifndef _WIN32
// Linux doesn't define _wcsicmp but the same
// functionality can be found in wcscasecmp
#define _wcsicmp wcscasecmp
#endif

// color for main frame
amf::AMFSurfacePtr pColor1;
amf::AMFSurfacePtr pColor2;

// color for the frame injection
amf::AMFSurfacePtr pColor3;
amf::AMFSurfacePtr pColor4;

#define NO_ARGS 1

amf::AMFCriticalSection  cs;
amf::AMFSurfacePtr       spReconDumpSurfaceRef;

class EncPollingThread : public PollingThread
{
public:
    EncPollingThread(amf::AMFContext* pContext, amf::AMFComponent* pEncoder, const wchar_t* pFileName, const wchar_t* pFileNameOutRecon);
    ~EncPollingThread();
protected:
    void ProcessData(amf::AMFData* pData) override;
    void PrintResults() override;

    amf::AMFDataStreamPtr   m_pReconFile;
};

#ifdef _WIN32
//int _tmain(int argc, _TCHAR* argv[])
amf_int main(amf_int /* argc */, char** /* argv */)
#else
int main(int /* argc */, char* /* argv */ [])
#endif
{
    AMF_RESULT res = AMF_OK; // error checking can be added later
    res = g_AMFFactory.Init();
    if(res != AMF_OK)
    {
        wprintf(L"AMF Failed to initialize");
        return 1;
    }

    ::amf_increase_timer_precision();

    amf::AMFTraceEnableWriter(AMF_TRACE_WRITER_CONSOLE, true);
    amf::AMFTraceEnableWriter(AMF_TRACE_WRITER_DEBUG_OUTPUT, true);
    
    // initialize AMF
    amf::AMFContextPtr context;
    amf::AMFComponentPtr encoder;
    amf::AMFSurfacePtr surfaceIn;

    // context
    res = g_AMFFactory.GetFactory()->CreateContext(&context);
    AMF_RETURN_IF_FAILED(res, L"CreateContext() failed");

    if(memoryTypeIn == amf::AMF_MEMORY_VULKAN)
    {
        res = amf::AMFContext1Ptr(context)->InitVulkan(NULL);
        AMF_RETURN_IF_FAILED(res, L"InitVulkan(NULL) failed");
        PrepareFillFromHost(context, memoryTypeIn, formatIn, widthIn, heightIn, true);
    }
#ifdef _WIN32
    else if(memoryTypeIn == amf::AMF_MEMORY_DX9)
    {

        res = context->InitDX9(NULL); // can be DX9 or DX9Ex device
        AMF_RETURN_IF_FAILED(res, L"InitDX9(NULL) failed");
    }
    else if(memoryTypeIn == amf::AMF_MEMORY_DX11)
    {
        res = context->InitDX11(NULL); // can be DX11 device
        AMF_RETURN_IF_FAILED(res, L"InitDX11(NULL) failed");
        PrepareFillFromHost(context, memoryTypeIn, formatIn, widthIn, heightIn, true);
    }
#endif
    
    // component: encoder
    res = g_AMFFactory.GetFactory()->CreateComponent(context, pCodec, &encoder);
    AMF_RETURN_IF_FAILED(res, L"CreateComponent(%s) failed", pCodec);

    // set encoder params
    if(isH264)
    { 
        res = SetEncoderDefaultsAVC(encoder, bitRateIn, frameRateIn, bMaximumSpeed, bEnable4K, false);
        AMF_RETURN_IF_FAILED(res, L"SetEncoderDefaultsAVC() failed");

        res = encoder->SetProperty(AMF_VIDEO_ENCODER_PICTURE_TRANSFER_MODE, picMgmtMode);
        AMF_RETURN_IF_FAILED(res, L"SetProperty(AMF_VIDEO_ENCODER_PICTURE_TRANSFER_MODE failed");
    }
    else
    {
        res = SetEncoderDefaultsHEVC(encoder, bitRateIn, frameRateIn, false, bMaximumSpeed, bEnable4K, false);
        AMF_RETURN_IF_FAILED(res, L"SetEncoderDefaultsHEVC() failed");

        res = encoder->SetProperty(AMF_VIDEO_ENCODER_HEVC_PICTURE_TRANSFER_MODE, picMgmtMode);
        AMF_RETURN_IF_FAILED(res, L"SetProperty(AMF_VIDEO_ENCODER_HEVC_PICTURE_TRANSFER_MODE failed");
    }

    res = encoder->Init(formatIn, widthIn, heightIn);
    AMF_RETURN_IF_FAILED(res, L"encoder->Init() failed");


    // allocate reference frame for frame injection
    amf::AMFSurfacePtr surfaceRef;
    res = context->AllocSurface(memoryTypeIn, formatIn, widthIn, heightIn, &surfaceRef);
    AMF_RETURN_IF_FAILED(res, L"AllocSurface() failed");

    if (memoryTypeIn == amf::AMF_MEMORY_VULKAN)
    {
        FillSurfaceVulkan(context, surfaceRef, true);
    }
#ifdef _WIN32
    else if (memoryTypeIn == amf::AMF_MEMORY_DX9)
    {
        FillSurfaceDX9(context, surfaceRef, true);
    }
    else
    {
        FillSurfaceDX11(context, surfaceRef, true);
    }
#endif

    EncPollingThread thread(context, encoder, fileNameOut, fileNameOutRecon);
    thread.Start();

    // encode some frames
    amf_int32 submitted = 0;
    while(submitted < frameCount)
    {
        if(surfaceIn == NULL)
        {
            surfaceIn = NULL;
            res = context->AllocSurface(memoryTypeIn, formatIn, widthIn, heightIn, &surfaceIn);
            AMF_RETURN_IF_FAILED(res, L"AllocSurface() failed");
            
            if(memoryTypeIn == amf::AMF_MEMORY_VULKAN)
            {
                FillSurfaceVulkan(context, surfaceIn, false);
            }
#ifdef _WIN32
            else if(memoryTypeIn  == amf::AMF_MEMORY_DX9)
            { 
                FillSurfaceDX9(context, surfaceIn, false);
            }
            else
            {
                FillSurfaceDX11(context, surfaceIn, false);
            }
#endif
        }
        // encode
        amf_pts start_time = amf_high_precision_clock();
        surfaceIn->SetProperty(START_TIME_PROPERTY, start_time);

        // inject reference frame,
        // the first frame in the stream is always an intra frame,
        // even if you pass in a reference frame, it will not be used
        if ((AMF_VIDEO_ENCODER_PICTURE_TRANSFER_MODE_ON == picMgmtMode ||
            AMF_VIDEO_ENCODER_HEVC_PICTURE_TRANSFER_MODE_ON == picMgmtMode) && submitted != 0)
        {
            const wchar_t* propId = (isH264) ? AMF_VIDEO_ENCODER_REFERENCE_PICTURE : AMF_VIDEO_ENCODER_HEVC_REFERENCE_PICTURE;
            if (injectReconDump)
            {
                while (spReconDumpSurfaceRef == NULL)
                {
                    amf_sleep(1); // wait for recon dump
                }
                amf::AMFLock  lock(&cs);
                res = surfaceIn->SetProperty(propId, spReconDumpSurfaceRef);
                spReconDumpSurfaceRef = NULL;
            }
            else
            {
                amf::AMFLock  lock(&cs);
                res = surfaceIn->SetProperty(propId, surfaceRef);
            }
            AMF_ASSERT_OK(res, L"Failed to set property AMF_VIDEO_ENCODER_REFERENCE_PICTURE / AMF_VIDEO_ENCODER_HEVC_REFERENCE_PICTURE");
        }

        res = encoder->SubmitInput(surfaceIn);
        if(res == AMF_NEED_MORE_INPUT) // handle full queue
        {
            // do nothing
        }
		else if (res == AMF_INPUT_FULL || res == AMF_DECODER_NO_FREE_SURFACES)
		{
			amf_sleep(1); // input queue is full: wait, poll and submit again
		}
        else
        {
            AMF_RETURN_IF_FAILED(res, L"SubmitInput() failed");

            surfaceIn = NULL;
            submitted++;
        }
    }
    // drain encoder; input queue can be full
    while(true)
    {
        res = encoder->Drain();
        if(res != AMF_INPUT_FULL) // handle full queue
        {
            break;
        }
        amf_sleep(1); // input queue is full: wait and try again
    }

    // Need to request stop before waiting for stop
    if (thread.RequestStop() == false)
    {
        AMFTraceError(AMF_FACILITY, L"thread.RequestStop() Failed");
    }

    if (thread.WaitForStop() == false)
    {
        AMFTraceError(AMF_FACILITY, L"thread.WaitForStop() Failed");
    }

    pColor1 = NULL;
    pColor2 = NULL;
    pColor3 = NULL;
    pColor4 = NULL;

    // cleanup in this order
    surfaceIn = NULL;
    surfaceRef = NULL;
    spReconDumpSurfaceRef = NULL;
    encoder->Terminate();
    encoder = NULL;
    context->Terminate();
    context = NULL; // context is the last

    g_AMFFactory.Terminate();
    return 0;
}


EncPollingThread::EncPollingThread(amf::AMFContext* pContext, amf::AMFComponent* pEncoder, const wchar_t* pFileName, const wchar_t* pFileNameOutRecon)
    : PollingThread(pContext, pEncoder, pFileName, true)
{
    AMF_RESULT res = AMF_OK;
    res = amf::AMFDataStream::OpenDataStream(pFileNameOutRecon, amf::AMFSO_WRITE, amf::AMFFS_SHARE_READ, &m_pReconFile);
    AMF_ASSERT_OK(res, L"Failed to open file %s", pFileNameOutRecon);
}

EncPollingThread::~EncPollingThread()
{
    if (m_pReconFile != NULL)
    {
        m_pReconFile->Close();
        m_pReconFile = NULL;
    }
}

void EncPollingThread::ProcessData(amf::AMFData* pData)
{
    AMF_RESULT res = AMF_OK;
    
    AdjustTimes(pData);

    res = pData->Convert(amf::AMF_MEMORY_HOST); // convert to system memory
    AMF_ASSERT_OK(res, L"Failed to convert GPU memory to system memory");

    amf::AMFBufferPtr spBuffer(pData); // query for buffer interface
    if (m_pFile != NULL)
    {
        res = m_pFile->Write(spBuffer->GetNative(), spBuffer->GetSize(), NULL);
        AMF_ASSERT_OK(res, L"Failed to write to output file");
    }

    // when output is available, get pointer to the reconstructed picture
    amf::AMFInterfacePtr spValue;
    const wchar_t* propId = (isH264) ? AMF_VIDEO_ENCODER_RECONSTRUCTED_PICTURE : AMF_VIDEO_ENCODER_HEVC_RECONSTRUCTED_PICTURE;
    res = pData->GetProperty(propId, &spValue);

    if (res != AMF_NOT_FOUND)
    {
        AMF_ASSERT_OK(res, L"Failed to get property AMF_VIDEO_ENCODER_RECONSTRUCTED_PICTURE / AMF_VIDEO_ENCODER_HEVC_RECONSTRUCTED_PICTURE");

        amf::AMFSurfacePtr spSurface(spValue);
        amf::AMF_MEMORY_TYPE origMemory = spSurface->GetMemoryType();
        res = spSurface->Convert(amf::AMF_MEMORY_HOST);
        AMF_ASSERT_OK(res, L"Failed to convert GPU memory to system memory");

        if (spSurface != NULL)
        {
            res = WritePlane(spSurface->GetPlane(amf::AMF_PLANE_Y), m_pReconFile); // get y-plane pixels
            AMF_ASSERT_OK(res, L"Failed to write to recon file");
            res = WritePlane(spSurface->GetPlane(amf::AMF_PLANE_UV), m_pReconFile); // get uv-plane pixels
            AMF_ASSERT_OK(res, L"Failed to write to recon file");

            // convert the surface back to it's original memory 
            // format in case we need to pass it back to the encoder
            res = spSurface->Convert(origMemory);
            AMF_ASSERT_OK(res, L"Failed to convert system memory back to GPU memory");

            amf::AMFLock  lock(&cs);
            spReconDumpSurfaceRef = spSurface;
        }
    }

    m_WriteDuration += amf_high_precision_clock() - m_LastPollTime;
}

void EncPollingThread::PrintResults()
{
    PrintTimes("encode ", frameCount);
}
