#include "hakmem_ace_stats.h"
#include <stdlib.h>

static _Atomic uint64_t g_mid_hit = 0;
static _Atomic uint64_t g_mid_miss = 0;
static _Atomic uint64_t g_large_hit = 0;
static _Atomic uint64_t g_large_miss = 0;
static _Atomic uint64_t g_l1_fallback = 0;
static _Atomic uint64_t g_mid_free = 0;
static _Atomic uint64_t g_large_free = 0;

static _Atomic uint32_t g_sample_mask = 0; // 0=off
static _Atomic uint64_t g_seq = 0;

void hkm_ace_stats_init(void) {
    const char* env = getenv("HAKMEM_ACE_SAMPLE");
    if (env) {
        int n = atoi(env);
        if (n > 0 && n < 31) {
            uint32_t mask = (1u << n) - 1u; // sample every 2^n calls
            atomic_store(&g_sample_mask, mask);
        }
    }
}

void hkm_ace_set_sample_mask(uint32_t mask) {
    atomic_store(&g_sample_mask, mask);
}

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

void hkm_ace_stat_mid_attempt(int hit) {
    if (!should_sample()) return;
    if (hit) atomic_fetch_add(&g_mid_hit, 1);
    else     atomic_fetch_add(&g_mid_miss, 1);
}

void hkm_ace_stat_large_attempt(int hit) {
    if (!should_sample()) return;
    if (hit) atomic_fetch_add(&g_large_hit, 1);
    else     atomic_fetch_add(&g_large_miss, 1);
}

void hkm_ace_stat_l1_fallback(void) {
    if (!should_sample()) return;
    atomic_fetch_add(&g_l1_fallback, 1);
}

void hkm_ace_stat_mid_free(void) {
    if (!should_sample()) return;
    atomic_fetch_add(&g_mid_free, 1);
}

void hkm_ace_stat_large_free(void) {
    if (!should_sample()) return;
    atomic_fetch_add(&g_large_free, 1);
}

void hkm_ace_stats_snapshot(hkm_ace_stats_snapshot_t* out, int reset) {
    if (!out) return;
    out->mid_hit = atomic_load(&g_mid_hit);
    out->mid_miss = atomic_load(&g_mid_miss);
    out->large_hit = atomic_load(&g_large_hit);
    out->large_miss = atomic_load(&g_large_miss);
    out->l1_fallback = atomic_load(&g_l1_fallback);
    out->mid_free = atomic_load(&g_mid_free);
    out->large_free = atomic_load(&g_large_free);
    if (reset) {
        atomic_store(&g_mid_hit, 0);
        atomic_store(&g_mid_miss, 0);
        atomic_store(&g_large_hit, 0);
        atomic_store(&g_large_miss, 0);
        atomic_store(&g_l1_fallback, 0);
        atomic_store(&g_mid_free, 0);
        atomic_store(&g_large_free, 0);
    }
}