/**
 * FreeRDP: A Remote Desktop Protocol Implementation
 * RemoteFX Codec Library - RLGR
 *
 * Copyright 2011 Vic Lee
 *
 * 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.
 */

/**
 * This implementation of RLGR refers to
 * [MS-RDPRFX] 3.1.8.1.7.3 RLGR1/RLGR3 Pseudocode
 */

#include <freerdp/config.h>

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <winpr/assert.h>
#include <winpr/cast.h>
#include <winpr/crt.h>
#include <winpr/print.h>
#include <winpr/sysinfo.h>
#include <winpr/bitstream.h>
#include <winpr/intrin.h>

#include "rfx_bitstream.h"

#include "rfx_rlgr.h"

/* Constants used in RLGR1/RLGR3 algorithm */
#define KPMAX (80) /* max value for kp or krp */
#define LSGR (3)   /* shift count to convert kp to k */
#define UP_GR (4)  /* increase in kp after a zero run in RL mode */
#define DN_GR (6)  /* decrease in kp after a nonzero symbol in RL mode */
#define UQ_GR (3)  /* increase in kp after nonzero symbol in GR mode */
#define DQ_GR (3)  /* decrease in kp after zero symbol in GR mode */

/* Returns the least number of bits required to represent a given value */
#define GetMinBits(_val, _nbits) \
	do                           \
	{                            \
		UINT32 _v = (_val);      \
		(_nbits) = 0;            \
		while (_v)               \
		{                        \
			_v >>= 1;            \
			(_nbits)++;          \
		}                        \
	} while (0)

/*
 * Update the passed parameter and clamp it to the range [0, KPMAX]
 * Return the value of parameter right-shifted by LSGR
 */
static inline uint32_t UpdateParam(uint32_t* param, int32_t deltaP)
{
	WINPR_ASSERT(param);
	if (deltaP < 0)
	{
		const uint32_t udeltaP = WINPR_ASSERTING_INT_CAST(uint32_t, -deltaP);
		if (udeltaP > *param)
			*param = 0;
		else
			*param -= udeltaP;
	}
	else
		*param += WINPR_ASSERTING_INT_CAST(uint32_t, deltaP);

	if ((*param) > KPMAX)
		(*param) = KPMAX;
	return (*param) >> LSGR;
}

static BOOL g_LZCNT = FALSE;

static INIT_ONCE rfx_rlgr_init_once = INIT_ONCE_STATIC_INIT;

static BOOL CALLBACK rfx_rlgr_init(PINIT_ONCE once, PVOID param, PVOID* context)
{
	WINPR_UNUSED(once);
	WINPR_UNUSED(param);
	WINPR_UNUSED(context);

	g_LZCNT = IsProcessorFeaturePresentEx(PF_EX_LZCNT);
	return TRUE;
}

static inline UINT32 lzcnt_s(UINT32 x)
{
	if (!x)
		return 32;

	if (!g_LZCNT)
	{
		UINT32 y = 0;
		UINT32 n = 32;
		y = x >> 16;
		if (y != 0)
		{
			WINPR_ASSERT(n >= 16);
			n = n - 16;
			x = y;
		}
		y = x >> 8;
		if (y != 0)
		{
			WINPR_ASSERT(n >= 8);
			n = n - 8;
			x = y;
		}
		y = x >> 4;
		if (y != 0)
		{
			WINPR_ASSERT(n >= 4);
			n = n - 4;
			x = y;
		}
		y = x >> 2;
		if (y != 0)
		{
			WINPR_ASSERT(n >= 2);
			n = n - 2;
			x = y;
		}
		y = x >> 1;
		if (y != 0)
		{
			WINPR_ASSERT(n >= 2);
			return n - 2;
		}

		WINPR_ASSERT(n >= x);
		return n - x;
	}

	return __lzcnt(x);
}

int rfx_rlgr_decode(RLGR_MODE mode, const BYTE* WINPR_RESTRICT pSrcData, UINT32 SrcSize,
                    INT16* WINPR_RESTRICT pDstData, UINT32 rDstSize)
{
	uint32_t vk = 0;
	size_t run = 0;
	size_t cnt = 0;
	size_t size = 0;
	size_t offset = 0;
	INT16 mag = 0;
	UINT32 k = 0;
	UINT32 kp = 0;
	UINT32 kr = 0;
	UINT32 krp = 0;
	UINT16 code = 0;
	UINT32 sign = 0;
	UINT32 nIdx = 0;
	UINT32 val1 = 0;
	UINT32 val2 = 0;
	INT16* pOutput = nullptr;
	wBitStream* bs = nullptr;
	wBitStream s_bs = WINPR_C_ARRAY_INIT;
	const SSIZE_T DstSize = rDstSize;

	if (!InitOnceExecuteOnce(&rfx_rlgr_init_once, rfx_rlgr_init, nullptr, nullptr))
		return -1;

	k = 1;
	kp = k << LSGR;

	kr = 1;
	krp = kr << LSGR;

	if ((mode != RLGR1) && (mode != RLGR3))
		mode = RLGR1;

	if (!pSrcData || !SrcSize)
		return -1;

	if (!pDstData || !DstSize)
		return -1;

	pOutput = pDstData;

	bs = &s_bs;

	BitStream_Attach(bs, pSrcData, SrcSize);
	BitStream_Fetch(bs);

	while ((BitStream_GetRemainingLength(bs) > 0) && ((pOutput - pDstData) < DstSize))
	{
		if (k)
		{
			/* Run-Length (RL) Mode */

			run = 0;

			/* count number of leading 0s */

			cnt = lzcnt_s(bs->accumulator);

			size_t nbits = BitStream_GetRemainingLength(bs);

			if (cnt > nbits)
				cnt = WINPR_ASSERTING_INT_CAST(uint32_t, nbits);

			vk = WINPR_ASSERTING_INT_CAST(uint32_t, cnt);

			while ((cnt == 32) && (BitStream_GetRemainingLength(bs) > 0))
			{
				BitStream_Shift32(bs);

				cnt = lzcnt_s(bs->accumulator);

				nbits = BitStream_GetRemainingLength(bs);

				if (cnt > nbits)
					cnt = nbits;

				WINPR_ASSERT(cnt + vk <= UINT32_MAX);
				vk += WINPR_ASSERTING_INT_CAST(uint32_t, cnt);
			}

			BitStream_Shift(bs, (vk % 32));

			if (BitStream_GetRemainingLength(bs) < 1)
				break;

			BitStream_Shift(bs, 1);

			while (vk--)
			{
				const UINT32 add = (1 << k); /* add (1 << k) to run length */
				run += add;

				/* update k, kp params */

				kp += UP_GR;

				if (kp > KPMAX)
					kp = KPMAX;

				k = kp >> LSGR;
			}

			/* next k bits contain run length remainder */

			if (BitStream_GetRemainingLength(bs) < k)
				break;

			bs->mask = ((1 << k) - 1);
			run += ((bs->accumulator >> (32 - k)) & bs->mask);
			BitStream_Shift(bs, k);

			/* read sign bit */

			if (BitStream_GetRemainingLength(bs) < 1)
				break;

			sign = (bs->accumulator & 0x80000000) ? 1 : 0;
			BitStream_Shift(bs, 1);

			/* count number of leading 1s */

			cnt = lzcnt_s(~(bs->accumulator));

			nbits = BitStream_GetRemainingLength(bs);

			if (cnt > nbits)
				cnt = nbits;

			vk = WINPR_ASSERTING_INT_CAST(uint32_t, cnt);

			while ((cnt == 32) && (BitStream_GetRemainingLength(bs) > 0))
			{
				BitStream_Shift32(bs);

				cnt = lzcnt_s(~(bs->accumulator));

				nbits = BitStream_GetRemainingLength(bs);

				if (cnt > nbits)
					cnt = nbits;

				WINPR_ASSERT(cnt + vk <= UINT32_MAX);
				vk += WINPR_ASSERTING_INT_CAST(uint32_t, cnt);
			}

			BitStream_Shift(bs, (vk % 32));

			if (BitStream_GetRemainingLength(bs) < 1)
				break;

			BitStream_Shift(bs, 1);

			/* next kr bits contain code remainder */

			if (BitStream_GetRemainingLength(bs) < kr)
				break;

			bs->mask = ((1 << kr) - 1);
			if (kr > 0)
				code = (UINT16)((bs->accumulator >> (32 - kr)) & bs->mask);
			else
				code = 0;
			BitStream_Shift(bs, kr);

			/* add (vk << kr) to code */

			code |= (vk << kr);

			if (!vk)
			{
				/* update kr, krp params */

				if (krp > 2)
					krp -= 2;
				else
					krp = 0;

				kr = krp >> LSGR;
			}
			else if (vk != 1)
			{
				/* update kr, krp params */

				krp += vk;

				if (krp > KPMAX)
					krp = KPMAX;

				kr = krp >> LSGR;
			}

			/* update k, kp params */

			if (kp > DN_GR)
				kp -= DN_GR;
			else
				kp = 0;

			k = kp >> LSGR;

			/* compute magnitude from code */

			if (sign)
				mag = WINPR_ASSERTING_INT_CAST(int16_t, (code + 1)) * -1;
			else
				mag = WINPR_ASSERTING_INT_CAST(int16_t, code + 1);

			/* write to output stream */

			offset = WINPR_ASSERTING_INT_CAST(size_t, (pOutput)-pDstData);
			size = run;

			if ((offset + size) > rDstSize)
				size = WINPR_ASSERTING_INT_CAST(size_t, DstSize) - offset;

			if (size)
			{
				ZeroMemory(pOutput, size * sizeof(INT16));
				pOutput += size;
			}

			if ((pOutput - pDstData) < DstSize)
			{
				*pOutput = mag;
				pOutput++;
			}
		}
		else
		{
			/* Golomb-Rice (GR) Mode */

			/* count number of leading 1s */

			cnt = lzcnt_s(~(bs->accumulator));

			size_t nbits = BitStream_GetRemainingLength(bs);

			if (cnt > nbits)
				cnt = nbits;

			vk = WINPR_ASSERTING_INT_CAST(uint32_t, cnt);

			while ((cnt == 32) && (BitStream_GetRemainingLength(bs) > 0))
			{
				BitStream_Shift32(bs);

				cnt = lzcnt_s(~(bs->accumulator));

				nbits = BitStream_GetRemainingLength(bs);

				if (cnt > nbits)
					cnt = nbits;

				WINPR_ASSERT(cnt + vk <= UINT32_MAX);
				vk += WINPR_ASSERTING_INT_CAST(uint32_t, cnt);
			}

			BitStream_Shift(bs, (vk % 32));

			if (BitStream_GetRemainingLength(bs) < 1)
				break;

			BitStream_Shift(bs, 1);

			/* next kr bits contain code remainder */

			if (BitStream_GetRemainingLength(bs) < kr)
				break;

			bs->mask = ((1 << kr) - 1);
			if (kr > 0)
				code = (UINT16)((bs->accumulator >> (32 - kr)) & bs->mask);
			else
				code = 0;
			BitStream_Shift(bs, kr);

			/* add (vk << kr) to code */

			code |= (vk << kr);

			if (!vk)
			{
				/* update kr, krp params */

				if (krp > 2)
					krp -= 2;
				else
					krp = 0;

				kr = (krp >> LSGR) & UINT32_MAX;
			}
			else if (vk != 1)
			{
				/* update kr, krp params */

				krp += vk;

				if (krp > KPMAX)
					krp = KPMAX;

				kr = krp >> LSGR;
			}

			if (mode == RLGR1) /* RLGR1 */
			{
				if (!code)
				{
					/* update k, kp params */

					kp += UQ_GR;

					if (kp > KPMAX)
						kp = KPMAX;

					k = kp >> LSGR;

					mag = 0;
				}
				else
				{
					/* update k, kp params */

					if (kp > DQ_GR)
						kp -= DQ_GR;
					else
						kp = 0;

					k = kp >> LSGR;

					/*
					 * code = 2 * mag - sign
					 * sign + code = 2 * mag
					 */

					if (code & 1)
						mag = WINPR_ASSERTING_INT_CAST(INT16, (code + 1) >> 1) * -1;
					else
						mag = WINPR_ASSERTING_INT_CAST(INT16, code >> 1);
				}

				if ((pOutput - pDstData) < DstSize)
				{
					*pOutput = mag;
					pOutput++;
				}
			}
			else if (mode == RLGR3) /* RLGR3 */
			{
				nIdx = 0;

				if (code)
				{
					mag = WINPR_ASSERTING_INT_CAST(int16_t, code);
					nIdx = 32 - lzcnt_s(WINPR_ASSERTING_INT_CAST(uint32_t, mag));
				}

				if (BitStream_GetRemainingLength(bs) < nIdx)
					break;

				bs->mask = ((1 << nIdx) - 1);
				if (nIdx > 0)
					val1 = ((bs->accumulator >> (32 - nIdx)) & bs->mask);
				else
					val1 = 0;
				BitStream_Shift(bs, nIdx);

				val2 = code - val1;

				if (val1 && val2)
				{
					/* update k, kp params */

					if (kp > 2 * DQ_GR)
						kp -= (2 * DQ_GR);
					else
						kp = 0;

					k = kp >> LSGR;
				}
				else if (!val1 && !val2)
				{
					/* update k, kp params */

					kp += (2 * UQ_GR);

					if (kp > KPMAX)
						kp = KPMAX;

					k = kp >> LSGR;
				}

				if (val1 & 1)
					mag = WINPR_ASSERTING_INT_CAST(int16_t, (val1 + 1) >> 1) * -1;
				else
					mag = WINPR_ASSERTING_INT_CAST(int16_t, val1 >> 1);

				if ((pOutput - pDstData) < DstSize)
				{
					*pOutput = mag;
					pOutput++;
				}

				if (val2 & 1)
					mag = WINPR_ASSERTING_INT_CAST(int16_t, (val2 + 1) >> 1) * -1;
				else
					mag = WINPR_ASSERTING_INT_CAST(int16_t, val2 >> 1);

				if ((pOutput - pDstData) < DstSize)
				{
					*pOutput = WINPR_ASSERTING_INT_CAST(int16_t, mag);
					pOutput++;
				}
			}
		}
	}

	offset = WINPR_ASSERTING_INT_CAST(size_t, (pOutput - pDstData));

	if (offset < rDstSize)
	{
		size = WINPR_ASSERTING_INT_CAST(size_t, DstSize) - offset;
		ZeroMemory(pOutput, size * 2);
		pOutput += size;
	}

	offset = WINPR_ASSERTING_INT_CAST(size_t, (pOutput - pDstData));

	if ((DstSize < 0) || (offset != (size_t)DstSize))
		return -1;

	return 1;
}

/* Returns the next coefficient (a signed int) to encode, from the input stream */
#define GetNextInput(_n)    \
	do                      \
	{                       \
		if (data_size > 0)  \
		{                   \
			(_n) = *data++; \
			data_size--;    \
		}                   \
		else                \
		{                   \
			(_n) = 0;       \
		}                   \
	} while (0)

/* Emit bitPattern to the output bitstream */
#define OutputBits(numBits, bitPattern) rfx_bitstream_put_bits(bs, bitPattern, numBits)

/* Emit a bit (0 or 1), count number of times, to the output bitstream */
static inline void OutputBit(RFX_BITSTREAM* bs, uint32_t count, UINT8 bit)
{
	UINT16 _b = ((bit) ? 0xFFFF : 0);
	const uint32_t rem = count % 16;
	for (uint32_t x = 0; x < count - rem; x += 16)
		rfx_bitstream_put_bits(bs, _b, 16);

	if (rem > 0)
		rfx_bitstream_put_bits(bs, _b, rem);
}

/* Converts the input value to (2 * abs(input) - sign(input)), where sign(input) = (input < 0 ? 1 :
 * 0) and returns it */
static inline UINT32 Get2MagSign(INT32 input)
{
	if (input >= 0)
		return WINPR_ASSERTING_INT_CAST(UINT32, 2 * input);
	return WINPR_ASSERTING_INT_CAST(UINT32, -2 * input - 1);
}

/* Outputs the Golomb/Rice encoding of a non-negative integer */
#define CodeGR(krp, val) rfx_rlgr_code_gr(bs, krp, val)

static void rfx_rlgr_code_gr(RFX_BITSTREAM* bs, uint32_t* krp, UINT32 val)
{
	uint32_t kr = *krp >> LSGR;

	/* unary part of GR code */

	const uint32_t vk = val >> kr;
	OutputBit(bs, vk, 1);
	OutputBit(bs, 1, 0);

	/* remainder part of GR code, if needed */
	if (kr)
	{
		OutputBits(kr, val & ((1 << kr) - 1));
	}

	/* update krp, only if it is not equal to 1 */
	if (vk == 0)
	{
		(void)UpdateParam(krp, -2);
	}
	else if (vk > 1)
	{
		(void)UpdateParam(krp, WINPR_CXX_COMPAT_CAST(int32_t, vk));
	}
}

int rfx_rlgr_encode(RLGR_MODE mode, const INT16* WINPR_RESTRICT data, UINT32 data_size,
                    BYTE* WINPR_RESTRICT buffer, UINT32 buffer_size)
{
	uint32_t k = 0;
	uint32_t kp = 0;
	uint32_t krp = 0;
	RFX_BITSTREAM* bs = nullptr;

	if (!(bs = (RFX_BITSTREAM*)winpr_aligned_calloc(1, sizeof(RFX_BITSTREAM), 32)))
		return 0;

	rfx_bitstream_attach(bs, buffer, buffer_size);

	/* initialize the parameters */
	k = 1;
	kp = 1 << LSGR;
	krp = 1 << LSGR;

	/* process all the input coefficients */
	while (data_size > 0)
	{
		int input = 0;

		if (k)
		{
			uint32_t numZeros = 0;
			uint32_t runmax = 0;
			BYTE sign = 0;

			/* RUN-LENGTH MODE */

			/* collect the run of zeros in the input stream */
			numZeros = 0;
			GetNextInput(input);
			while (input == 0 && data_size > 0)
			{
				numZeros++;
				GetNextInput(input);
			}

			// emit output zeros
			runmax = 1 << k;
			while (numZeros >= runmax)
			{
				OutputBit(bs, 1, 0); /* output a zero bit */
				numZeros -= runmax;
				k = UpdateParam(&kp, UP_GR); /* update kp, k */
				runmax = 1 << k;
			}

			/* output a 1 to terminate runs */
			OutputBit(bs, 1, 1);

			/* output the remaining run length using k bits */
			OutputBits(k, numZeros);

			/* note: when we reach here and the last byte being encoded is 0, we still
			   need to output the last two bits, otherwise mstsc will crash */

			/* encode the nonzero value using GR coding */
			const UINT32 mag =
			    (UINT32)(input < 0 ? -input : input); /* absolute value of input coefficient */
			sign = (input < 0 ? 1 : 0);         /* sign of input coefficient */

			OutputBit(bs, 1, sign);          /* output the sign bit */
			CodeGR(&krp, mag ? mag - 1 : 0); /* output GR code for (mag - 1) */

			k = UpdateParam(&kp, -DN_GR);
		}
		else
		{
			/* GOLOMB-RICE MODE */

			if (mode == RLGR1)
			{
				UINT32 twoMs = 0;

				/* RLGR1 variant */

				/* convert input to (2*magnitude - sign), encode using GR code */
				GetNextInput(input);
				twoMs = Get2MagSign(input);
				CodeGR(&krp, twoMs);

				/* update k, kp */
				/* NOTE: as of Aug 2011, the algorithm is still wrongly documented
				   and the update direction is reversed */
				if (twoMs)
				{
					k = UpdateParam(&kp, -DQ_GR);
				}
				else
				{
					k = UpdateParam(&kp, UQ_GR);
				}
			}
			else /* mode == RLGR3 */
			{
				UINT32 twoMs1 = 0;
				UINT32 twoMs2 = 0;
				UINT32 sum2Ms = 0;
				UINT32 nIdx = 0;

				/* RLGR3 variant */

				/* convert the next two input values to (2*magnitude - sign) and */
				/* encode their sum using GR code */

				GetNextInput(input);
				twoMs1 = Get2MagSign(input);
				GetNextInput(input);
				twoMs2 = Get2MagSign(input);
				sum2Ms = twoMs1 + twoMs2;

				CodeGR(&krp, sum2Ms);

				/* encode binary representation of the first input (twoMs1). */
				GetMinBits(sum2Ms, nIdx);
				OutputBits(nIdx, twoMs1);

				/* update k,kp for the two input values */

				if (twoMs1 && twoMs2)
				{
					k = UpdateParam(&kp, -2 * DQ_GR);
				}
				else if (!twoMs1 && !twoMs2)
				{
					k = UpdateParam(&kp, 2 * UQ_GR);
				}
			}
		}
	}

	rfx_bitstream_flush(bs);
	uint32_t processed_size = rfx_bitstream_get_processed_bytes(bs);
	winpr_aligned_free(bs);

	return WINPR_ASSERTING_INT_CAST(int, processed_size);
}