#include <winpr/sysinfo.h>

#include <winpr/pool.h>

#include <freerdp/primitives.h>

#include <freerdp/log.h>

#include <freerdp/codec/yuv.h>

#define TAG FREERDP_TAG("codec")

struct _YUV_CONTEXT

{

	UINT32 width, height;

	BOOL useThreads;

	UINT32 nthreads;

	UINT32 heightStep;

	PTP_POOL threadPool;

	TP_CALLBACK_ENVIRON ThreadPoolEnv;

};

struct _YUV_PROCESS_WORK_PARAM

{

	YUV_CONTEXT* context;

	const BYTE* pYUVData[3];

	UINT32 iStride[3];

	DWORD DstFormat;

	BYTE* dest;

	UINT32 nDstStep;

	UINT32 y;

	UINT32 height;

};

typedef struct _YUV_PROCESS_WORK_PARAM YUV_PROCESS_WORK_PARAM;

static void CALLBACK yuv_process_work_callback(PTP_CALLBACK_INSTANCE instance, void* context,

                                               PTP_WORK work)

{

	prim_size_t roi;

	YUV_PROCESS_WORK_PARAM* param = (YUV_PROCESS_WORK_PARAM*)context;

	primitives_t* prims = primitives_get();

	roi.width = param->context->width;

	roi.height = param->height;

	if (prims->YUV420ToRGB_8u_P3AC4R(param->pYUVData, param->iStride, param->dest, param->nDstStep,

	                                 param->DstFormat, &roi) != PRIMITIVES_SUCCESS)

	{

		WLog_ERR(TAG, "error when decoding lines");

	}

}

void yuv_context_reset(YUV_CONTEXT* context, UINT32 width, UINT32 height)

{

	context->width = width;

	context->height = height;

	context->heightStep = (height / context->nthreads);

}

YUV_CONTEXT* yuv_context_new(BOOL encoder)

{

	SYSTEM_INFO sysInfos;

	YUV_CONTEXT* ret = calloc(1, sizeof(*ret));

	if (!ret)

		return NULL;

	/** do it here to avoid a race condition between threads */

	primitives_get();

	GetNativeSystemInfo(&sysInfos);

	ret->useThreads = (sysInfos.dwNumberOfProcessors > 1);

	if (ret->useThreads)

	{

		ret->nthreads = sysInfos.dwNumberOfProcessors;

		ret->threadPool = CreateThreadpool(NULL);

		if (!ret->threadPool)

		{

			goto error_threadpool;

		}

		InitializeThreadpoolEnvironment(&ret->ThreadPoolEnv);

		SetThreadpoolCallbackPool(&ret->ThreadPoolEnv, ret->threadPool);

	}

	else

	{

		ret->nthreads = 1;

	}

	return ret;

error_threadpool:

	free(ret);

	return NULL;

}

void yuv_context_free(YUV_CONTEXT* context)

{

	if (context->useThreads)

	{

		CloseThreadpool(context->threadPool);

		DestroyThreadpoolEnvironment(&context->ThreadPoolEnv);

	}

	free(context);

}

BOOL yuv_context_decode(YUV_CONTEXT* context, const BYTE* pYUVData[3], UINT32 iStride[3],

                        DWORD DstFormat, BYTE* dest, UINT32 nDstStep)

{

	UINT32 y, nobjects, i;

	PTP_WORK* work_objects = NULL;

	YUV_PROCESS_WORK_PARAM* params;

	UINT32 waitCount = 0;

	BOOL ret = TRUE;

	primitives_t* prims = primitives_get();

	if (!context->useThreads || (primitives_flags(prims) & PRIM_FLAGS_HAVE_EXTGPU))

	{

		prim_size_t roi;

		roi.width = context->width;

		roi.height = context->height;

		return prims->YUV420ToRGB_8u_P3AC4R(pYUVData, iStride, dest, nDstStep, DstFormat, &roi) ==

		       PRIMITIVES_SUCCESS;

	}

	/* case where we use threads */

	nobjects = (context->height + context->heightStep - 1) / context->heightStep;

	work_objects = (PTP_WORK*)calloc(nobjects, sizeof(PTP_WORK));

	if (!work_objects)

	{

		return FALSE;

	}

	params = (YUV_PROCESS_WORK_PARAM*)calloc(nobjects, sizeof(*params));

	if (!params)

	{

		free(work_objects);

		return FALSE;

	}

	for (i = 0, y = 0; i < nobjects; i++, y += context->heightStep, waitCount++)

	{

		params[i].context = context;

		params[i].DstFormat = DstFormat;

		params[i].pYUVData[0] = pYUVData[0] + (y * iStride[0]);

		params[i].pYUVData[1] = pYUVData[1] + ((y / 2) * iStride[1]);

		params[i].pYUVData[2] = pYUVData[2] + ((y / 2) * iStride[2]);

		params[i].iStride[0] = iStride[0];

		params[i].iStride[1] = iStride[1];

		params[i].iStride[2] = iStride[2];

		params[i].nDstStep = nDstStep;

		params[i].dest = dest + (nDstStep * y);

		params[i].y = y;

		if (y + context->heightStep <= context->height)

			params[i].height = context->heightStep;

		else

			params[i].height = context->height % context->heightStep;

		work_objects[i] = CreateThreadpoolWork(yuv_process_work_callback, (void*)&params[i],

		                                       &context->ThreadPoolEnv);

		if (!work_objects[i])

		{

			ret = FALSE;

			break;

		}

		SubmitThreadpoolWork(work_objects[i]);

	}

	for (i = 0; i < waitCount; i++)

	{

		WaitForThreadpoolWorkCallbacks(work_objects[i], FALSE);

		CloseThreadpoolWork(work_objects[i]);

	}

	free(work_objects);

	free(params);

	return ret;

}