// pg_stat_ch executor hooks implementation

#include <sys/resource.h>

#include "postgres.h"

#include "access/parallel.h"
#include "access/xact.h"
#include "commands/dbcommands.h"
#include "common/ip.h"
#include "common/pg_prng.h"
#include "executor/executor.h"
#include "executor/instrument.h"
#include "libpq/libpq-be.h"
#include "miscadmin.h"
#include "postmaster/bgworker.h"
#include "tcop/utility.h"
#include "utils/elog.h"
#include "utils/guc.h"
#include "utils/lsyscache.h"
#include "utils/timestamp.h"

#include "parser/analyze.h"

#if PG_VERSION_NUM >= 140000
#include "nodes/queryjumble.h"
#endif

#if PG_VERSION_NUM >= 150000
#include "jit/jit.h"
#endif

#include "hooks/query_normalize_state.h"

#include "config/guc.h"
#include "hooks/hooks.h"
#include "hooks/query_normalize.h"
#include "hooks/string_utils.h"
#include "queue/event.h"
#include "queue/shmem.h"

// Previous hook values for chaining
static post_parse_analyze_hook_type prev_post_parse_analyze = NULL;
static ExecutorStart_hook_type prev_executor_start = NULL;
static ExecutorRun_hook_type prev_executor_run = NULL;
static ExecutorFinish_hook_type prev_executor_finish = NULL;
static ExecutorEnd_hook_type prev_executor_end = NULL;
static ProcessUtility_hook_type prev_process_utility = NULL;
static emit_log_hook_type prev_emit_log_hook = NULL;

// Track nesting level to identify top-level queries
static int nesting_level = 0;

// CPU time tracking via getrusage
static struct rusage rusage_start;

// Deadlock prevention for emit_log_hook
static bool disable_error_capture = false;

// Track whether the current query started at top level
static bool current_query_is_top_level = false;

// Track query start time for duration calculation
static TimestampTz query_start_ts = 0;

// System initialization flag - set after hooks are installed and shmem is ready
static bool system_init = false;

static int GetClientAddress(char* buf, int buf_size);
static void ResolveNames(PschEvent* event);

// Cache for session-stable values to avoid repeated catalog lookups on every query.
// Following pg_stat_monitor's pattern of caching client IP (pg_stat_monitor.c:73-96).
// Database name and client address never change within a session. Username is
// re-resolved when userid changes (handles SET ROLE). This also carries the
// session-local registry of parse-time query text looked up later by
// ExecutorEnd or ProcessUtility.
typedef struct PschBackendState {
  bool initialized;
  char datname[NAMEDATALEN];
  uint8 datname_len;
  Oid cached_userid;
  char username[NAMEDATALEN];
  uint8 username_len;
  char client_addr[46];  // INET6_ADDRSTRLEN
  uint8 client_addr_len;
  PschNormalizedQueryCache normalize_cache;
} PschBackendState;
static PschBackendState backend_state = {0};

// Resolve and cache the current username. On initial resolve, falls back to
// "<unknown>" if resolution fails. On SET ROLE re-resolve, keeps the existing
// cached value on failure (better to show the old name than "<unknown>") and
// leaves cached_userid unchanged so future calls can retry resolution.
static void CacheUsername(Oid userid, bool fallback_on_null) {
  const char* username = GetUserNameFromId(userid, true);
  if (username != NULL) {
    backend_state.username_len =
        PschCopyName(backend_state.username, sizeof(backend_state.username), username);
    backend_state.cached_userid = userid;
  } else if (fallback_on_null) {
    backend_state.username_len =
        PschCopyName(backend_state.username, sizeof(backend_state.username), "<unknown>");
    backend_state.cached_userid = userid;
  }
}

// Ensure the backend cache is populated. Called on each query; the first call
// resolves datname and client_addr (session-stable), and all calls check whether
// userid changed (SET ROLE) to re-resolve username.
static void EnsureBackendCache(void) {
  Oid userid = GetUserId();

  if (!backend_state.initialized) {
    // Can't resolve catalog names outside a transaction
    if (!IsTransactionState()) {
      return;
    }

    // Database name (session-stable)
    const char* datname = get_database_name(MyDatabaseId);
    backend_state.datname_len = PschCopyName(backend_state.datname, sizeof(backend_state.datname),
                                             datname != NULL ? datname : "<unknown>");

    // Client address (session-stable)
    backend_state.client_addr_len = (uint8)(
        GetClientAddress(backend_state.client_addr, sizeof(backend_state.client_addr)));

    // Username (may change via SET ROLE)
    CacheUsername(userid, true);

    backend_state.initialized = true;
    return;
  }

  // Re-resolve username if userid changed (SET ROLE)
  if (backend_state.cached_userid != userid) {
    if (IsTransactionState()) {
      CacheUsername(userid, false);
    }
  }
}

static PschCmdType ConvertCmdType(CmdType cmd) {
  switch (cmd) {
    case CMD_SELECT:
      return PSCH_CMD_SELECT;
    case CMD_UPDATE:
      return PSCH_CMD_UPDATE;
    case CMD_INSERT:
      return PSCH_CMD_INSERT;
    case CMD_DELETE:
      return PSCH_CMD_DELETE;
#if PG_VERSION_NUM >= 150000
    case CMD_MERGE:
      return PSCH_CMD_MERGE;
#endif
    case CMD_UTILITY:
      return PSCH_CMD_UTILITY;
    case CMD_NOTHING:
      return PSCH_CMD_NOTHING;
    default:
      return PSCH_CMD_UNKNOWN;
  }
}

static int64 TimeDiffMicrosec(struct timeval end, struct timeval start) {
  return ((int64)(end.tv_sec - start.tv_sec) * 1000000LL) +
         (int64)(end.tv_usec - start.tv_usec);
}

// Unpack SQLSTATE code from PostgreSQL's packed format to string
static void UnpackSqlState(int sql_state, char* buf) {
  for (int i = 0; i < 5; i++) {
    buf[i] = PGUNSIXBIT(sql_state);
    sql_state >>= 6;
  }
  buf[5] = '\0';
}

static int GetApplicationName(char* buf, int buf_size) {
  // The application_name GUC is the authoritative live value (the beentry's
  // st_appname is populated from it, never the other way around).
  if (application_name != NULL && application_name[0] != '\0') {
    return (int)(PschCopyTrimmed(buf, buf_size, application_name));
  }

  buf[0] = '\0';
  return 0;
}

static int GetClientAddress(char* buf, int buf_size) {
  buf[0] = '\0';

  // MyProcPort->raddr is what pgstat_bestart copies into the beentry's
  // st_clientaddr; it is immutable for the life of the session. NULL in
  // background/aux processes, which have no client connection.
  if (MyProcPort == NULL) {
    return 0;
  }

  char remote_host[NI_MAXHOST] = {0};

  int ret = pg_getnameinfo_all(&MyProcPort->raddr.addr, MyProcPort->raddr.salen, remote_host,
                               sizeof(remote_host), NULL, 0, NI_NUMERICHOST | NI_NUMERICSERV);

  if (ret != 0 || remote_host[0] == '\0') {
    return 0;
  }

  // Handle local connections
  if (strcmp(remote_host, "[local]") == 0) {
    size_t src_len = strlcpy(buf, "127.0.0.1", buf_size);
    return (int)(Min(src_len, (size_t)(buf_size - 1)));
  }

  size_t src_len = strlcpy(buf, remote_host, buf_size);
  return (int)(Min(src_len, (size_t)(buf_size - 1)));
}

// Check whether an event should be captured based on duration thresholds
// and sampling rate. Queries at or above min_duration_us are always captured;
// faster queries are randomly sampled at the configured rate.
static bool ShouldSampleEvent(uint64 duration_us) {
  if (psch_min_duration_us == 0 && psch_sample_rate >= 1.0) {
    return true;
  }
  if (duration_us >= (uint64)(psch_min_duration_us)) {
    return true;
  }
  if (psch_sample_rate <= 0.0) {
    return false;
  }
  if (psch_sample_rate >= 1.0) {
    return true;
  }
  return pg_prng_double(&pg_global_prng_state) < psch_sample_rate;
}

static void CopyBufferUsage(PschEvent* event, const BufferUsage* buf) {
  event->shared_blks_hit = buf->shared_blks_hit;
  event->shared_blks_read = buf->shared_blks_read;
  event->shared_blks_dirtied = buf->shared_blks_dirtied;
  event->shared_blks_written = buf->shared_blks_written;
  event->local_blks_hit = buf->local_blks_hit;
  event->local_blks_read = buf->local_blks_read;
  event->local_blks_dirtied = buf->local_blks_dirtied;
  event->local_blks_written = buf->local_blks_written;
  event->temp_blks_read = buf->temp_blks_read;
  event->temp_blks_written = buf->temp_blks_written;
}

// Copy I/O timing to event (version-aware)
static void CopyIoTiming(PschEvent* event, const BufferUsage* buf) {
#if PG_VERSION_NUM >= 170000
  event->shared_blk_read_time_us = INSTR_TIME_GET_MICROSEC(buf->shared_blk_read_time);
  event->shared_blk_write_time_us = INSTR_TIME_GET_MICROSEC(buf->shared_blk_write_time);
  event->local_blk_read_time_us = INSTR_TIME_GET_MICROSEC(buf->local_blk_read_time);
  event->local_blk_write_time_us = INSTR_TIME_GET_MICROSEC(buf->local_blk_write_time);
#else
  // PG16 and earlier: blk_read_time/blk_write_time (no local block timing)
  event->shared_blk_read_time_us = INSTR_TIME_GET_MICROSEC(buf->blk_read_time);
  event->shared_blk_write_time_us = INSTR_TIME_GET_MICROSEC(buf->blk_write_time);
#endif
#if PG_VERSION_NUM >= 150000
  event->temp_blk_read_time_us = INSTR_TIME_GET_MICROSEC(buf->temp_blk_read_time);
  event->temp_blk_write_time_us = INSTR_TIME_GET_MICROSEC(buf->temp_blk_write_time);
#endif
}

static void CopyWalUsage(PschEvent* event, const WalUsage* wal) {
  event->wal_records = wal->wal_records;
  event->wal_fpi = wal->wal_fpi;
  event->wal_bytes = wal->wal_bytes;
}

// Initialize PschEvent by zeroing only the fixed-size prefix instead of the full
// struct (~4.5KB).  After the field reorder in event.h, everything before
// err_message is fixed-size, so one memset covers the numeric fields, names,
// client context, and stored lengths. The two large text buffers only need a
// leading '\0'; later code either overwrites them or leaves len=0.
static void InitEventPartial(PschEvent* event) {
  const size_t fixed_prefix_size = offsetof(PschEvent, err_message);
  memset(event, 0, fixed_prefix_size);
  event->err_message[0] = '\0';
  event->query[0] = '\0';
}

static void InitBaseEvent(PschEvent* event, TimestampTz ts_start, bool top_level,
                          PschCmdType cmd_type) {
  InitEventPartial(event);
  event->ts_start = ts_start;
  event->dbid = MyDatabaseId;
  event->userid = GetUserId();
  event->pid = MyProcPid;
  event->top_level = top_level;
  event->cmd_type = cmd_type;
  ResolveNames(event);
}

static void CopyClientContext(PschEvent* event) {
  event->application_name_len = (uint8)(
      GetApplicationName(event->application_name, sizeof(event->application_name)));

  EnsureBackendCache();
  if (backend_state.initialized) {
    memcpy(event->client_addr, backend_state.client_addr, backend_state.client_addr_len + 1);
    event->client_addr_len = backend_state.client_addr_len;
  } else {
    event->client_addr_len =
        (uint8)(GetClientAddress(event->client_addr, sizeof(event->client_addr)));
  }
}

// Copy query text into the event buffer from the parse-time LRU cache.
//
// We only export text that was stashed during post_parse_analyze_hook:
// - constant-bearing statements export normalized text
// - constant-free statements export the unchanged statement slice
// - on a cache miss, query text is left empty instead of falling back to
//   raw SQL at execution time
static void CopyQueryText(PschEvent* event, uint64 query_id) {
  PschLookupNormalizedQuery(&backend_state.normalize_cache, query_id, event->query,
                            sizeof(event->query), &event->query_len);
}

// Resolve database and user names, using the session cache when available.
// Falls back to catalog lookups if cache hasn't been initialized yet (e.g.,
// emit_log_hook fires before the first executor hook).
static void ResolveNames(PschEvent* event) {
  EnsureBackendCache();

  if (backend_state.initialized) {
    memcpy(event->datname, backend_state.datname, backend_state.datname_len + 1);
    event->datname_len = backend_state.datname_len;
    memcpy(event->username, backend_state.username, backend_state.username_len + 1);
    event->username_len = backend_state.username_len;
    return;
  }

  // Fallback: resolve fresh (cache not yet initialized, e.g. emit_log_hook early)
  const char* datname = NULL;
  const char* username = NULL;
  if (IsTransactionState()) {
    datname = get_database_name(event->dbid);
    username = GetUserNameFromId(event->userid, true);
  }

  event->datname_len = PschCopyName(event->datname, sizeof(event->datname),
                                    datname != NULL ? datname : "<unknown>");
  event->username_len = PschCopyName(event->username, sizeof(event->username),
                                     username != NULL ? username : "<unknown>");
}

// Copy JIT instrumentation to event (PG15+)
static void CopyJitInstrumentation(PschEvent* event pg_attribute_unused(),
                                   QueryDesc* query_desc pg_attribute_unused()) {
#if PG_VERSION_NUM >= 150000
  if (query_desc->estate->es_jit != NULL) {
    JitInstrumentation* jit = &query_desc->estate->es_jit->instr;
    event->jit_functions = (int32)(jit->created_functions);
    event->jit_generation_time_us =
        (int32)(INSTR_TIME_GET_MICROSEC(jit->generation_counter));
    event->jit_inlining_time_us =
        (int32)(INSTR_TIME_GET_MICROSEC(jit->inlining_counter));
    event->jit_optimization_time_us =
        (int32)(INSTR_TIME_GET_MICROSEC(jit->optimization_counter));
    event->jit_emission_time_us =
        (int32)(INSTR_TIME_GET_MICROSEC(jit->emission_counter));
#if PG_VERSION_NUM >= 170000
    event->jit_deform_time_us = (int32)(INSTR_TIME_GET_MICROSEC(jit->deform_counter));
#endif
  }
#endif
}

// Copy parallel worker info to event (PG18+)
static void CopyParallelWorkerInfo(PschEvent* event pg_attribute_unused(),
                                   QueryDesc* query_desc pg_attribute_unused()) {
#if PG_VERSION_NUM >= 180000
  if (query_desc->estate != NULL) {
    event->parallel_workers_planned =
        (int16)(query_desc->estate->es_parallel_workers_to_launch);
    event->parallel_workers_launched =
        (int16)(query_desc->estate->es_parallel_workers_launched);
  }
#endif
}

static void BuildEventFromQueryDesc(QueryDesc* query_desc, PschEvent* event, int64 cpu_user_us,
                                    int64 cpu_sys_us) {
  InitBaseEvent(event, query_start_ts, current_query_is_top_level,
                ConvertCmdType(query_desc->operation));
  event->queryid = query_desc->plannedstmt->queryId;
  event->rows = query_desc->estate->es_processed;
  event->cpu_user_time_us = cpu_user_us;
  event->cpu_sys_time_us = cpu_sys_us;

  // Instrumentation data (duration, buffer, WAL)
  if (query_desc->totaltime != NULL) {
#if PG_VERSION_NUM >= 190000
    event->duration_us = (uint64)(INSTR_TIME_GET_MICROSEC(query_desc->totaltime->total));
#else
    event->duration_us = (uint64)(query_desc->totaltime->total * 1000000.0);
#endif
    CopyBufferUsage(event, &query_desc->totaltime->bufusage);
    CopyIoTiming(event, &query_desc->totaltime->bufusage);
    CopyWalUsage(event, &query_desc->totaltime->walusage);
  } else {
    event->duration_us = (uint64)(GetCurrentTimestamp() - query_start_ts);
  }

  CopyJitInstrumentation(event, query_desc);
  CopyParallelWorkerInfo(event, query_desc);
  CopyClientContext(event);
  CopyQueryText(event, query_desc->plannedstmt->queryId);
}

// post_parse_analyze_hook — decide query text at parse time.
// The JumbleState (with constant locations) is only available here, so we
// must generate any normalized form now and stash the final exported text for
// ExecutorEnd.
static void PschPostParseAnalyze(ParseState* pstate, Query* query, JumbleState* jstate) {
  if (prev_post_parse_analyze != NULL) {
    prev_post_parse_analyze(pstate, query, jstate);
  }

  if (!psch_enabled || IsParallelWorker()) {
    return;
  }

  // Nothing to cache without a queryId (the executor/utility paths also skip
  // queryId == 0, so there is no consumer for this text).
  if (query->queryId == UINT64CONST(0)) {
    return;
  }

  const char* query_text = pstate->p_sourcetext;
  int query_loc = query->stmt_location;
  int query_len = query->stmt_len;

  // CleanQuerytext slices a multi-statement source string down to the current
  // statement before we either normalize literals or store the unchanged text.
  query_text = CleanQuerytext(query_text, &query_loc, &query_len);

  // Allocate the normalized/trimmed text in CurrentMemoryContext (typically the
  // query context). PschRememberNormalizedQuery copies it into the cache's own
  // long-lived context, so this allocation can be short-lived.
  char* exported_query = NULL;
  if (jstate != NULL && jstate->clocations_count > 0) {
    exported_query = PschNormalizeQuery(query_text, query_loc, &query_len, jstate);
  } else {
    exported_query = PschCopyTrimmedStatement(query_text, query_len);
    if (exported_query != NULL) {
      query_len = (int)(strlen(exported_query));
    }
  }

  if (exported_query != NULL) {
    PschRememberNormalizedQuery(&backend_state.normalize_cache, query->queryId, exported_query,
                                query_len);
    pfree(exported_query);
  }
}

static void PschExecutorStart(QueryDesc* query_desc, int eflags) {
  if (IsParallelWorker()) {
    if (prev_executor_start != NULL) {
      prev_executor_start(query_desc, eflags);
    } else {
      standard_ExecutorStart(query_desc, eflags);
    }
    return;
  }

  // Record if this is a top-level query (before nesting_level changes in Run)
  if (nesting_level == 0) {
    current_query_is_top_level = true;
    query_start_ts = GetCurrentTimestamp();
    // Capture CPU time baseline for top-level queries
    if (psch_enabled) {
      getrusage(RUSAGE_SELF, &rusage_start);
    }
  } else {
    current_query_is_top_level = false;
  }

  if (prev_executor_start != NULL) {
    prev_executor_start(query_desc, eflags);
  } else {
    standard_ExecutorStart(query_desc, eflags);
  }

  if (psch_enabled && query_desc->plannedstmt->queryId != UINT64CONST(0)) {
    if (query_desc->totaltime == NULL) {
      MemoryContext oldcxt = MemoryContextSwitchTo(query_desc->estate->es_query_cxt);
#if PG_VERSION_NUM < 140000
      query_desc->totaltime = InstrAlloc(1, INSTRUMENT_ALL);
#else
      query_desc->totaltime = InstrAlloc(1, INSTRUMENT_ALL, false);
#endif
      MemoryContextSwitchTo(oldcxt);
    }
  }
}

#if PG_VERSION_NUM >= 180000
static void PschExecutorRun(QueryDesc* query_desc, ScanDirection direction, uint64 count) {
#else
static void PschExecutorRun(QueryDesc* query_desc, ScanDirection direction, uint64 count,
                            bool execute_once) {
#endif
  if (IsParallelWorker()) {
#if PG_VERSION_NUM >= 180000
    if (prev_executor_run != NULL) {
      prev_executor_run(query_desc, direction, count);
    } else {
      standard_ExecutorRun(query_desc, direction, count);
    }
#else
    if (prev_executor_run != NULL) {
      prev_executor_run(query_desc, direction, count, execute_once);
    } else {
      standard_ExecutorRun(query_desc, direction, count, execute_once);
    }
#endif
    return;
  }

  nesting_level++;
  PG_TRY();
  {
#if PG_VERSION_NUM >= 180000
    if (prev_executor_run != NULL) {
      prev_executor_run(query_desc, direction, count);
    } else {
      standard_ExecutorRun(query_desc, direction, count);
    }
#else
    if (prev_executor_run != NULL) {
      prev_executor_run(query_desc, direction, count, execute_once);
    } else {
      standard_ExecutorRun(query_desc, direction, count, execute_once);
    }
#endif
  }
  PG_FINALLY();
  { nesting_level--; }
  PG_END_TRY();
}

static void PschExecutorFinish(QueryDesc* query_desc) {
  if (IsParallelWorker()) {
    if (prev_executor_finish != NULL) {
      prev_executor_finish(query_desc);
    } else {
      standard_ExecutorFinish(query_desc);
    }
    return;
  }

  nesting_level++;
  PG_TRY();
  {
    if (prev_executor_finish != NULL) {
      prev_executor_finish(query_desc);
    } else {
      standard_ExecutorFinish(query_desc);
    }
  }
  PG_FINALLY();
  { nesting_level--; }
  PG_END_TRY();
}

static void PschExecutorEnd(QueryDesc* query_desc) {
  if (!psch_enabled || IsParallelWorker() || query_desc->plannedstmt->queryId == UINT64CONST(0)) {
    if (prev_executor_end != NULL) {
      prev_executor_end(query_desc);
    } else {
      standard_ExecutorEnd(query_desc);
    }
    return;
  }

  if (query_desc->totaltime != NULL) {
    InstrEndLoop(query_desc->totaltime);
  }

  // Compute duration early for sampling filter
  uint64 duration_us;
  if (query_desc->totaltime != NULL) {
#if PG_VERSION_NUM >= 190000
    duration_us = (uint64)(INSTR_TIME_GET_MICROSEC(query_desc->totaltime->total));
#else
    duration_us = (uint64)(query_desc->totaltime->total * 1000000.0);
#endif
  } else {
    duration_us = (uint64)(GetCurrentTimestamp() - query_start_ts);
  }

  if (!ShouldSampleEvent(duration_us)) {
    if (prev_executor_end != NULL) {
      prev_executor_end(query_desc);
    } else {
      standard_ExecutorEnd(query_desc);
    }
    return;
  }

  // Compute CPU time delta from getrusage
  int64 cpu_user_us = 0;
  int64 cpu_sys_us = 0;
  struct rusage rusage_end;
  if (getrusage(RUSAGE_SELF, &rusage_end) == 0) {
    cpu_user_us = TimeDiffMicrosec(rusage_end.ru_utime, rusage_start.ru_utime);
    cpu_sys_us = TimeDiffMicrosec(rusage_end.ru_stime, rusage_start.ru_stime);
  }

  PschEvent event;
  BuildEventFromQueryDesc(query_desc, &event, cpu_user_us, cpu_sys_us);
  PschEnqueueEvent(&event);

  if (prev_executor_end != NULL) {
    prev_executor_end(query_desc);
  } else {
    standard_ExecutorEnd(query_desc);
  }
}

// Build a PschEvent for utility statements (no QueryDesc available)
static void BuildEventForUtility(PschEvent* event, uint64 query_id, TimestampTz start_ts,
                                 uint64 duration_us, bool is_top_level, uint64 rows,
                                 BufferUsage* bufusage, WalUsage* walusage, int64 cpu_user_us,
                                 int64 cpu_sys_us) {
  InitBaseEvent(event, start_ts, is_top_level, PSCH_CMD_UTILITY);
  event->queryid = query_id;
  event->duration_us = duration_us;
  event->rows = rows;
  event->cpu_user_time_us = cpu_user_us;
  event->cpu_sys_time_us = cpu_sys_us;

  CopyBufferUsage(event, bufusage);
  CopyIoTiming(event, bufusage);
  CopyWalUsage(event, walusage);
  CopyClientContext(event);
  CopyQueryText(event, query_id);
}

// Helper macro to call ProcessUtility (previous hook or standard)
#if PG_VERSION_NUM >= 140000
#define CALL_PROCESS_UTILITY()                                                                     \
  do {                                                                                             \
    if (prev_process_utility) {                                                                    \
      prev_process_utility(pstmt, queryString, readOnlyTree, context, params, queryEnv, dest, qc); \
    } else {                                                                                       \
      standard_ProcessUtility(pstmt, queryString, readOnlyTree, context, params, queryEnv, dest,   \
                              qc);                                                                 \
    }                                                                                              \
  } while (0)
#else
#define CALL_PROCESS_UTILITY()                                                          \
  do {                                                                                  \
    if (prev_process_utility) {                                                         \
      prev_process_utility(pstmt, queryString, context, params, queryEnv, dest, qc);    \
    } else {                                                                            \
      standard_ProcessUtility(pstmt, queryString, context, params, queryEnv, dest, qc); \
    }                                                                                   \
  } while (0)
#endif

static bool ShouldTrackUtility(Node* parsetree) {
  if (!psch_enabled || IsParallelWorker()) {
    return false;
  }
  // Skip EXECUTE/PREPARE/DEALLOCATE to avoid double-counting
  if (IsA(parsetree, ExecuteStmt) || IsA(parsetree, PrepareStmt) ||
      IsA(parsetree, DeallocateStmt)) {
    return false;
  }
  // Skip transaction control statements (BEGIN, COMMIT, ROLLBACK, SAVEPOINT,
  // etc.).  They carry no meaningful telemetry — no query text, no buffer
  // stats, no duration — and at high TPS they consume a significant share of
  // queue capacity (2 out of 7 events per pgbench TPC-B transaction).
  if (IsA(parsetree, TransactionStmt)) {
    return false;
  }
  return true;
}

static uint64 GetUtilityRowCount(QueryCompletion* qc) {
  if (qc == NULL) {
    return 0;
  }
  switch (qc->commandTag) {
    case CMDTAG_COPY:
    case CMDTAG_FETCH:
    case CMDTAG_SELECT:
    case CMDTAG_REFRESH_MATERIALIZED_VIEW:
      return qc->nprocessed;
    default:
      return 0;
  }
}

static void ExecuteUtilityWithNesting(PlannedStmt* pstmt, const char* queryString,
#if PG_VERSION_NUM >= 140000
                                      bool readOnlyTree,
#endif
                                      ProcessUtilityContext context, ParamListInfo params,
                                      QueryEnvironment* queryEnv, DestReceiver* dest,
                                      QueryCompletion* qc) {
  nesting_level++;
  PG_TRY();
  { CALL_PROCESS_UTILITY(); }
  PG_FINALLY();
  { nesting_level--; }
  PG_END_TRY();
}

// ProcessUtility hook - captures DDL and utility statements
#if PG_VERSION_NUM >= 140000
static void PschProcessUtility(PlannedStmt* pstmt, const char* queryString, bool readOnlyTree,
                               ProcessUtilityContext context, ParamListInfo params,
                               QueryEnvironment* queryEnv, DestReceiver* dest,
                               QueryCompletion* qc) {
#else
static void PschProcessUtility(PlannedStmt* pstmt, const char* queryString,
                               ProcessUtilityContext context, ParamListInfo params,
                               QueryEnvironment* queryEnv, DestReceiver* dest,
                               QueryCompletion* qc) {
#endif
  if (!ShouldTrackUtility(pstmt->utilityStmt)) {
    CALL_PROCESS_UTILITY();
    return;
  }

  // Capture state before execution
  bool is_top_level = (nesting_level == 0);
  TimestampTz start_ts = GetCurrentTimestamp();
  uint64 query_id = pstmt->queryId;
  BufferUsage bufusage_start = pgBufferUsage;
  WalUsage walusage_start = pgWalUsage;
  struct rusage rusage_util_start;
  getrusage(RUSAGE_SELF, &rusage_util_start);
  instr_time start_time;
  INSTR_TIME_SET_CURRENT(start_time);

#if PG_VERSION_NUM >= 140000
  ExecuteUtilityWithNesting(pstmt, queryString, readOnlyTree, context, params, queryEnv, dest, qc);
#else
  ExecuteUtilityWithNesting(pstmt, queryString, context, params, queryEnv, dest, qc);
#endif

  instr_time duration;
  INSTR_TIME_SET_CURRENT(duration);
  INSTR_TIME_SUBTRACT(duration, start_time);

  uint64 duration_us = INSTR_TIME_GET_MICROSEC(duration);
  if (!ShouldSampleEvent(duration_us)) {
    return;
  }

  BufferUsage bufusage_delta;
  WalUsage walusage_delta;
  MemSet(&bufusage_delta, 0, sizeof(BufferUsage));
  MemSet(&walusage_delta, 0, sizeof(WalUsage));
  BufferUsageAccumDiff(&bufusage_delta, &pgBufferUsage, &bufusage_start);
  WalUsageAccumDiff(&walusage_delta, &pgWalUsage, &walusage_start);

  int64 cpu_user_us = 0;
  int64 cpu_sys_us = 0;
  struct rusage rusage_util_end;
  if (getrusage(RUSAGE_SELF, &rusage_util_end) == 0) {
    cpu_user_us = TimeDiffMicrosec(rusage_util_end.ru_utime, rusage_util_start.ru_utime);
    cpu_sys_us = TimeDiffMicrosec(rusage_util_end.ru_stime, rusage_util_start.ru_stime);
  }

  PschEvent event;
  BuildEventForUtility(&event, query_id, start_ts, duration_us, is_top_level,
                       GetUtilityRowCount(qc), &bufusage_delta, &walusage_delta, cpu_user_us,
                       cpu_sys_us);
  PschEnqueueEvent(&event);
}

#undef CALL_PROCESS_UTILITY

// Check if log capture should occur for this error.
// Returns false during early initialization, in background workers, or when disabled.
static bool ShouldCaptureLog(ErrorData* edata) {
  // Basic preconditions
  if (edata == NULL || !system_init || !psch_enabled || disable_error_capture) {
    return false;
  }

  // Check error level threshold
  if (edata->elevel < psch_log_min_elevel) {
    return false;
  }

  // PostgreSQL bootstrapping checks - MyProc indicates PGPROC allocation complete
  if (MyProc == NULL || IsParallelWorker()) {
    return false;
  }

  // Session initialization checks (critical for debug5 safety):
  // - MyDatabaseId: database must be assigned
  // - IsUnderPostmaster: not single-user mode or bootstrap
  // - psch_shared_state: shared memory must be ready
  // - MyBgworkerEntry: skip background workers (not user queries)
  if (MyDatabaseId == InvalidOid || !IsUnderPostmaster || psch_shared_state == NULL ||
      MyBgworkerEntry != NULL) {
    return false;
  }

  return true;
}

// Build and enqueue an error event from ErrorData.
//
// We intentionally leave event.query empty here. emit_log_hook only exposes
// debug_query_string and cursor position, not the exact statement identity used
// by ExecutorEnd/ProcessUtility, so reconstructing normalized SQL required
// fuzzy matching and extra backend-local state. Error events still carry the
// message, SQLSTATE, and client/session metadata.
// This runs inside errfinish (via PschEmitLogHook), which holds the recursion
// guard and wraps us in PG_TRY: if anything here throws (e.g. the catalog
// fallback in ResolveNames failing to allocate under OOM), the nested error is
// bounded and swallowed by the hook instead of recursing.
static void CaptureLogEvent(ErrorData* edata) {
  PschEvent event;
  InitBaseEvent(&event, GetCurrentTimestamp(), (nesting_level == 0), PSCH_CMD_UNKNOWN);

  UnpackSqlState(edata->sqlerrcode, event.err_sqlstate);
  event.err_elevel = (uint8)(edata->elevel);

  if (edata->message != NULL) {
    event.err_message_len = (uint16)(
        PschCopyTrimmed(event.err_message, PSCH_MAX_ERR_MSG_LEN, edata->message));
  }

  CopyClientContext(&event);

  PschEnqueueEvent(&event);
}

// emit_log_hook - captures log messages at configured level and above
//
// CHAINING ORDER: We chain to the previous hook FIRST to allow other extensions
// (e.g., log formatters, filters) to transform ErrorData before we capture it.
// This ensures we capture the final, potentially modified log message.
static void PschEmitLogHook(ErrorData* edata) {
  // Chain to previous hook first (allows log transformation by other extensions)
  if (prev_emit_log_hook != NULL) {
    prev_emit_log_hook(edata);
  }

  if (!ShouldCaptureLog(edata)) {
    return;
  }

  // Hold the recursion guard across the whole capture. emit_log_hook runs
  // inside errfinish, where a throw starts a new, nested error report that
  // re-enters this hook BEFORE any longjmp: unguarded, an OOM error whose
  // capture itself fails to allocate recurses until the errordata stack
  // overflows into PANIC -> abort() -> database-wide crash recovery.
  //
  // The guard bounds that recursion at one level (the nested invocation
  // returns at ShouldCaptureLog), and the PG_CATCH does two jobs: it restores
  // the guard -- the nested error's longjmp would otherwise skip the clear
  // below and permanently disable capture in this backend -- and it swallows
  // the nested error so the original errfinish resumes and delivers the real
  // report. The cost of a throw during capture is one lost telemetry event.
  MemoryContext oldcxt = CurrentMemoryContext;
  disable_error_capture = true;
  PG_TRY();
  {
    CaptureLogEvent(edata);
  }
  PG_CATCH();
  {
    MemoryContextSwitchTo(oldcxt);
    FlushErrorState();
  }
  PG_END_TRY();
  disable_error_capture = false;
}

// Set or clear the error-capture guard, returning the previous value so
// callers nest correctly (e.g. PschEnqueueEvent suppressing capture while
// already inside the emit_log_hook capture, which holds the guard).
bool PschSuppressErrorCapture(bool suppress) {
  bool prev = disable_error_capture;
  disable_error_capture = suppress;
  return prev;
}

void PschInstallHooks(void) {
#if PG_VERSION_NUM >= 140000
  EnableQueryId();
#endif

  prev_post_parse_analyze = post_parse_analyze_hook;
  post_parse_analyze_hook = PschPostParseAnalyze;

  prev_executor_start = ExecutorStart_hook;
  ExecutorStart_hook = PschExecutorStart;

  prev_executor_run = ExecutorRun_hook;
  ExecutorRun_hook = PschExecutorRun;

  prev_executor_finish = ExecutorFinish_hook;
  ExecutorFinish_hook = PschExecutorFinish;

  prev_executor_end = ExecutorEnd_hook;
  ExecutorEnd_hook = PschExecutorEnd;

  prev_process_utility = ProcessUtility_hook;
  ProcessUtility_hook = PschProcessUtility;

  prev_emit_log_hook = emit_log_hook;
  emit_log_hook = PschEmitLogHook;

  // Mark system as initialized - emit_log_hook will now capture messages
  system_init = true;
}
