#include "hakmem_size_hist.h"
#include <stdatomic.h>
#include <stdlib.h>

#define HAK_HIST_MAX_KB 2048

static _Atomic uint64_t g_bins[HAK_HIST_MAX_KB+1];
static _Atomic uint32_t g_mask = 0; // sampling mask
static _Atomic uint64_t g_seq = 0;

void hkm_size_hist_init(void) {
    const char* e = getenv("HAKMEM_HIST_SAMPLE");
    uint32_t mask = 0; // 0=off
    if (e) {
        int n = atoi(e);
        if (n > 0 && n < 31) mask = (1u << n) - 1u;
    }
    atomic_store(&g_mask, mask);
}

static inline int should_sample(void) {
    uint32_t mask = atomic_load(&g_mask);
    if (mask == 0) return 0;
    uint64_t x = atomic_fetch_add(&g_seq, 1);
    return ((x & mask) == 0);
}

void hkm_size_hist_record(size_t size) {
    if (!should_sample()) return;
    size_t kb = (size + 1023) >> 10;
    if (kb > HAK_HIST_MAX_KB) kb = HAK_HIST_MAX_KB;
    atomic_fetch_add(&g_bins[kb], 1);
}

void hkm_size_hist_snapshot(uint64_t* out, int n, int reset) {
    if (!out || n <= 0) return;
    if (n > (HAK_HIST_MAX_KB+1)) n = HAK_HIST_MAX_KB+1;
    for (int i = 0; i < n; i++) {
        out[i] = atomic_load(&g_bins[i]);
        if (reset) atomic_store(&g_bins[i], 0);
    }
}