/*-------------------------------------------------------------------------
 * Copyright (c) Microsoft Corporation.  All rights reserved.
 *
 * src/aggregation/bson_sorted_accumulator.c
 *
 * Functions related to custom aggregates for group by
 * accumulator with sort specification:
 * first (implicit)
 * last (implicit)
 * top (explicit)
 * bottom (explicit)
 * firstN (implicit)
 * lastN (implicit)
 * topN (explicit)
 * bottomN (explicit)
 *-------------------------------------------------------------------------
 */

#include <postgres.h>
#include <funcapi.h>


#include "io/bson_core.h"
#include "query/bson_compare.h"
#include "query/bson_dollar_operators.h"
#include "operators/bson_expression.h"
#include "aggregation/bson_sorted_accumulator.h"
#include "utils/array.h"
#include "utils/fmgr_utils.h"


/* --------------------------------------------------------- */
/* Forward declaration */
/* --------------------------------------------------------- */
static void ParseInputExpressionAndPersistValue(AggregationExpressionData *expressionData,
												const bson_value_t *expressionValue,
												ParseAggregationExpressionContext *context);


/*
 * Converts a BsonOrderAggState into a serialized form to allow the internal type to be bytea
 * Resulting bytes look like:
 * | Varlena Header | numAggValues | numSortKeys | BsonOrderAggValue * numAggValues | sortDirections | inputExpression (pgbson)
 *
 * The function uses the CurrentMemoryContext (ecxt_per_tuple_memory context) to allocate the new set bytea for result, because we want
 * this to be free'd as soon as possible, without leaking memory for entire duration of aggregate.
 * Postgres makes sure to copy it into aggregate context if needed.
 *
 * For more info look at these places in PG where the state pointer is tracked and cleaned up:
 * nodeAgg.c: https://github.com/postgres/postgres/blob/9c24111c4dad850bc2625c2113bb3d2dfd592efc/src/backend/executor/nodeAgg.c#L779
 * nodeWindowAgg.c: https://github.com/postgres/postgres/blob/9c24111c4dad850bc2625c2113bb3d2dfd592efc/src/backend/executor/nodeWindowAgg.c#L367
 *
 */
pg_attribute_no_sanitize_alignment() bytea *
SerializeOrderState(BsonOrderAggState * state)
{
	int valueSize = 0;
	for (int i = 0; i < state->currentCount; i++)
	{
		/* Flag for not null */
		valueSize += 1;

		/* Get total size of currentResult. Null values are marked with 0 byte */
		if (state->currentResult[i] != NULL)
		{
			/* flag for value not null */
			valueSize += 1;
			if (state->currentResult[i]->value != NULL)
			{
				valueSize += VARSIZE(state->currentResult[i]->value);
			}

			int sortKeyValuesSize = 0;

			/* Accumulate the total size of the sortKeys, nulls are marked with a 0 byte */
			for (int j = 0; j < state->numSortKeys; j++)
			{
				/* Flag for sort key not null */
				sortKeyValuesSize += 1;
				if (state->currentResult[i]->sortKeyValues[j] == 0)
				{
					continue;
				}
				pgbson *bson = DatumGetPgBson(state->currentResult[i]->sortKeyValues[j]);
				sortKeyValuesSize += VARSIZE(bson);
			}
			valueSize += sortKeyValuesSize;
		}
	}

	/* We account for the null flag byte by default. */
	int inputExpressionValueSize = 1;
	if (state->inputExpression != NULL)
	{
		inputExpressionValueSize += VARSIZE(state->inputExpression);
	}

	/*
	 * Make sure the byte array is big enough to fit the whole state if numAggValues is 1:
	 * varlena Header | numAggValues | numSortKeys  | BsonOrderAggValues * numAggValues | sortDirections (bool) | inputExpression (pgbson)
	 *
	 * We can't reuse the same block with muliple values as we would potentially
	 * overwrite the existing values if it was shifted to later in the list.
	 * Sometimes we need to increase the size with single values, e.g., when
	 * the new currentValue becomes bigger. In that case we repalloc.
	 */
	int requiredByteSize = VARHDRSZ + /*varlena Header */
						   sizeof(int64) + /* numAggValues */
						   sizeof(int64) + /* currentCount */
						   sizeof(int) + /* numSortKeys */
						   valueSize + /* BsonOrderAggValues */
						   state->numSortKeys * sizeof(bool) + /* sortDirections */
						   inputExpressionValueSize; /* inputExpression (pgbson) + null flag */

	/* We should always allocate a new block of memory as the old one is still being used in the state that is referred later here.
	 * We don't need to do anything with the old memory (free/repalloc) since Postgres will free it as
	 * it will track that the memory block changed in the transition function.
	 * If we intent to free or repalloc the old memory, when Postgres tries to free it would crash as the old memory would be invalid.
	 */
	char *bytes = (char *) MemoryContextAlloc(CurrentMemoryContext, requiredByteSize);
	SET_VARSIZE(bytes, requiredByteSize);

	/* Copy into the current value variable */
	char *byteAllocationPointer = (char *) VARDATA(bytes);

	/* Set the number of Aggregation Values */
	memcpy(byteAllocationPointer, &state->numAggValues, sizeof(int64));
	byteAllocationPointer += sizeof(int64);

	/* Set the number of Current Values */
	memcpy(byteAllocationPointer, &state->currentCount, sizeof(int64));
	byteAllocationPointer += sizeof(int64);

	/* Set the number of sortKeys after the numAggValues */
	memcpy(byteAllocationPointer, &state->numSortKeys, sizeof(int));
	byteAllocationPointer += sizeof(int);

	/* Write each result */
	for (int i = 0; i < state->currentCount; i++)
	{
		/* Set Not Null for currentResult */
		*byteAllocationPointer = (char) (state->currentResult[i] != NULL);
		byteAllocationPointer += 1;
		if (state->currentResult[i] != NULL)
		{
			/* Set Not Null for currentResult */
			*byteAllocationPointer = (char) (state->currentResult[i]->value != NULL);
			byteAllocationPointer += 1;
			if (state->currentResult[i]->value != NULL)
			{
				int size = VARSIZE(state->currentResult[i]->value);
				memcpy(byteAllocationPointer, state->currentResult[i]->value, size);
				byteAllocationPointer += size;
			}

			/* Write each sortKey */
			for (int j = 0; j < state->numSortKeys; j++)
			{
				if (state->currentResult[i]->sortKeyValues[j] == 0)
				{
					*byteAllocationPointer = 0;
					byteAllocationPointer += 1;
					continue;
				}

				*byteAllocationPointer = 1;
				byteAllocationPointer += 1;
				pgbson *currentSortKey = DatumGetPgBson(
					state->currentResult[i]->sortKeyValues[j]);
				memcpy(byteAllocationPointer, currentSortKey, VARSIZE(currentSortKey));

				byteAllocationPointer += VARSIZE(currentSortKey);
			}
		}
	}

	/* Write sortDirection */
	for (int i = 0; i < state->numSortKeys; i++)
	{
		*((bool *) (byteAllocationPointer)) = state->sortDirections[i];
		byteAllocationPointer += sizeof(bool);
	}

	/* Write pgbson input expression spec. */
	*byteAllocationPointer = (char) (state->inputExpression != NULL);
	byteAllocationPointer += 1;

	if (state->inputExpression != NULL)
	{
		int size = VARSIZE(state->inputExpression);
		memcpy(byteAllocationPointer, state->inputExpression, size);
		byteAllocationPointer += size;
	}

	return (bytea *) bytes;
}


/*
 * Converts a BsonOrderAggState from a serialized form to allow the internal type to be bytea
 * Incoming bytes look like:
 * varlena Header | numAggValues | numSortKeys  | BsonOrderAggValues * numAggValues | sortDirections (bool) | inputExpression (pgbson)
 */
pg_attribute_no_sanitize_alignment() void
DeserializeOrderState(bytea *byteArray,
					  BsonOrderAggState *state)
{
	/* Use a char* instead for finer control over offsets */
	char *bytes = (char *) VARDATA(byteArray);

	/* Extract the number of numAggValues */
	memcpy(&state->numAggValues, bytes, sizeof(int64));
	bytes += sizeof(int64);

	/* Extract the number of current Values */
	memcpy(&state->currentCount, bytes, sizeof(int64));
	bytes += sizeof(int64);

	/* Check to see if we should allocate one extra position for the new data */
	if (state->currentCount >= state->numAggValues)
	{
		state->currentResult = palloc(sizeof(BsonOrderAggValue *) * state->currentCount);
	}
	else
	{
		state->currentResult = palloc(sizeof(BsonOrderAggValue *) * (state->currentCount +
																	 1));
	}

	/* Extract the number of sortKeys from after numAggValues */
	memcpy(&state->numSortKeys, bytes, sizeof(int));
	bytes += sizeof(int);

	/* Extract values from each current result */
	for (int i = 0; i < state->currentCount; i++)
	{
		if (*bytes == 0)
		{
			/* NULL value */
			state->currentResult[i] = NULL;
			bytes++;
		}
		else
		{
			/* Skip 1 byte for currentResult not null. */
			bytes++;
			state->currentResult[i] = palloc0(sizeof(BsonOrderAggValue));
			if (*bytes == 0)
			{
				state->currentResult[i]->value = NULL;
				bytes++;
			}
			else
			{
				bytes++;
				state->currentResult[i]->value = (pgbson *) (bytes);
				bytes += VARSIZE(state->currentResult[i]->value);
			}

			/* Extract each sortKey */
			for (int j = 0; j < state->numSortKeys; j++)
			{
				if (*bytes == 0)
				{
					state->currentResult[i]->sortKeyValues[j] = (Datum) 0;
					bytes++;
				}
				else
				{
					bytes++;
					state->currentResult[i]->sortKeyValues[j] = PointerGetDatum(
						(pgbson *) bytes);
					bytes += VARSIZE(state->currentResult[i]->sortKeyValues[j]);
				}
			}
		}
	}

	for (int i = 0; i < state->numSortKeys; i++)
	{
		state->sortDirections[i] = *((bool *) bytes);
		bytes += sizeof(bool);
	}

	if (*bytes == 0)
	{
		/* inputExpression is NULL */
		state->inputExpression = NULL;
		bytes++;
	}
	else
	{
		/* skip nullity byte flag */
		bytes++;
		state->inputExpression = (pgbson *) bytes;
		bytes += VARSIZE(state->inputExpression);
	}
}


/*
 * Applies the "state transition" (SFUNC) for sorted accumulators.
 * The args in PG_FUNCTION_ARGS:
 *      0) The current aggregation state
 *      1) The document/evaluated expression which will eventually be returned
 *
 * For isSingle == true:
 *      2) An array of bsons representing the sort specs.
 * For isSingle == true && storeInputExpression == true
 *      3) The input expression to be applied on the sorted results.
 *
 * For isSingle == false:
 *      2) The number of documents to be returned or 'N'
 *      3) An array of bsons representing the sort specs.
 * For isSingle == false && storeInputExpression == true
 *      4) The input expression to be applied on the sorted results.
 *
 * invertSort is 'false' for ascending and 'true' for descending.
 * isSingle is 'false' for N param accumulators (i.e. firstN/lastN/topN/bottomN)
 *      and 'true' for single accumulators (i.e. first/last/top/bottom).
 */
Datum
BsonOrderTransition(PG_FUNCTION_ARGS, bool invertSort, bool isSingle, bool
					storeInputExpression)
{
	MemoryContext aggregateContext;
	int aggContext = AggCheckCallContext(fcinfo, &aggregateContext);

	if (!aggContext)
	{
		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg(
							"Aggregate function invoked in non-aggregate context")));
	}

	/* We store input expression in $setWindowFields context for $top(N)/$bottom(N) */
	if (aggContext == AGG_CONTEXT_WINDOW)
	{
		storeInputExpression = true;
	}

	BsonOrderAggState inputAggregateState = { 0 };
	bytea *result;

	/* The 2nd argument is the input document*/
	pgbson *inputDocument = PG_GETARG_MAYBE_NULL_PGBSON(1);
	int64 numResults = 1;
	ArrayType *val_array = NULL;

	if (isSingle)
	{
		/* The 3rd argument is a list of sort specs */
		val_array = PG_GETARG_ARRAYTYPE_P(2);
	}
	else
	{
		/* The third parameter specifies the number of results that should be returned */
		numResults = PG_GETARG_INT64(2);
		Assert(numResults > 0);

		/* The 4th argument is a list of sort specs */
		val_array = PG_GETARG_ARRAYTYPE_P(3);
	}

	Datum *sortSpecs;
	bool *nulls;
	int numSortKeys;
	deconstruct_array(val_array,
					  ARR_ELEMTYPE(val_array), -1, false, TYPALIGN_INT,
					  &sortSpecs, &nulls, &numSortKeys);

	Assert(numSortKeys != 0);
	if (numSortKeys > 32)
	{
		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg(
							"Too many compound keys. A maximum of 32 keys is allowed.")));
	}


	bool validateSort = false;

	/* If this is the first call, allocate the state, otherwise deserialize it */
	if (!PG_ARGISNULL(0))
	{
		DeserializeOrderState(PG_GETARG_BYTEA_P(0),
							  &inputAggregateState);
	}
	else
	{
		inputAggregateState.numAggValues = numResults;
		inputAggregateState.currentCount = 0;

		/* Only need one value right now. */
		inputAggregateState.currentResult = palloc0(sizeof(BsonOrderAggValue *));
		inputAggregateState.numSortKeys = numSortKeys;

		for (int i = 0; i < numSortKeys; i++)
		{
			pgbson *sortKeySpec = DatumGetPgBson(sortSpecs[i]);
			pgbsonelement sortKeyElement;
			PgbsonToSinglePgbsonElement(sortKeySpec, &sortKeyElement);
			sortKeyElement.pathLength = 0;
			bool ascOrder = (BsonValueAsInt32(&sortKeyElement.bsonValue) == 1) ?
							true : false;
			inputAggregateState.sortDirections[i] = ascOrder;
		}

		/* Check if this is a single element sort call. */
		if (isSingle)
		{
			/* The 4th argument is the pgbson input expression. */
			inputAggregateState.inputExpression = storeInputExpression ?
												  PG_GETARG_MAYBE_NULL_PGBSON(3) : NULL;
		}
		else
		{
			/* The 5th argument is the pgbson input expression. */
			inputAggregateState.inputExpression = storeInputExpression ?
												  PG_GETARG_MAYBE_NULL_PGBSON(4) : NULL;
		}
	}


	int currentParsePos = 0;
	int currentPos;
	BsonOrderAggValue *newValue = palloc0(sizeof(BsonOrderAggValue));

	/* TODO: support collation with sorted accumulators*/
	char *collationString = NULL;

	/* Find which position newValue needs ot be inserted into */
	for (currentPos = inputAggregateState.currentCount - 1; currentPos >= 0; currentPos--)
	{
		bool replaceCurrentAggregate = false;

		/* Compare the newValue sortValues to currentPos */
		for (int i = 0; i < inputAggregateState.numSortKeys; i++)
		{
			/* Get the possible null bsons to compare */
			pgbson *currentAggregateBson =
				(inputAggregateState.currentResult[currentPos]->sortKeyValues[i] == 0) ?
				NULL :
				DatumGetPgBson(
					inputAggregateState.currentResult[currentPos]->sortKeyValues[i]);

			if (currentParsePos <= i)
			{
				if (inputDocument == NULL)
				{
					newValue->sortKeyValues[currentParsePos++] = 0;
				}
				else
				{
					newValue->sortKeyValues[currentParsePos++] = BsonOrderby(
						inputDocument,
						DatumGetPgBson(sortSpecs[i]),
						validateSort, collationString);
				}
			}

			pgbson *newBson = (newValue->sortKeyValues[i] == 0) ? NULL : DatumGetPgBson(
				newValue->sortKeyValues[i]);

			int comparisonResult = CompareNullablePgbson(currentAggregateBson, newBson);
			bool ascOrder = inputAggregateState.sortDirections[i];

			/* If the call is coming from LAST, we have to invert the comparison to take the correct one */
			if (invertSort)
			{
				comparisonResult *= -1;
			}

			if (comparisonResult > 0)
			{
				if (ascOrder)
				{
					replaceCurrentAggregate = true;
				}
				break;
			}
			else if (comparisonResult < 0)
			{
				if (!ascOrder)
				{
					replaceCurrentAggregate = true;
				}
				break;
			}
		}
		if (!replaceCurrentAggregate)
		{
			break;
		}
		else if (currentPos == inputAggregateState.numAggValues - 1)
		{
			/* Last item in the results can be removed */
			pfree(inputAggregateState.currentResult[currentPos]);
			inputAggregateState.currentResult[currentPos] = NULL;
		}
		else
		{
			/* Shift result current at currentPos to the next position */
			inputAggregateState.currentResult[currentPos + 1] =
				inputAggregateState.currentResult[currentPos];
			inputAggregateState.currentResult[currentPos] = NULL;
		}
	}

	/* Insert new item after the current position */
	currentPos++;

	/* Check to see if the currentPos for the new item is in the result set count.
	 */
	if (currentPos != inputAggregateState.numAggValues)
	{
		newValue->value = inputDocument;

		/* Finish updating any remaining SortKeys that need to be parsed */
		while (currentParsePos < inputAggregateState.numSortKeys)
		{
			if (inputDocument == NULL)
			{
				newValue->sortKeyValues[currentParsePos] = 0;
			}
			else
			{
				newValue->sortKeyValues[currentParsePos] =
					BsonOrderby(inputDocument,
								DatumGetPgBson(
									sortSpecs
									[currentParsePos]),
								validateSort,
								collationString);
			}
			currentParsePos++;
		}
		inputAggregateState.currentResult[currentPos] = newValue;
		if (inputAggregateState.numAggValues > inputAggregateState.currentCount)
		{
			inputAggregateState.currentCount += 1;
		}
		result = SerializeOrderState(&inputAggregateState);
	}
	else
	{
		/* Free up newly allocated object and return old pointer. */
		result = PG_ARGISNULL(0) ? NULL : PG_GETARG_BYTEA_P(0);
	}

	return PointerGetDatum(result);
}


/*
 * Applies the "state transition" (SFUNC) for accumulator.
 * The args are:
 *      0) The current aggregation state
 *      1) The document/evaluated expression which will eventually be returned
 *
 * For isSingle == 'true' && storeInputExpression == 'true
 *      2) The input expression to be applied on the results.
 *
 * For isSingle == 'false' && storeInputExpression == 'true
 *      2) The number of documents to be returned or 'N'
 *      3) The input expression to be applied on the results.
 */
Datum
BsonOrderTransitionOnSorted(PG_FUNCTION_ARGS, bool invertSort, bool isSingle)
{
	MemoryContext aggregateContext;
	int aggContext = AggCheckCallContext(fcinfo, &aggregateContext);

	if (!aggContext)
	{
		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg(
							"Aggregate function invoked in non-aggregate context")));
	}

	bool storeInputExpression = false;

	/* We store input expression in $setWindowFields context */
	if (aggContext == AGG_CONTEXT_WINDOW)
	{
		storeInputExpression = true;
	}

	bytea *byteArray = PG_ARGISNULL(0) ? NULL : PG_GETARG_BYTEA_P(0);
	pgbson *newValue = PG_GETARG_MAYBE_NULL_PGBSON(1);
	int64 returnCount = 1;
	char *sourcePtr;
	int64 currentCount = 0;
	int64 copySize = 0;
	pgbson *inputExpression;
	int64 inputExpressionSize = 0;

	/*
	 *  If no sort keys are provided, or the data is pre-sorted we do the following:
	 *      (1) Initialize the current value with the first N elements. And after that
	 *          (1.a.) keep ignoring any new element for invertSort = false
	 *          (1.b.) keep replacing the lease recent value with the new one for invertSort = true
	 *  Format is:
	 *  VARHDR | ByteSize | ReturnCount | CurrentCount | InputExpressionSize | InputExpression | Result 1 | Size 1 | Result 2 | Size 2 | ...
	 */

	/* Get the current count and setup*/
	if (byteArray == NULL)
	{
		copySize = 0;
		sourcePtr = NULL;
		currentCount = 1;
		if (!isSingle)
		{
			/* The third parameter specifies the number of results that should be returned */
			returnCount = PG_GETARG_INT64(2);
			inputExpression = storeInputExpression ? PG_GETARG_MAYBE_NULL_PGBSON(3) :
							  NULL;
		}
		else
		{
			inputExpression = storeInputExpression ? PG_GETARG_MAYBE_NULL_PGBSON(2) :
							  NULL;
		}

		if (inputExpression != NULL)
		{
			inputExpressionSize = VARSIZE(inputExpression);
		}
	}
	else
	{
		sourcePtr = (char *) VARDATA(byteArray);

		/* Copy size is total size - VARHDRSZ, the 4 ints we are extracting, and ptr. */
		int64 totalSizeValue;
		memcpy(&totalSizeValue, sourcePtr, sizeof(int64));
		copySize = totalSizeValue - sizeof(int64) * 4 - VARHDRSZ;
		sourcePtr += sizeof(int64);
		memcpy(&returnCount, sourcePtr, sizeof(int64));
		sourcePtr += sizeof(int64);
		memcpy(&currentCount, sourcePtr, sizeof(int64));
		sourcePtr += sizeof(int64);
		memcpy(&inputExpressionSize, sourcePtr, sizeof(int64));
		sourcePtr += sizeof(int64);

		if (inputExpressionSize != 0)
		{
			inputExpression = (pgbson *) sourcePtr;
			sourcePtr += inputExpressionSize;
			copySize -= inputExpressionSize;
		}
		else
		{
			inputExpression = NULL;
		}

		if (currentCount > returnCount)
		{
			/* Should never encounter this state unless somethin is very wrong */
			ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg(
								"Invalid state for aggregate function BsonOrderTransitionOnSorted")));
		}
		else if (currentCount == returnCount && !invertSort)
		{
			/* First N is already at N elements so just return */
			PG_RETURN_POINTER(byteArray);
		}
		else if (currentCount == returnCount && invertSort)
		{
			/*
			 * When the aggregation state has already accumulated `returnCount` documents
			 * and `invertSort == true` (e.g., for a $lastN without sort stage and bottomN behavior), we need to evict
			 * the oldest document (the first one) to make room for the incoming new document.
			 *
			 * This block of code performs that eviction:
			 * - It first verifies that the remaining data contains at least one complete BSON document
			 *   by checking its minimum size and validity.
			 * - Then it reads the size of the first BSON document.
			 * - Finally, it advances the source pointer past the first document (including its trailing
			 *   size marker), effectively removing it from the aggregation state.
			 */

			char *entryPtr = sourcePtr;
			if (copySize < (int64) VARHDRSZ + (int64) sizeof(uint32))
			{
				ereport(ERROR,
						(errcode(ERRCODE_INTERNAL_ERROR),
						 errmsg("Insufficient data for first aggregate entry")));
			}

			int32 bsonSize = VARSIZE_ANY((pgbson *) entryPtr);
			if (bsonSize + (int64) sizeof(uint32) > copySize)
			{
				ereport(ERROR,
						(errcode(ERRCODE_INTERNAL_ERROR),
						 errmsg("Corrupted BSON entry in aggregation state")));
			}

			entryPtr += bsonSize + sizeof(uint32);
			copySize -= (entryPtr - sourcePtr);
			sourcePtr = entryPtr;
		}
		else
		{
			/* Just append the new element to the current results */
			currentCount++;
		}
	}

	/* Output is VARHDR | ByteSize | ReturnCount | CurrentCount | InputExpressionSize | InputExpression | values... */
	uint32 totalSize = VARHDRSZ + sizeof(int64) * 4 + copySize + inputExpressionSize;

	/* Add in newValue to total size */
	if (newValue != NULL)
	{
		totalSize += VARSIZE(newValue) + sizeof(uint32);
	}
	else
	{
		/* Write NULL as byte with value 0 */
		totalSize += 1 + sizeof(uint32);
	}

	bytea *returnData = byteArray;
	if (byteArray == NULL || totalSize > VARSIZE(byteArray))
	{
		/* Don't need to pfree byteArray as it get cleaned up by postgres */
		returnData = (bytea *) palloc0(totalSize);
		SET_VARSIZE(returnData, totalSize);
	}

	char *returnDataPtr = (char *) VARDATA(returnData);
	int64 totalSizeAsInt64 = (int64) totalSize;
	memcpy(returnDataPtr, &totalSizeAsInt64, sizeof(int64));
	returnDataPtr += sizeof(int64);
	memcpy(returnDataPtr, &returnCount, sizeof(int64));
	returnDataPtr += sizeof(int64);
	memcpy(returnDataPtr, &currentCount, sizeof(int64));
	returnDataPtr += sizeof(int64);
	memcpy(returnDataPtr, &inputExpressionSize, sizeof(int64));
	returnDataPtr += sizeof(int64);

	/* Copy input expression to returnDataPtr */
	if (inputExpression != NULL)
	{
		memcpy(returnDataPtr, inputExpression, inputExpressionSize);
		returnDataPtr += inputExpressionSize;
	}


	/* Need to copy first so we don't overwrite with newValue if we are reusing the array. */
	/* Use memmove to account for the fact that we may be overwriting the data we are copying. */
	/* Values for previously existing return values are no longer reliable after this. */
	/* Copy over the first values we've already seen and adjust pointer to next position. */
	if (copySize != 0)
	{
		memmove(returnDataPtr, sourcePtr, copySize);
		returnDataPtr += copySize;
	}

	/* Add in new value */
	if (newValue != NULL)
	{
		memcpy(returnDataPtr, newValue, VARSIZE(newValue));
		returnDataPtr += VARSIZE(newValue);
		uint32 newValueSize = VARSIZE(newValue);
		memcpy(returnDataPtr, &newValueSize, sizeof(uint32));
	}
	else
	{
		*returnDataPtr = 0;
		returnDataPtr += 1;
		uint32 nullMarkerSize = (uint32) 1;
		memcpy(returnDataPtr, &nullMarkerSize, sizeof(uint32));
	}

	PG_RETURN_POINTER(returnData);
}


/*
 * Applies the "state transition" (SFUNC) for accumulator.
 *
 * The args of PG_FUNCTION_ARGS are:
 *      0) left side partial aggregate
 *      1) right side partial aggregate
 */
Datum
BsonOrderCombine(PG_FUNCTION_ARGS, bool invertSort)
{
	MemoryContext aggregateContext;
	if (!AggCheckCallContext(fcinfo, &aggregateContext))
	{
		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg(
							"Aggregate function invoked in non-aggregate context")));
	}

	if (PG_ARGISNULL(0) && PG_ARGISNULL(1))
	{
		PG_RETURN_NULL();
	}

	/* Check if either result is NULL and return Non-NULL */
	if (PG_ARGISNULL(0))
	{
		return PG_GETARG_DATUM(1);
	}

	if (PG_ARGISNULL(1))
	{
		return PG_GETARG_DATUM(0);
	}

	BsonOrderAggState leftState = { 0 }, rightState = { 0 };
	DeserializeOrderState(PG_GETARG_BYTEA_P(0), &leftState);
	DeserializeOrderState(PG_GETARG_BYTEA_P(1), &rightState);

	int minPos = 0; /* Minimum position that new result can be inserted at */
	int leftPos, rightPos; /* current result position for left/right states*/
	bool updatedLeft = false; /* Flag indicates whether the leftState was updated */

	/* Allocate additional memory for merging right results. */
	if (leftState.currentCount != leftState.numAggValues)
	{
		int newSize = leftState.numAggValues;
		if (leftState.currentCount + rightState.currentCount < leftState.numAggValues)
		{
			newSize = leftState.currentCount + rightState.currentCount;
		}
		BsonOrderAggValue **newPtr = palloc0(sizeof(BsonOrderAggValue *) * newSize);

		/* Copy over pointers to values to new array */
		memcpy(newPtr, leftState.currentResult, sizeof(BsonOrderAggValue *) *
			   leftState.currentCount);

		/* Free current pointer and update */
		pfree(leftState.currentResult);
		leftState.currentResult = newPtr;
	}

	/* Loop through each object in rightState and insert it into leftState*/
	for (rightPos = 0; rightPos < rightState.currentCount; rightPos++)
	{
		if (rightState.currentResult[rightPos] == NULL)
		{
			/* No more results in right */
			break;
		}
		for (leftPos = leftState.currentCount - 1; leftPos >= minPos; leftPos--)
		{
			bool replaceCurrentAggregate = false;
			if (leftState.currentResult[leftPos] == NULL)
			{
				/* Current result isn't set so skip it. */
				continue;
			}

			for (int i = 0; i < leftState.numSortKeys; i++)
			{
				/* Get the possible null bsons to compare */
				pgbson *leftBson = (leftState.currentResult[leftPos]->sortKeyValues[i] ==
									0) ? NULL :
								   DatumGetPgBson(
					leftState.currentResult[leftPos]->sortKeyValues[i]);
				pgbson *rightBson =
					(rightState.currentResult[rightPos]->sortKeyValues[i] == 0) ? NULL :
					DatumGetPgBson(
						rightState.currentResult[rightPos]->sortKeyValues[i]);

				int comparisonResult = CompareNullablePgbson(leftBson, rightBson);
				bool ascOrder = leftState.sortDirections[i];

				/* Check if we have to invert the comparison to take the correct one */
				if (invertSort)
				{
					comparisonResult *= -1;
				}

				if (comparisonResult > 0)
				{
					if (ascOrder)
					{
						replaceCurrentAggregate = true;
					}
					break;
				}
				else if (comparisonResult < 0)
				{
					if (!ascOrder)
					{
						replaceCurrentAggregate = true;
					}
					break;
				}
			}

			if (!replaceCurrentAggregate)
			{
				break;
			}
			else if (leftPos == leftState.numAggValues - 1)
			{
				/* Last Item in leftState results can be removed as it will be replaced */
				pfree(leftState.currentResult[leftPos]);
				leftState.currentResult[leftPos] = NULL;
			}
			else
			{
				/* Shift result currently at leftPos to the next position */
				leftState.currentResult[leftPos + 1] =
					leftState.currentResult[leftPos];
				leftState.currentResult[leftPos] = NULL;
			}
		}

		/* Insert current result after this position */
		leftPos++;

		if (leftPos != leftState.numAggValues)
		{
			/* If currentCount isn't already numAggValues we are adding
			 * an additional value.  Otherwise we are replacing one.
			 */
			if (leftState.currentCount != leftState.numAggValues)
			{
				leftState.currentCount++;
			}

			/* Replace the object currently pointed at with the right one*/
			leftState.currentResult[leftPos] = rightState.currentResult[rightPos];

			/* Clear rightState pointer to prevent double free */
			rightState.currentResult[rightPos] = NULL;

			/* Set flag to indicate that we need to reserialize leftState */
			updatedLeft = true;
		}

		/* All remaining rightState elements will have to be inserted after current leftPos */
		minPos = leftPos + 1;
	}


	bytea *result = NULL;

	if (updatedLeft)
	{
		result = SerializeOrderState(&leftState);
	}
	else
	{
		result = PG_GETARG_BYTEA_P(0);
	}

	return PointerGetDatum(result);
}


/*
 * Applies the FINALFUNC for accumulator.
 *
 * The args of PG_FUNCTION_ARGS are:
 *      0) Current state
 * if IsSingle is true a single value will be returned
 * otherwise an array of values will be returned.
 */
Datum
BsonOrderFinal(PG_FUNCTION_ARGS, bool isSingle, bool invert)
{
	MemoryContext aggregateContext;
	int aggContext = AggCheckCallContext(fcinfo, &aggregateContext);
	if (!aggContext)
	{
		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg(
							"Aggregate function invoked in non-aggregate context")));
	}
	bool returnNull = false;
	BsonOrderAggState state = { 0 };

	if (PG_ARGISNULL(0))
	{
		returnNull = true;
	}
	else
	{
		DeserializeOrderState(PG_GETARG_BYTEA_P(0), &state);
		if (isSingle)
		{
			/* Validate there is a value to return. */
			bson_iter_t pathSpecIter;
			PgbsonInitIterator(state.currentResult[0]->value, &pathSpecIter);
			if (!bson_iter_next(&pathSpecIter))
			{
				returnNull = true;
			}
		}
	}

	/* If there were no regular rows, or the last result was NULL, the result is $null */
	if (returnNull)
	{
		pgbson *finalPgbson;

		if (isSingle)
		{
			/* Returns $null for empty sets */
			pgbsonelement finalValue;
			finalValue.path = "";
			finalValue.pathLength = 0;
			finalValue.bsonValue.value_type = BSON_TYPE_NULL;
			finalPgbson = PgbsonElementToPgbson(&finalValue);
		}
		else
		{
			/* empty array */
			pgbson_writer writer;
			PgbsonWriterInit(&writer);
			PgbsonWriterAppendEmptyArray(&writer, "", 0);
			finalPgbson = PgbsonWriterGetPgbson(&writer);
		}

		PG_RETURN_POINTER(finalPgbson);
	}

	AggregationExpressionData *aggregationExpressionState = NULL;
	StringView path;
	ExpressionVariableContext *variableContext = NULL;

	/* Cache aggregationExpressionState for $setWindowFields to avoid re-parsing the expression for each row/partition.
	 * aggregationExpressionState holds the parsed input expression which stays the same for all rows/partitions.
	 */
	if (aggContext == AGG_CONTEXT_WINDOW)
	{
		if (state.inputExpression != NULL)
		{
			pgbsonelement element;
			PgbsonToSinglePgbsonElement(state.inputExpression, &element);

			path = (StringView) {
				.length = element.pathLength,
				.string = element.path,
			};

			/*
			 * Use numArgs as 0 as we want to return true from IsSafeToReuseFmgrFunctionExtraMultiArgs().
			 * This is to avoid the check for the node being a const or a Param of type extern as the caller knows its safe to reuse
			 * as this is the common expression for all rows inside the bytea returned by transition function.
			 */
			int argPositions[1] = { 0 };
			int numArgs = 0;
			ParseAggregationExpressionContext context = { 0 };

			SetCachedFunctionStateMultiArgs(
				aggregationExpressionState,
				AggregationExpressionData,
				argPositions,
				numArgs,
				ParseInputExpressionAndPersistValue,
				&element.bsonValue,
				&context);

			if (aggregationExpressionState == NULL)
			{
				aggregationExpressionState = palloc0(sizeof(AggregationExpressionData));
				ParseInputExpressionAndPersistValue(
					aggregationExpressionState,
					&element.bsonValue,
					&context);
			}
		}
	}

	pgbson *result = NULL;
	if (isSingle)
	{
		result = state.currentResult[0]->value;
		if (aggContext == AGG_CONTEXT_WINDOW)
		{
			if (state.inputExpression != NULL)
			{
				/* Apply the inputExpression to the result documents to calculate result for $top/$bottom */
				bool isNullOnEmpty = true;
				pgbson_writer writer;
				PgbsonWriterInit(&writer);
				EvaluateAggregationExpressionDataToWriter(aggregationExpressionState,
														  state.currentResult[0]->value,
														  path,
														  &writer,
														  variableContext, isNullOnEmpty);

				result = PgbsonWriterGetPgbson(&writer);
			}
		}
	}
	else
	{
		/* Create array of results */
		pgbson_writer writer;
		pgbson_array_writer arrayWriter;
		PgbsonWriterInit(&writer);
		PgbsonWriterStartArray(&writer, "", 0, &arrayWriter);

		int currentPoint = invert ? state.currentCount - 1 : 0;
		int toPoint = invert ? -1 : state.currentCount;
		int direction = invert ? -1 : 1;

		for (int i = currentPoint; i != toPoint; i += direction)
		{
			if (state.currentResult[i] == NULL)
			{
				/* No additional results found*/
				break;
			}

			/* Check for Null value*/
			if (state.currentResult[i]->value != NULL)
			{
				if (aggContext == AGG_CONTEXT_WINDOW)
				{
					if (state.inputExpression != NULL)
					{
						/* Apply the inputExpression to the result documents to calculate result for $topN/$bottomN */
						bool isNullOnEmpty = true;
						pgbson_writer innerWriter;
						PgbsonWriterInit(&innerWriter);
						EvaluateAggregationExpressionDataToWriter(
							aggregationExpressionState,
							state.currentResult[i]->value,
							path,
							&innerWriter,
							variableContext,
							isNullOnEmpty);
						pgbson *resultpg = PgbsonWriterGetPgbson(&innerWriter);
						pgbsonelement element;
						PgbsonToSinglePgbsonElement(resultpg, &element);
						PgbsonArrayWriterWriteValue(&arrayWriter, &element.bsonValue);
					}
				}
				else
				{
					PgbsonArrayWriterWriteDocument(&arrayWriter,
												   state.currentResult[i]->value);
				}
			}
			else
			{
				PgbsonArrayWriterWriteNull(&arrayWriter);
			}
		}
		PgbsonWriterEndArray(&writer, &arrayWriter);
		result = PgbsonWriterGetPgbson(&writer);
	}

	PG_RETURN_POINTER(result);
}


/*
 * Applies the FINALFUNC for already sorted accumulator.
 *
 * The args of PG_FUNCTION_ARGS are:
 *      0) Current state
 * if IsSingle is true a single value will be returned
 * otherwise an array of values will be returned.
 */
Datum
BsonOrderFinalOnSorted(PG_FUNCTION_ARGS, bool isSingle)
{
	MemoryContext aggregateContext;
	int aggContext = AggCheckCallContext(fcinfo, &aggregateContext);

	if (!aggContext)
	{
		ereport(ERROR, (errcode(ERRCODE_INTERNAL_ERROR), errmsg(
							"Aggregate function invoked in non-aggregate context")));
	}

	bool returnNull = false;
	bytea *byteArray = PG_ARGISNULL(0) ? NULL : PG_GETARG_BYTEA_P(0);
	char *sourcePtr;
	int64 currentCount = 0;
	pgbson *inputExpression = NULL;
	int64 inputExpressionSize = 0;

	if (byteArray == NULL)
	{
		returnNull = true;
	}
	else
	{
		sourcePtr = (char *) VARDATA(byteArray);

		/* Skip over bytes and returnCount */
		sourcePtr += sizeof(int64);
		sourcePtr += sizeof(int64);
		memcpy(&currentCount, sourcePtr, sizeof(int64));
		sourcePtr += sizeof(int64);
		memcpy(&inputExpressionSize, sourcePtr, sizeof(int64));
		sourcePtr += sizeof(int64);

		/* Extract inputExpression from sourcePtr */
		if (inputExpressionSize != 0)
		{
			inputExpression = (pgbson *) sourcePtr;
			sourcePtr += inputExpressionSize;
		}

		if (isSingle)
		{
			/* validate the return value */
			if (currentCount == 0)
			{
				returnNull = true;
			}
			else
			{
				pgbson *currentValue = (pgbson *) sourcePtr;
				bson_iter_t pathSpecIter;
				PgbsonInitIterator(currentValue, &pathSpecIter);
				if (!bson_iter_next(&pathSpecIter))
				{
					returnNull = true;
				}
			}
		}
	}

	AggregationExpressionData *aggregationExpressionState = NULL;
	StringView path;
	ExpressionVariableContext *variableContext = NULL;

	/* Cache aggregationExpressionState for $setWindowFields to avoid re-parsing the expression for each row/partition.
	 * aggregationExpressionState holds the parsed input expression which stays the same for all rows/partitions.
	 */
	if (aggContext == AGG_CONTEXT_WINDOW)
	{
		if (inputExpression != NULL)
		{
			pgbsonelement element;
			PgbsonToSinglePgbsonElement(inputExpression, &element);

			path = (StringView) {
				.length = element.pathLength,
				.string = element.path,
			};

			/*
			 * Use numArgs as 0 as we want to return true from IsSafeToReuseFmgrFunctionExtraMultiArgs().
			 * This is to avoid the check for the node being a const or a Param of type extern as the caller knows its safe to reuse
			 * as this is the common expression for all rows inside the bytea returned by transition function.
			 */
			int argPositions[1] = { 0 };
			int numArgs = 0;
			ParseAggregationExpressionContext context = { 0 };

			SetCachedFunctionStateMultiArgs(
				aggregationExpressionState,
				AggregationExpressionData,
				argPositions,
				numArgs,
				ParseInputExpressionAndPersistValue,
				&element.bsonValue,
				&context);

			if (aggregationExpressionState == NULL)
			{
				aggregationExpressionState = palloc0(sizeof(AggregationExpressionData));
				ParseInputExpressionAndPersistValue(
					aggregationExpressionState,
					&element.bsonValue,
					&context);
			}
		}
	}

	/* If there were no regular rows, or the last result was NULL, the result is $null */
	if (returnNull)
	{
		pgbson *finalPgbson;

		if (isSingle)
		{
			/* Returns $null for empty sets */
			pgbsonelement finalValue;
			finalValue.path = "";
			finalValue.pathLength = 0;
			finalValue.bsonValue.value_type = BSON_TYPE_NULL;
			finalPgbson = PgbsonElementToPgbson(&finalValue);
		}
		else
		{
			/* empty array */
			pgbson_writer writer;
			PgbsonWriterInit(&writer);
			PgbsonWriterAppendEmptyArray(&writer, "", 0);
			finalPgbson = PgbsonWriterGetPgbson(&writer);
		}

		PG_RETURN_POINTER(finalPgbson);
	}
	else if (isSingle)
	{
		if (aggContext == AGG_CONTEXT_WINDOW)
		{
			if (inputExpression != NULL)
			{
				/* Apply the inputExpression to the result documents to calculate result for $first/$last */
				bool isNullOnEmpty = true;
				pgbson_writer writer;
				PgbsonWriterInit(&writer);
				EvaluateAggregationExpressionDataToWriter(aggregationExpressionState,
														  (pgbson *) sourcePtr,
														  path,
														  &writer,
														  variableContext, isNullOnEmpty);

				pgbson *result = PgbsonWriterGetPgbson(&writer);
				PG_RETURN_POINTER(result);
			}
		}

		PG_RETURN_POINTER((pgbson *) sourcePtr);
	}
	else
	{
		/* Create array of results */
		pgbson_writer writer;
		pgbson_array_writer arrayWriter;
		PgbsonWriterInit(&writer);
		PgbsonWriterStartArray(&writer, "", 0, &arrayWriter);
		for (uint32 i = 0; i < currentCount; i++)
		{
			if (*sourcePtr == 0)
			{
				PgbsonArrayWriterWriteNull(&arrayWriter);

				/* Move to next value */
				sourcePtr += sizeof(uint32) + 1;
			}
			else
			{
				if (aggContext == AGG_CONTEXT_WINDOW)
				{
					if (inputExpression != NULL)
					{
						/* Apply the inputExpression to the result documents to calculate result for $firstN/$lastN */
						bool isNullOnEmpty = true;
						pgbson_writer innerWriter;
						PgbsonWriterInit(&innerWriter);
						EvaluateAggregationExpressionDataToWriter(
							aggregationExpressionState,
							(pgbson *) sourcePtr,
							path,
							&innerWriter,
							variableContext,
							isNullOnEmpty);
						pgbson *resultpg = PgbsonWriterGetPgbson(&innerWriter);
						pgbsonelement element;
						PgbsonToSinglePgbsonElement(resultpg, &element);
						PgbsonArrayWriterWriteValue(&arrayWriter, &element.bsonValue);
					}
				}
				else
				{
					PgbsonArrayWriterWriteDocument(&arrayWriter, (pgbson *) sourcePtr);
				}

				sourcePtr += VARSIZE(sourcePtr);
				sourcePtr += sizeof(uint32);
			}
		}
		PgbsonWriterEndArray(&writer, &arrayWriter);
		PG_RETURN_POINTER(PgbsonWriterGetPgbson(&writer));
	}
}


/*
 * Create a copy of the inputExpression bson_value_t to make sure it can be persisted in the cache across partitions in $setWindowFields.
 * BsonValueToDocumentPgbson creates a new pgbson out of the value in the current memory context which is ExecutorContext
 * when cache is being created so the value is guaranteed to live across different partitions and rows.
 */
static void
ParseInputExpressionAndPersistValue(AggregationExpressionData *expressionData,
									const bson_value_t *expressionValue,
									ParseAggregationExpressionContext *context)
{
	pgbson *pgbson = BsonValueToDocumentPgbson(expressionValue);
	pgbsonelement element;
	PgbsonToSinglePgbsonElement(pgbson, &element);
	bson_value_t persistedValue = element.bsonValue;
	ParseAggregationExpressionData(expressionData, &persistedValue, context);
}
