/*-------------------------------------------------------------------------
 *
 * logtofile_execution_memory.c
 *      Partial hooks to measure memory footprint of execution
 *
 * Copyright (c) 2026, Francisco Miguel Biete Banon
 *
 * This code is released under the PostgreSQL licence, as given at
 *  http://www.postgresql.org/about/licence/
 *-------------------------------------------------------------------------
 */
#include "logtofile_execution_memory.h"

#include "logtofile_vars.h"

#include <utils/memutils.h>

/* forward declaration private functions */
inline static Size pgauditlogtofile_MemoryContextTotalAllocated(MemoryContext ctx) __attribute__((always_inline));
inline static MemoryContext pgauditlogtofile_get_query_memory_context(QueryDesc *queryDesc) __attribute__((always_inline));
inline static void pgauditlogtofile_update_peak_memory(Size current) __attribute__((always_inline));

/**
 * @brief ExecutorStart hook to record the memory usage at the start of a statement.
 * @param queryDesc
 * @param eflags
 */
void PgAuditLogToFile_ExecutorStart_Memory(QueryDesc *queryDesc, __attribute__((unused)) int eflags)
{
  MemoryContext ctx = pgauditlogtofile_get_query_memory_context(queryDesc);

  pgaudit_ltf_statement_memory_start = pgauditlogtofile_MemoryContextTotalAllocated(ctx);
  pgaudit_ltf_statement_memory_peak = pgaudit_ltf_statement_memory_start;
  pgaudit_ltf_statement_memory_end = 0;
}

/**
 * @brief ExecutorEnd hook to calculate and log the statement memory usage.
 * @param queryDesc
 */
void PgAuditLogToFile_ExecutorEnd_Memory(QueryDesc *queryDesc)
{
  MemoryContext ctx = pgauditlogtofile_get_query_memory_context(queryDesc);

  pgaudit_ltf_statement_memory_end = pgauditlogtofile_MemoryContextTotalAllocated(ctx);
  pgauditlogtofile_update_peak_memory(pgaudit_ltf_statement_memory_end);
}

/**
 * @brief ExecutorRun hook to capture peak of memory usage during run
 * @param queryDesc
 */
#if (PG_VERSION_NUM >= 180000)
void PgAuditLogToFile_ExecutorRun_Memory(QueryDesc *queryDesc,
                                         __attribute__((unused)) ScanDirection direction,
                                         __attribute__((unused)) uint64 count)
#else
void PgAuditLogToFile_ExecutorRun_Memory(QueryDesc *queryDesc,
                                         __attribute__((unused)) ScanDirection direction,
                                         __attribute__((unused)) uint64 count,
                                         __attribute__((unused)) bool execute_once)
#endif
{
  MemoryContext ctx = pgauditlogtofile_get_query_memory_context(queryDesc);
  Size current = pgauditlogtofile_MemoryContextTotalAllocated(ctx);

  pgauditlogtofile_update_peak_memory(current);
}

/* private functions */

/**
 * @brief Obtains memory allocated
 * @param ctx
 * @return Size
 */
static Size
pgauditlogtofile_MemoryContextTotalAllocated(MemoryContext ctx)
{
  if (ctx == NULL)
    return 0;

  return MemoryContextMemAllocated(ctx, true);
}

/**
 * @brief Obtains the query context
 * @param queryDesc
 * @return MemoryContext
 */
static MemoryContext
pgauditlogtofile_get_query_memory_context(QueryDesc *queryDesc)
{
  return (queryDesc && queryDesc->estate) ? queryDesc->estate->es_query_cxt : NULL;
}

/**
 * @brief Update the peak memory value if required
 * @param current
 */
static void
pgauditlogtofile_update_peak_memory(Size current)
{
  if (current > pgaudit_ltf_statement_memory_peak)
    pgaudit_ltf_statement_memory_peak = current;
}
