#include "postgres.h"
#include "fmgr.h"
#include "utils/array.h"
#include "utils/memutils.h"
#include "nodes/execnodes.h"

PG_MODULE_MAGIC;

/* Structure for storing our pointers to the
 * ArrayBuildState for the array we are building
 * and the MemoryContext in which it is being
 * built.  Note that this structure is 
 * considered a bytea externally and therefore
 * must open with an int32 defining the length. */
typedef struct {
	int32				 vl_len;
	ArrayBuildState		*astate;
	MemoryContext		 arrctx;
} aaccum_info;

/* The state-transistion function for our aggregate. */
PG_FUNCTION_INFO_V1(aaccum_sfunc);
Datum
aaccum_sfunc(PG_FUNCTION_ARGS)
{
	aaccum_info		*ainfo;
	AggState		*aggstate;

	/* Make sure we are in an aggregate. */
	if (!fcinfo->context || !IsA(fcinfo->context, AggState))
		ereport(ERROR,
				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
				 errmsg("Can not call aaccum_sfunc as a non-aggregate")));

	aggstate = (AggState*) fcinfo->context;

	/* Initial call passes NULL in for our state variable. 
	 * Allocate memory and get set up. */
	if (PG_ARGISNULL(0)) {
		/* Allocate memory to hold the pointers to the ArrayBuildState
		 * and the MemoryContext where we are building the array.  Note
		 * that we can do this in the CurrentMemoryContext because when
		 * we return the storage "bytea" will be copied into the AggState
		 * context by the caller and passed back to us on the next call. */
		ainfo = (aaccum_info*) palloc(sizeof(aaccum_info));
		ainfo->vl_len = sizeof(aaccum_info);
		ainfo->astate = NULL;

		/* New context created which will store our array accumulation.
		 * The parent is the AggContext for this query since it needs to
		 * persist for the same timeframe as the state value. 
		 * The state value holds the pointers to the ArrayBuildState and this 
		 * MemoryContext through the aaccum_info structure. */
		ainfo->arrctx = AllocSetContextCreate(aggstate->aggcontext, "ArrayAccumCtx",
											  ALLOCSET_DEFAULT_MINSIZE,
											  ALLOCSET_DEFAULT_INITSIZE,
											  ALLOCSET_DEFAULT_MAXSIZE);
	} else {
		/* Our state variable is non-null, therefore it must be an existing
		 * ainfo structure. */
		ainfo = (aaccum_info*) PG_GETARG_BYTEA_P(0);
	}

	/* Pull the element to be added and pass it along with the ArrayBuildState
	 * and ArrayAccumCtx MemoryContext to accumArrayResult, checking if it is
	 * NULL or not. */
	ainfo->astate = accumArrayResult(ainfo->astate, 
									 PG_ARGISNULL(1) ? (Datum) 0 : PG_GETARG_DATUM(1),
									 PG_ARGISNULL(1),
									 get_fn_expr_argtype(fcinfo->flinfo, 1),
									 ainfo->arrctx);

	/* Caller will copy storage into the AggContext after the first call and then
	 * should not touch it as we will always return the same pointer passed in. */
	PG_RETURN_BYTEA_P(ainfo);
}

/* The final function for our aggregate. */
PG_FUNCTION_INFO_V1(aaccum_ffunc);
Datum
aaccum_ffunc(PG_FUNCTION_ARGS)
{
	aaccum_info		*ainfo;

	/* Check if we are passed in a NULL */
	if (PG_ARGISNULL(0)) PG_RETURN_ARRAYTYPE_P(NULL);

	/* Make sure we are in an aggregate. */
	if (!fcinfo->context || !IsA(fcinfo->context, AggState))
		ereport(ERROR,
				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
				 errmsg("Can not call aaccum_sfunc as a non-aggregate")));

	ainfo = (aaccum_info*) PG_GETARG_BYTEA_P(0);

	/* makeArrayResult will delete ainfo->arrctx for us. */
	PG_RETURN_ARRAYTYPE_P(makeArrayResult(ainfo->astate, ainfo->arrctx));
}