/**
 * FreeRDP: A Remote Desktop Protocol Implementation
 * H.264 Bitmap Compression
 *
 * Copyright 2014 Mike McDonald <Mike.McDonald@software.dell.com>
 * Copyright 2017 David Fort <contact@hardening-consulting.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <winpr/crt.h>
#include <winpr/print.h>
#include <winpr/library.h>
#include <winpr/bitstream.h>
#include <winpr/synch.h>

#include <freerdp/primitives.h>
#include <freerdp/codec/h264.h>
#include <freerdp/log.h>

#include "h264.h"

#define TAG FREERDP_TAG("codec")

static BOOL avc444_ensure_buffer(H264_CONTEXT* h264, DWORD nDstHeight);

BOOL avc420_ensure_buffer(H264_CONTEXT* h264, UINT32 stride, UINT32 width, UINT32 height)
{
	if (!h264)
		return FALSE;

	if (stride == 0)
		stride = width;

	if (stride % 16 != 0)
		stride += 16 - stride % 16;

	if (height % 16 != 0)
		height += 16 - height % 16;

	if (!h264->pYUVData[0] || !h264->pYUVData[1] || !h264->pYUVData[2] ||
	    (width != h264->width) || (height != h264->height) || (stride != h264->iStride[0]))
	{
		h264->iStride[0] = stride;
		h264->iStride[1] = (stride + 1) / 2;
		h264->iStride[2] = (stride + 1) / 2;
		h264->width = width;
		h264->height = height;
		_aligned_free(h264->pYUVData[0]);
		_aligned_free(h264->pYUVData[1]);
		_aligned_free(h264->pYUVData[2]);
		h264->pYUVData[0] = _aligned_malloc(h264->iStride[0] * height, 16);
		h264->pYUVData[1] = _aligned_malloc(h264->iStride[1] * height, 16);
		h264->pYUVData[2] = _aligned_malloc(h264->iStride[2] * height, 16);

		if (!h264->pYUVData[0] || !h264->pYUVData[1] || !h264->pYUVData[2])
			return FALSE;
	}

	return TRUE;
}


static BOOL check_rect(const H264_CONTEXT* h264, const RECTANGLE_16* rect,
                       UINT32 nDstWidth, UINT32 nDstHeight)
{
	/* Check, if the output rectangle is valid in decoded h264 frame. */
	if ((rect->right > h264->width) || (rect->left > h264->width))
		return FALSE;

	if ((rect->top > h264->height) || (rect->bottom > h264->height))
		return FALSE;

	/* Check, if the output rectangle is valid in destination buffer. */
	if ((rect->right > nDstWidth) || (rect->left > nDstWidth))
		return FALSE;

	if ((rect->bottom > nDstHeight) || (rect->top > nDstHeight))
		return FALSE;

	return TRUE;
}

static BOOL avc_yuv_to_rgb(H264_CONTEXT* h264, const RECTANGLE_16* regionRects,
                           UINT32 numRegionRects, UINT32 nDstWidth,
                           UINT32 nDstHeight, UINT32 nDstStep, BYTE* pDstData,
                           DWORD DstFormat, BOOL use444)
{
	UINT32 x;
	BYTE* pDstPoint;
	prim_size_t roi;
	INT32 width, height;
	const BYTE* pYUVPoint[3];
	primitives_t* prims = primitives_get();

	for (x = 0; x < numRegionRects; x++)
	{
		const RECTANGLE_16* rect = &(regionRects[x]);
		const UINT32* iStride;
		BYTE** ppYUVData;

		if (use444)
		{
			iStride = h264->iYUV444Stride;
			ppYUVData = h264->pYUV444Data;
		}
		else
		{
			iStride = h264->iStride;
			ppYUVData = h264->pYUVData;
		}

		if (!check_rect(h264, rect, nDstWidth, nDstHeight))
			return FALSE;

		width = rect->right - rect->left;
		height = rect->bottom - rect->top;
		pDstPoint = pDstData + rect->top * nDstStep + rect->left * 4;
		pYUVPoint[0] = ppYUVData[0] + rect->top * iStride[0] + rect->left;
		pYUVPoint[1] = ppYUVData[1];
		pYUVPoint[2] = ppYUVData[2];

		if (use444)
		{
			pYUVPoint[1] += rect->top * iStride[1] + rect->left;
			pYUVPoint[2] += rect->top * iStride[2] + rect->left;
		}
		else
		{
			pYUVPoint[1] += rect->top / 2 * iStride[1] + rect->left / 2;
			pYUVPoint[2] += rect->top / 2 * iStride[2] + rect->left / 2;
		}

		roi.width = width;
		roi.height = height;

		if (use444)
		{
			if (prims->YUV444ToRGB_8u_P3AC4R(
			        pYUVPoint, iStride, pDstPoint, nDstStep,
			        DstFormat, &roi) != PRIMITIVES_SUCCESS)
			{
				return FALSE;
			}
		}
		else
		{
			if (prims->YUV420ToRGB_8u_P3AC4R(
			        pYUVPoint, iStride, pDstPoint, nDstStep,
			        DstFormat, &roi) != PRIMITIVES_SUCCESS)
				return FALSE;
		}
	}

	return TRUE;
}

INT32 avc420_decompress(H264_CONTEXT* h264, const BYTE* pSrcData, UINT32 SrcSize,
                        BYTE* pDstData, DWORD DstFormat, UINT32 nDstStep,
                        UINT32 nDstWidth, UINT32 nDstHeight,
                        RECTANGLE_16* regionRects, UINT32 numRegionRects)
{
	int status;

	if (!h264)
		return -1001;

	status = h264->subsystem->Decompress(h264, pSrcData, SrcSize);

	if (status == 0)
		return 1;

	if (status < 0)
		return status;

	if (!avc_yuv_to_rgb(h264, regionRects, numRegionRects, nDstWidth,
	                    nDstHeight, nDstStep, pDstData, DstFormat, FALSE))
		return -1002;

	return 1;
}

INT32 avc420_compress(H264_CONTEXT* h264, const BYTE* pSrcData, DWORD SrcFormat,
                      UINT32 nSrcStep, UINT32 nSrcWidth, UINT32 nSrcHeight,
                      BYTE** ppDstData, UINT32* pDstSize)
{
	prim_size_t roi;
	primitives_t* prims = primitives_get();

	if (!h264)
		return -1;

	if (!h264->subsystem->Compress)
		return -1;

	if (!avc420_ensure_buffer(h264, nSrcStep, nSrcWidth, nSrcHeight))
		return -1;

	roi.width = nSrcWidth;
	roi.height = nSrcHeight;

	if (prims->RGBToYUV420_8u_P3AC4R(pSrcData, SrcFormat, nSrcStep, h264->pYUVData, h264->iStride,
	                                 &roi) != PRIMITIVES_SUCCESS)
		return -1;

	{
		const BYTE* pYUVData[3] = {h264->pYUVData[0], h264->pYUVData[1], h264->pYUVData[2]};
		return h264->subsystem->Compress(h264, pYUVData, h264->iStride, ppDstData, pDstSize);
	}
}

INT32 avc444_compress(H264_CONTEXT* h264, const BYTE* pSrcData, DWORD SrcFormat,
                      UINT32 nSrcStep, UINT32 nSrcWidth, UINT32 nSrcHeight,
                      BYTE version, BYTE* op, BYTE** ppDstData, UINT32* pDstSize,
                      BYTE** ppAuxDstData, UINT32* pAuxDstSize)
{
	prim_size_t roi;
	primitives_t* prims = primitives_get();
	BYTE* coded;
	UINT32 codedSize;

	if (!h264)
		return -1;

	if (!h264->subsystem->Compress)
		return -1;

	if (!avc420_ensure_buffer(h264, nSrcStep, nSrcWidth, nSrcHeight))
		return -1;

	if (!avc444_ensure_buffer(h264, nSrcHeight))
		return -1;

	roi.width = nSrcWidth;
	roi.height = nSrcHeight;

	switch (version)
	{
		case 1:
			if (prims->RGBToAVC444YUV(pSrcData, SrcFormat, nSrcStep, h264->pYUV444Data, h264->iStride,
			                          h264->pYUVData, h264->iStride, &roi) != PRIMITIVES_SUCCESS)
				return -1;

			break;

		case 2:
			if (prims->RGBToAVC444YUVv2(pSrcData, SrcFormat, nSrcStep, h264->pYUV444Data, h264->iStride,
			                            h264->pYUVData, h264->iStride, &roi) != PRIMITIVES_SUCCESS)
				return -1;

			break;

		default:
			return -1;
	}

	{
		const BYTE* pYUV444Data[3] = {h264->pYUV444Data[0], h264->pYUV444Data[1], h264->pYUV444Data[2]};

		if (h264->subsystem->Compress(h264, pYUV444Data, h264->iStride, &coded, &codedSize) < 0)
			return -1;
	}

	memcpy(h264->lumaData, coded, codedSize);
	*ppDstData = h264->lumaData;
	*pDstSize = codedSize;
	{
		const BYTE* pYUVData[3] = {h264->pYUVData[0], h264->pYUVData[1], h264->pYUVData[2]};

		if (h264->subsystem->Compress(h264, pYUVData, h264->iStride, &coded, &codedSize) < 0)
			return -1;
	}
	*ppAuxDstData = coded;
	*pAuxDstSize = codedSize;
	*op = 0;
	return 0;
}

static BOOL avc444_ensure_buffer(H264_CONTEXT* h264,
                                 DWORD nDstHeight)
{
	UINT32 x;
	const UINT32* piMainStride = h264->iStride;
	UINT32* piDstSize = h264->iYUV444Size;
	UINT32* piDstStride = h264->iYUV444Stride;
	BYTE** ppYUVDstData = h264->pYUV444Data;
	UINT32 padDstHeight = nDstHeight + 16; /* Need alignment to 16x16 blocks */

	if ((piMainStride[0] != piDstStride[0]) ||
	    (piDstSize[0] != piMainStride[0] * padDstHeight))
	{
		for (x = 0; x < 3; x++)
		{
			piDstStride[x] = piMainStride[0];
			piDstSize[x] = piDstStride[x] * padDstHeight;
			_aligned_free(ppYUVDstData[x]);
			ppYUVDstData[x] = _aligned_malloc(piDstSize[x], 16);

			if (!ppYUVDstData[x])
				goto fail;

			memset(ppYUVDstData[x], 0, piDstSize[x]);
		}

		_aligned_free(h264->lumaData);
		h264->lumaData = _aligned_malloc(piDstSize[0] * 4, 16);
	}

	for (x = 0; x < 3; x++)
	{
		if (!ppYUVDstData[x] || (piDstSize[x] == 0) || (piDstStride[x] == 0))
		{
			WLog_Print(h264->log, WLOG_ERROR, "YUV buffer not initialized! check your decoder settings");
			goto fail;
		}
	}

	if (!h264->lumaData)
		goto fail;

	return TRUE;
fail:
	_aligned_free(ppYUVDstData[0]);
	_aligned_free(ppYUVDstData[1]);
	_aligned_free(ppYUVDstData[2]);
	_aligned_free(h264->lumaData);
	ppYUVDstData[0] = NULL;
	ppYUVDstData[1] = NULL;
	ppYUVDstData[2] = NULL;
	h264->lumaData = NULL;
	return FALSE;
}

static BOOL avc444_process_rects(H264_CONTEXT* h264, const BYTE* pSrcData,
                                 UINT32 SrcSize, BYTE* pDstData, UINT32 DstFormat, UINT32 nDstStep,
                                 UINT32 nDstWidth, UINT32 nDstHeight,
                                 const RECTANGLE_16* rects, UINT32 nrRects,
                                 avc444_frame_type type)
{
	const primitives_t* prims = primitives_get();
	UINT32 x;
	UINT32* piDstStride = h264->iYUV444Stride;
	BYTE** ppYUVDstData = h264->pYUV444Data;
	const UINT32* piStride = h264->iStride;
	const BYTE** ppYUVData = (const BYTE**)h264->pYUVData;

	if (h264->subsystem->Decompress(h264, pSrcData, SrcSize) < 0)
		return FALSE;

	if (!avc444_ensure_buffer(h264, nDstHeight))
		return FALSE;

	for (x = 0; x < nrRects; x++)
	{
		const RECTANGLE_16* rect = &rects[x];
		const UINT32 alignedWidth = h264->width + ((h264->width % 16 != 0) ? 16 - h264->width % 16 : 0);
		const UINT32 alignedHeight = h264->height + ((h264->height % 16 != 0) ? 16 - h264->height % 16 : 0);

		if (!check_rect(h264, rect, nDstWidth, nDstHeight))
			continue;

		if (prims->YUV420CombineToYUV444(type, ppYUVData, piStride,
		                                 alignedWidth, alignedHeight,
		                                 ppYUVDstData, piDstStride,
		                                 rect) != PRIMITIVES_SUCCESS)
			return FALSE;
	}

	if (!avc_yuv_to_rgb(h264, rects, nrRects, nDstWidth,
	                    nDstHeight, nDstStep, pDstData, DstFormat, TRUE))
		return FALSE;

	return TRUE;
}

#if defined(AVC444_FRAME_STAT)
static UINT64 op1 = 0;
static double op1sum = 0;
static UINT64 op2 = 0;
static double op2sum = 0;
static UINT64 op3 = 0;
static double op3sum = 0;
static double avg(UINT64* count, double old, double size)
{
	double tmp = size + *count * old;
	(*count)++;
	tmp = tmp / *count;
	return tmp;
}
#endif

INT32 avc444_decompress(H264_CONTEXT* h264, BYTE op,
                        RECTANGLE_16* regionRects, UINT32 numRegionRects,
                        const BYTE* pSrcData, UINT32 SrcSize,
                        RECTANGLE_16* auxRegionRects, UINT32 numAuxRegionRect,
                        const BYTE* pAuxSrcData, UINT32 AuxSrcSize,
                        BYTE* pDstData, DWORD DstFormat,
                        UINT32 nDstStep, UINT32 nDstWidth, UINT32 nDstHeight,
                        UINT32 codecId)
{
	INT32 status = -1;
	avc444_frame_type chroma = (codecId == RDPGFX_CODECID_AVC444) ? AVC444_CHROMAv1 : AVC444_CHROMAv2;

	if (!h264 || !regionRects ||
	    !pSrcData || !pDstData)
		return -1001;

	switch (op)
	{
		case 0: /* YUV420 in stream 1
		 * Chroma420 in stream 2 */
			if (!avc444_process_rects(h264, pSrcData, SrcSize, pDstData, DstFormat, nDstStep, nDstWidth,
			                          nDstHeight,
			                          regionRects, numRegionRects, AVC444_LUMA))
				status = -1;
			else if (!avc444_process_rects(h264, pAuxSrcData, AuxSrcSize, pDstData, DstFormat, nDstStep,
			                               nDstWidth, nDstHeight,
			                               auxRegionRects, numAuxRegionRect, chroma))
				status = -1;
			else
				status = 0;

			break;

		case 2: /* Chroma420 in stream 1 */
			if (!avc444_process_rects(h264, pSrcData, SrcSize, pDstData, DstFormat, nDstStep, nDstWidth,
			                          nDstHeight,
			                          regionRects, numRegionRects, chroma))
				status = -1;
			else
				status = 0;

			break;

		case 1: /* YUV420 in stream 1 */
			if (!avc444_process_rects(h264, pSrcData, SrcSize, pDstData, DstFormat, nDstStep, nDstWidth,
			                          nDstHeight,
			                          regionRects, numRegionRects, AVC444_LUMA))
				status = -1;
			else
				status = 0;

			break;

		default: /* WTF? */
			break;
	}

#if defined(AVC444_FRAME_STAT)

	switch (op)
	{
		case 0:
			op1sum = avg(&op1, op1sum, SrcSize + AuxSrcSize);
			break;

		case 1:
			op2sum = avg(&op2, op2sum, SrcSize);
			break;

		case 2:
			op3sum = avg(&op3, op3sum, SrcSize);
			break;

		default:
			break;
	}

	WLog_Print(h264->log, WLOG_INFO,
	           "luma=%"PRIu64" [avg=%lf] chroma=%"PRIu64" [avg=%lf] combined=%"PRIu64" [avg=%lf]",
	           op1, op1sum, op2, op2sum, op3, op3sum);
#endif
	return status;
}

#define MAX_SUBSYSTEMS 10
static INIT_ONCE subsystems_once = INIT_ONCE_STATIC_INIT;
static H264_CONTEXT_SUBSYSTEM* subSystems[MAX_SUBSYSTEMS];

#if defined(_WIN32) && defined(WITH_MEDIA_FOUNDATION)
extern H264_CONTEXT_SUBSYSTEM g_Subsystem_MF;
#endif

static BOOL CALLBACK h264_register_subsystems(PINIT_ONCE once, PVOID param, PVOID* context)
{
	int i = 0;
	ZeroMemory(subSystems, sizeof(subSystems));
#if defined(_WIN32) && defined(WITH_MEDIA_FOUNDATION)
	{
		subSystems[i] = &g_Subsystem_MF;
		i++;
	}
#endif
#ifdef WITH_OPENH264
	{
		extern H264_CONTEXT_SUBSYSTEM g_Subsystem_OpenH264;
		subSystems[i] = &g_Subsystem_OpenH264;
		i++;
	}
#endif
#ifdef WITH_FFMPEG
	{
		extern H264_CONTEXT_SUBSYSTEM g_Subsystem_libavcodec;
		subSystems[i] = &g_Subsystem_libavcodec;
		i++;
	}
#endif
#ifdef WITH_X264
	{
		extern H264_CONTEXT_SUBSYSTEM g_Subsystem_x264;
		subSystems[i] = &g_Subsystem_x264;
		i++;
	}
#endif
	return i > 0;
}


BOOL h264_context_init(H264_CONTEXT* h264)
{
	int i;

	if (!h264)
		return FALSE;

	h264->log = WLog_Get(TAG);

	if (!h264->log)
		return FALSE;

	h264->subsystem = NULL;
	InitOnceExecuteOnce(&subsystems_once, h264_register_subsystems, NULL, NULL);

	for (i = 0; i < MAX_SUBSYSTEMS; i++)
	{
		H264_CONTEXT_SUBSYSTEM* subsystem = subSystems[i];

		if (!subsystem || !subsystem->Init)
			break;

		if (subsystem->Init(h264))
		{
			h264->subsystem = subsystem;
			return TRUE;
		}
	}

	return FALSE;
}

BOOL h264_context_reset(H264_CONTEXT* h264, UINT32 width, UINT32 height)
{
	if (!h264)
		return FALSE;

	h264->width = width;
	h264->height = height;
	return TRUE;
}

H264_CONTEXT* h264_context_new(BOOL Compressor)
{
	H264_CONTEXT* h264;
	h264 = (H264_CONTEXT*) calloc(1, sizeof(H264_CONTEXT));

	if (h264)
	{
		h264->Compressor = Compressor;

		if (Compressor)
		{
			/* Default compressor settings, may be changed by caller */
			h264->BitRate = 1000000;
			h264->FrameRate = 30;
		}

		if (!h264_context_init(h264))
		{
			free(h264);
			return NULL;
		}
	}

	return h264;
}

void h264_context_free(H264_CONTEXT* h264)
{
	if (h264)
	{
		h264->subsystem->Uninit(h264);
		_aligned_free(h264->pYUV444Data[0]);
		_aligned_free(h264->pYUV444Data[1]);
		_aligned_free(h264->pYUV444Data[2]);
		_aligned_free(h264->lumaData);
		free(h264);
	}
}