/*------------------------------------------------------------------------------------
 *
 * test_custom_var_stats.c
 *		Test module for variable-sized custom pgstats
 *
 * Copyright (c) 2025-2026, PostgreSQL Global Development Group
 *
 * IDENTIFICATION
 *		src/test/modules/test_custom_var_stats/test_custom_var_stats.c
 *
 * ------------------------------------------------------------------------------------
 */
#include "postgres.h"

#include "access/htup_details.h"
#include "common/hashfn.h"
#include "funcapi.h"
#include "storage/dsm_registry.h"
#include "storage/fd.h"
#include "utils/builtins.h"
#include "utils/pgstat_internal.h"

PG_MODULE_MAGIC_EXT(
					.name = "test_custom_var_stats",
					.version = PG_VERSION
);

#define TEST_CUSTOM_VAR_MAGIC_NUMBER (0xBEEFBEEF)

/*--------------------------------------------------------------------------
 * Macros and constants
 *--------------------------------------------------------------------------
 */

/*
 * Kind ID for test_custom_var_stats statistics.
 */
#define PGSTAT_KIND_TEST_CUSTOM_VAR_STATS 25

/* File paths for auxiliary data serialization */
#define TEST_CUSTOM_AUX_DATA_DESC "pg_stat/test_custom_var_stats_desc.stats"

/*
 * Hash statistic name to generate entry index for pgstat lookup.
 */
#define PGSTAT_CUSTOM_VAR_STATS_IDX(name) hash_bytes_extended((const unsigned char *) name, strlen(name), 0)

/*--------------------------------------------------------------------------
 * Type definitions
 *--------------------------------------------------------------------------
 */

/* Backend-local pending statistics before flush to shared memory */
typedef struct PgStat_StatCustomVarEntry
{
	PgStat_Counter numcalls;	/* times statistic was incremented */
} PgStat_StatCustomVarEntry;

/* Shared memory statistics entry visible to all backends */
typedef struct PgStatShared_CustomVarEntry
{
	PgStatShared_Common header; /* standard pgstat entry header */
	PgStat_StatCustomVarEntry stats;	/* custom statistics data */
	dsa_pointer description;	/* pointer to description string in DSA */
} PgStatShared_CustomVarEntry;

/*--------------------------------------------------------------------------
 * Global Variables
 *--------------------------------------------------------------------------
 */

/* File handle for auxiliary data serialization */
static FILE *fd_description = NULL;

/* Current write offset in fd_description file */
static pgoff_t fd_description_offset = 0;

/* DSA area for storing variable-length description strings */
static dsa_area *custom_stats_description_dsa = NULL;

/*--------------------------------------------------------------------------
 * Function prototypes
 *--------------------------------------------------------------------------
 */

/* Flush callback: merge pending stats into shared memory */
static bool test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref,
												   bool nowait);

/* Serialization callback: write auxiliary entry data */
static void test_custom_stats_var_to_serialized_data(const PgStat_HashKey *key,
													 const PgStatShared_Common *header,
													 FILE *statfile);

/* Deserialization callback: read auxiliary entry data */
static bool test_custom_stats_var_from_serialized_data(const PgStat_HashKey *key,
													   PgStatShared_Common *header,
													   FILE *statfile);

/* Finish callback: end of statistics file operations */
static void test_custom_stats_var_finish(PgStat_StatsFileOp status);

/*--------------------------------------------------------------------------
 * Custom kind configuration
 *--------------------------------------------------------------------------
 */

static const PgStat_KindInfo custom_stats = {
	.name = "test_custom_var_stats",
	.fixed_amount = false,		/* variable number of entries */
	.write_to_file = true,		/* persist across restarts */
	.track_entry_count = true,	/* count active entries */
	.accessed_across_databases = true,	/* global statistics */
	.shared_size = sizeof(PgStatShared_CustomVarEntry),
	.shared_data_off = offsetof(PgStatShared_CustomVarEntry, stats),
	.shared_data_len = sizeof(((PgStatShared_CustomVarEntry *) 0)->stats),
	.pending_size = sizeof(PgStat_StatCustomVarEntry),
	.flush_pending_cb = test_custom_stats_var_flush_pending_cb,
	.to_serialized_data = test_custom_stats_var_to_serialized_data,
	.from_serialized_data = test_custom_stats_var_from_serialized_data,
	.finish = test_custom_stats_var_finish,
};

/*--------------------------------------------------------------------------
 * Module initialization
 *--------------------------------------------------------------------------
 */

void
_PG_init(void)
{
	/* Must be loaded via shared_preload_libraries */
	if (!process_shared_preload_libraries_in_progress)
		return;

	/* Register custom statistics kind */
	pgstat_register_kind(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, &custom_stats);
}

/*--------------------------------------------------------------------------
 * Statistics callback functions
 *--------------------------------------------------------------------------
 */

/*
 * test_custom_stats_var_flush_pending_cb
 *		Merge pending backend statistics into shared memory
 *
 * Called by pgstat collector to flush accumulated local statistics
 * to shared memory where other backends can read them.
 *
 * Returns false only if nowait=true and lock acquisition fails.
 */
static bool
test_custom_stats_var_flush_pending_cb(PgStat_EntryRef *entry_ref, bool nowait)
{
	PgStat_StatCustomVarEntry *pending_entry;
	PgStatShared_CustomVarEntry *shared_entry;

	pending_entry = (PgStat_StatCustomVarEntry *) entry_ref->pending;
	shared_entry = (PgStatShared_CustomVarEntry *) entry_ref->shared_stats;

	if (!pgstat_lock_entry(entry_ref, nowait))
		return false;

	/* Add pending counts to shared totals */
	shared_entry->stats.numcalls += pending_entry->numcalls;

	pgstat_unlock_entry(entry_ref);

	return true;
}

/*
 * test_custom_stats_var_to_serialized_data() -
 *
 * Serialize auxiliary data (descriptions) for custom statistics entries
 * to a secondary statistics file. This is called while writing the statistics
 * to disk.
 *
 * This callback writes a mix of data within the main pgstats file and a
 * secondary statistics file.  The following data is written to the main file for
 * each entry:
 * - An arbitrary magic number.
 * - An offset.  This is used to know the location we need to look at
 * to retrieve the information from the second file.
 *
 * The following data is written to the secondary statistics file:
 * - The entry key, cross-checked with the data from the main file
 * when reloaded.
 * - The length of the description.
 * - The description data itself.
 */
static void
test_custom_stats_var_to_serialized_data(const PgStat_HashKey *key,
										 const PgStatShared_Common *header,
										 FILE *statfile)
{
	char	   *description;
	size_t		len;
	const PgStatShared_CustomVarEntry *entry = (const PgStatShared_CustomVarEntry *) header;
	bool		found;
	uint32		magic_number = TEST_CUSTOM_VAR_MAGIC_NUMBER;

	/*
	 * First mark the main file with a magic number, keeping a trace that some
	 * auxiliary data will exist in the secondary statistics file.
	 */
	pgstat_write_chunk_s(statfile, &magic_number);

	/* Open statistics file for writing. */
	if (!fd_description)
	{
		fd_description = AllocateFile(TEST_CUSTOM_AUX_DATA_DESC, PG_BINARY_W);
		if (fd_description == NULL)
		{
			ereport(LOG,
					(errcode_for_file_access(),
					 errmsg("could not open statistics file \"%s\" for writing: %m",
							TEST_CUSTOM_AUX_DATA_DESC)));
			return;
		}

		/* Initialize offset for secondary statistics file. */
		fd_description_offset = 0;
	}

	/* Write offset to the main data file */
	pgstat_write_chunk_s(statfile, &fd_description_offset);

	/*
	 * First write the entry key to the secondary statistics file.  This will
	 * be cross-checked with the key read from main stats file at loading
	 * time.
	 */
	pgstat_write_chunk_s(fd_description, (PgStat_HashKey *) key);
	fd_description_offset += sizeof(PgStat_HashKey);

	if (!custom_stats_description_dsa)
		custom_stats_description_dsa = GetNamedDSA("test_custom_stat_dsa", &found);

	/* Handle entries without descriptions */
	if (!DsaPointerIsValid(entry->description) || !custom_stats_description_dsa)
	{
		/* length to description file */
		len = 0;
		pgstat_write_chunk_s(fd_description, &len);
		fd_description_offset += sizeof(size_t);
		return;
	}

	/*
	 * Retrieve description from DSA, then write the length followed by the
	 * description.
	 */
	description = dsa_get_address(custom_stats_description_dsa,
								  entry->description);
	len = strlen(description) + 1;
	pgstat_write_chunk_s(fd_description, &len);
	pgstat_write_chunk(fd_description, description, len);

	/*
	 * Update offset for next entry, counting for the length (size_t) of the
	 * description and the description contents.
	 */
	fd_description_offset += len + sizeof(size_t);
}

/*
 * test_custom_stats_var_from_serialized_data() -
 *
 * Read auxiliary data (descriptions) for custom statistics entries from
 * the secondary statistics file.  This is called while loading the statistics
 * at startup.
 *
 * See the top of test_custom_stats_var_to_serialized_data() for a
 * detailed description of the data layout read here.
 */
static bool
test_custom_stats_var_from_serialized_data(const PgStat_HashKey *key,
										   PgStatShared_Common *header,
										   FILE *statfile)
{
	PgStatShared_CustomVarEntry *entry;
	dsa_pointer dp;
	size_t		len;
	pgoff_t		offset;
	char	   *buffer;
	bool		found;
	uint32		magic_number = 0;
	PgStat_HashKey file_key;

	/* Check the magic number first, in the main file. */
	if (!pgstat_read_chunk_s(statfile, &magic_number))
	{
		elog(WARNING, "failed to read magic number from statistics file");
		return false;
	}

	if (magic_number != TEST_CUSTOM_VAR_MAGIC_NUMBER)
	{
		elog(WARNING, "found magic number %u from statistics file, should be %u",
			 magic_number, TEST_CUSTOM_VAR_MAGIC_NUMBER);
		return false;
	}

	/*
	 * Read the offset from the main stats file, to be able to read the
	 * auxiliary data from the secondary statistics file.
	 */
	if (!pgstat_read_chunk_s(statfile, &offset))
	{
		elog(WARNING, "failed to read metadata offset from statistics file");
		return false;
	}

	/* Open statistics file for reading if not already open */
	if (!fd_description)
	{
		fd_description = AllocateFile(TEST_CUSTOM_AUX_DATA_DESC, PG_BINARY_R);
		if (fd_description == NULL)
		{
			if (errno != ENOENT)
				ereport(LOG,
						(errcode_for_file_access(),
						 errmsg("could not open statistics file \"%s\" for reading: %m",
								TEST_CUSTOM_AUX_DATA_DESC)));
			pgstat_reset_of_kind(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS);
			return false;
		}
	}

	/* Read data from the secondary statistics file, at the specified offset */
	if (fseeko(fd_description, offset, SEEK_SET) != 0)
	{
		elog(WARNING, "could not seek in file \"%s\": %m",
			 TEST_CUSTOM_AUX_DATA_DESC);
		return false;
	}

	/* Read the hash key from the secondary statistics file */
	if (!pgstat_read_chunk_s(fd_description, &file_key))
	{
		elog(WARNING, "failed to read hash key from file");
		return false;
	}

	/* Check key consistency */
	if (file_key.kind != key->kind ||
		file_key.dboid != key->dboid ||
		file_key.objid != key->objid)
	{
		elog(WARNING, "found entry key %u/%u/%" PRIu64 " not matching with %u/%u/%" PRIu64,
			 file_key.kind, file_key.dboid, file_key.objid,
			 key->kind, key->dboid, key->objid);
		return false;
	}

	entry = (PgStatShared_CustomVarEntry *) header;

	/* Read the description length and its data */
	if (!pgstat_read_chunk_s(fd_description, &len))
	{
		elog(WARNING, "failed to read metadata length from statistics file");
		return false;
	}

	/* Handle empty descriptions */
	if (len == 0)
	{
		entry->description = InvalidDsaPointer;
		return true;
	}

	/* Initialize DSA if needed */
	if (!custom_stats_description_dsa)
		custom_stats_description_dsa = GetNamedDSA("test_custom_stat_dsa", &found);

	if (!custom_stats_description_dsa)
	{
		elog(WARNING, "could not access DSA for custom statistics descriptions");
		return false;
	}

	buffer = palloc(len);
	if (!pgstat_read_chunk(fd_description, buffer, len))
	{
		pfree(buffer);
		elog(WARNING, "failed to read description from file");
		return false;
	}

	/* Allocate space in DSA and copy the description */
	dp = dsa_allocate(custom_stats_description_dsa, len);
	memcpy(dsa_get_address(custom_stats_description_dsa, dp), buffer, len);
	entry->description = dp;
	pfree(buffer);

	return true;
}

/*
 * test_custom_stats_var_finish() -
 *
 * Cleanup function called at the end of statistics file operations.
 * Handles closing files and cleanup based on the operation type.
 */
static void
test_custom_stats_var_finish(PgStat_StatsFileOp status)
{
	switch (status)
	{
		case STATS_WRITE:
			if (!fd_description)
				return;

			fd_description_offset = 0;

			/* Check for write errors and cleanup if necessary */
			if (ferror(fd_description))
			{
				ereport(LOG,
						(errcode_for_file_access(),
						 errmsg("could not write to file \"%s\": %m",
								TEST_CUSTOM_AUX_DATA_DESC)));
				FreeFile(fd_description);
				unlink(TEST_CUSTOM_AUX_DATA_DESC);
			}
			else if (FreeFile(fd_description) < 0)
			{
				ereport(LOG,
						(errcode_for_file_access(),
						 errmsg("could not close file \"%s\": %m",
								TEST_CUSTOM_AUX_DATA_DESC)));
				unlink(TEST_CUSTOM_AUX_DATA_DESC);
			}
			break;

		case STATS_READ:
			if (fd_description)
				FreeFile(fd_description);

			/* Remove the file after reading */
			elog(DEBUG2, "removing file \"%s\"", TEST_CUSTOM_AUX_DATA_DESC);
			unlink(TEST_CUSTOM_AUX_DATA_DESC);
			break;

		case STATS_DISCARD:
			{
				int			ret;

				/* Attempt to remove the file */
				ret = unlink(TEST_CUSTOM_AUX_DATA_DESC);
				if (ret != 0)
				{
					if (errno == ENOENT)
						elog(LOG,
							 "didn't need to unlink file \"%s\" - didn't exist",
							 TEST_CUSTOM_AUX_DATA_DESC);
					else
						ereport(LOG,
								(errcode_for_file_access(),
								 errmsg("could not unlink file \"%s\": %m",
										TEST_CUSTOM_AUX_DATA_DESC)));
				}
				else
				{
					ereport(LOG,
							(errmsg_internal("unlinked file \"%s\"",
											 TEST_CUSTOM_AUX_DATA_DESC)));
				}
			}
			break;
	}

	fd_description = NULL;
}

/*--------------------------------------------------------------------------
 * Helper functions
 *--------------------------------------------------------------------------
 */

/*
 * test_custom_stats_var_fetch_entry
 *		Look up custom statistic by name
 *
 * Returns statistics entry from shared memory, or NULL if not found.
 */
static PgStat_StatCustomVarEntry *
test_custom_stats_var_fetch_entry(const char *stat_name)
{
	/* Fetch entry by hashed name */
	return (PgStat_StatCustomVarEntry *)
		pgstat_fetch_entry(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS,
						   InvalidOid,
						   PGSTAT_CUSTOM_VAR_STATS_IDX(stat_name),
						   NULL);
}

/*--------------------------------------------------------------------------
 * SQL-callable functions
 *--------------------------------------------------------------------------
 */

/*
 * test_custom_stats_var_create
 *		Create new custom statistic entry
 *
 * Initializes a statistics entry with the given name and description.
 */
PG_FUNCTION_INFO_V1(test_custom_stats_var_create);
Datum
test_custom_stats_var_create(PG_FUNCTION_ARGS)
{
	PgStat_EntryRef *entry_ref;
	PgStatShared_CustomVarEntry *shared_entry;
	char	   *stat_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
	char	   *description = text_to_cstring(PG_GETARG_TEXT_PP(1));
	dsa_pointer dp = InvalidDsaPointer;
	bool		found;

	/* Validate name length first */
	if (strlen(stat_name) >= NAMEDATALEN)
		ereport(ERROR,
				(errcode(ERRCODE_NAME_TOO_LONG),
				 errmsg("custom statistic name \"%s\" is too long", stat_name),
				 errdetail("Name must be less than %d characters.", NAMEDATALEN)));

	/* Initialize DSA and description provided */
	if (!custom_stats_description_dsa)
		custom_stats_description_dsa = GetNamedDSA("test_custom_stat_dsa", &found);

	if (!custom_stats_description_dsa)
		ereport(ERROR,
				(errmsg("could not access DSA for custom statistics descriptions")));

	/* Allocate space in DSA and copy description */
	dp = dsa_allocate(custom_stats_description_dsa, strlen(description) + 1);
	memcpy(dsa_get_address(custom_stats_description_dsa, dp),
		   description,
		   strlen(description) + 1);

	/* Create or get existing entry */
	entry_ref = pgstat_get_entry_ref_locked(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, InvalidOid,
											PGSTAT_CUSTOM_VAR_STATS_IDX(stat_name), true);

	if (!entry_ref)
		PG_RETURN_VOID();

	shared_entry = (PgStatShared_CustomVarEntry *) entry_ref->shared_stats;

	/* Zero-initialize statistics */
	memset(&shared_entry->stats, 0, sizeof(shared_entry->stats));

	/* Store description pointer */
	shared_entry->description = dp;

	pgstat_unlock_entry(entry_ref);

	PG_RETURN_VOID();
}

/*
 * test_custom_stats_var_update
 *		Increment custom statistic counter
 *
 * Increments call count in backend-local memory.  Changes are flushed
 * to shared memory by the statistics collector.
 */
PG_FUNCTION_INFO_V1(test_custom_stats_var_update);
Datum
test_custom_stats_var_update(PG_FUNCTION_ARGS)
{
	char	   *stat_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
	PgStat_EntryRef *entry_ref;
	PgStat_StatCustomVarEntry *pending_entry;

	/* Get pending entry in local memory */
	entry_ref = pgstat_prep_pending_entry(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, InvalidOid,
										  PGSTAT_CUSTOM_VAR_STATS_IDX(stat_name), NULL);

	pending_entry = (PgStat_StatCustomVarEntry *) entry_ref->pending;
	pending_entry->numcalls++;

	PG_RETURN_VOID();
}

/*
 * test_custom_stats_var_drop
 *		Remove custom statistic entry
 *
 * Drops the named statistic from shared memory.
 */
PG_FUNCTION_INFO_V1(test_custom_stats_var_drop);
Datum
test_custom_stats_var_drop(PG_FUNCTION_ARGS)
{
	char	   *stat_name = text_to_cstring(PG_GETARG_TEXT_PP(0));

	/* Drop entry and request GC if the entry could not be freed */
	if (!pgstat_drop_entry(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, InvalidOid,
						   PGSTAT_CUSTOM_VAR_STATS_IDX(stat_name)))
		pgstat_request_entry_refs_gc();

	PG_RETURN_VOID();
}

/*
 * test_custom_stats_var_report
 *		Retrieve custom statistic values
 *
 * Returns single row with statistic name, call count, and description if the
 * statistic exists, otherwise returns no rows.
 */
PG_FUNCTION_INFO_V1(test_custom_stats_var_report);
Datum
test_custom_stats_var_report(PG_FUNCTION_ARGS)
{
	FuncCallContext *funcctx;
	char	   *stat_name;
	PgStat_StatCustomVarEntry *stat_entry;

	if (SRF_IS_FIRSTCALL())
	{
		TupleDesc	tupdesc;
		MemoryContext oldcontext;

		/* Initialize SRF context */
		funcctx = SRF_FIRSTCALL_INIT();
		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

		/* Get composite return type */
		if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
			elog(ERROR, "test_custom_stats_var_report: return type is not composite");

		funcctx->tuple_desc = BlessTupleDesc(tupdesc);
		funcctx->max_calls = 1; /* single row result */

		MemoryContextSwitchTo(oldcontext);
	}

	funcctx = SRF_PERCALL_SETUP();

	if (funcctx->call_cntr < funcctx->max_calls)
	{
		Datum		values[3];
		bool		nulls[3] = {false, false, false};
		HeapTuple	tuple;
		PgStat_EntryRef *entry_ref;
		PgStatShared_CustomVarEntry *shared_entry;
		char	   *description = NULL;
		bool		found;

		stat_name = text_to_cstring(PG_GETARG_TEXT_PP(0));
		stat_entry = test_custom_stats_var_fetch_entry(stat_name);

		/* Return row only if entry exists */
		if (stat_entry)
		{
			/* Get entry ref to access shared entry */
			entry_ref = pgstat_get_entry_ref(PGSTAT_KIND_TEST_CUSTOM_VAR_STATS, InvalidOid,
											 PGSTAT_CUSTOM_VAR_STATS_IDX(stat_name), false, NULL);

			if (entry_ref)
			{
				shared_entry = (PgStatShared_CustomVarEntry *) entry_ref->shared_stats;

				/* Get description from DSA if available */
				if (DsaPointerIsValid(shared_entry->description))
				{
					if (!custom_stats_description_dsa)
						custom_stats_description_dsa = GetNamedDSA("test_custom_stat_dsa", &found);

					if (custom_stats_description_dsa)
						description = dsa_get_address(custom_stats_description_dsa, shared_entry->description);
				}
			}

			values[0] = PointerGetDatum(cstring_to_text(stat_name));
			values[1] = Int64GetDatum(stat_entry->numcalls);

			if (description)
				values[2] = PointerGetDatum(cstring_to_text(description));
			else
				nulls[2] = true;

			tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
			SRF_RETURN_NEXT(funcctx, HeapTupleGetDatum(tuple));
		}
	}

	SRF_RETURN_DONE(funcctx);
}
