/*
 * pqc-bench harness — cycle-count + optional PAPI hardware counter benchmarks.
 *
 * Usage: <binary> [nspins]
 *   nspins  number of outer loop-spin iterations (default: 1)
 *
 * Each spin runs all operations with NTESTS inner iterations and prints one
 * median/average pair per operation. With WITH_PAPI, additional lines are
 * emitted per hardware counter using the same parseable format.
 *
 * Build flags:
 *   -DWITH_PAPI   link against PAPI and emit hardware counter lines
 */

#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "kem.h"
#include "params.h"
#include "indcpa.h"
#include "polyvec.h"
#include "poly.h"
#include "randombytes.h"
#include "cpucycles.h"
#include "speed_print.h"

#ifdef WITH_PAPI
#include <papi.h>
#endif

#define NTESTS 1000

/* ── PAPI instrumentation ───────────────────────────────────────────────── */
#ifdef WITH_PAPI

typedef struct {
  int   code;
  const char *name;
} papi_event_def;

static const papi_event_def DESIRED_EVENTS[] = {
  { PAPI_TOT_INS, "instructions"    },
  { PAPI_L1_DCM,  "l1_misses"       },
  { PAPI_L2_TCM,  "l2_misses"       },
  { PAPI_L3_TCM,  "l3_misses"       },
  { PAPI_BR_MSP,  "branch_mispreds" },
};
#define MAX_EVENTS ((int)(sizeof(DESIRED_EVENTS) / sizeof(DESIRED_EVENTS[0])))

static int        papi_eventset         = PAPI_NULL;
static int        active_codes[MAX_EVENTS];
static const char *active_names[MAX_EVENTS];
static int        n_active              = 0;
static int        papi_ok               = 0; /* set to 1 if init succeeded */

static void papi_init(void) {
  int ret;

  ret = PAPI_library_init(PAPI_VER_CURRENT);
  if (ret != PAPI_VER_CURRENT) {
    fprintf(stderr, "PAPI_library_init: %s — hardware counters disabled\n",
            PAPI_strerror(ret));
    return;
  }

  if ((ret = PAPI_create_eventset(&papi_eventset)) != PAPI_OK) {
    fprintf(stderr, "PAPI_create_eventset: %s — hardware counters disabled\n",
            PAPI_strerror(ret));
    return;
  }

  for (int i = 0; i < MAX_EVENTS; i++) {
    if (PAPI_query_event(DESIRED_EVENTS[i].code) != PAPI_OK) {
      fprintf(stderr, "PAPI: event %s not available on this hardware, skipping\n",
              DESIRED_EVENTS[i].name);
      continue;
    }
    ret = PAPI_add_event(papi_eventset, DESIRED_EVENTS[i].code);
    if (ret != PAPI_OK) {
      fprintf(stderr, "PAPI_add_event(%s): %s — skipping\n",
              DESIRED_EVENTS[i].name, PAPI_strerror(ret));
      continue;
    }
    active_codes[n_active] = DESIRED_EVENTS[i].code;
    active_names[n_active] = DESIRED_EVENTS[i].name;
    n_active++;
  }

  if (n_active == 0) {
    fprintf(stderr, "PAPI: no events could be added — hardware counters disabled\n");
    return;
  }

  if ((ret = PAPI_start(papi_eventset)) != PAPI_OK) {
    fprintf(stderr, "PAPI_start: %s — hardware counters disabled\n",
            PAPI_strerror(ret));
    return;
  }

  papi_ok = 1;
}

/*
 * papi_print — print per-call counter values for one (op, counter) pair.
 * Both "median" and "average" are set to the same per-call value; the outer
 * loop-spin structure gives the aggregation tool a real distribution.
 * The IPC line uses a float value multiplied by 1000 for integer storage;
 * the analysis tool divides by 1000 to recover IPC.
 */
static void papi_print(const char *op, const char *counter,
                       long long total, int ntests)
{
  long long per_call = total / ntests;
  printf("%s_%s: \nmedian: %lld per_call\naverage: %lld per_call\n\n",
         op, counter, per_call, per_call);
}

/*
 * papi_bench — read counters around an already-executed NTESTS block.
 * Call papi_read_before() immediately before the loop and
 * papi_bench_report() immediately after.
 */
static long long _papi_before[MAX_EVENTS];
static long long _papi_after[MAX_EVENTS];

static inline void papi_read_before(void) {
  if (papi_ok) PAPI_read(papi_eventset, _papi_before);
}

static void papi_bench_report(const char *op) {
  if (!papi_ok) return;
  PAPI_read(papi_eventset, _papi_after);
  for (int e = 0; e < n_active; e++) {
    long long delta = _papi_after[e] - _papi_before[e];
    papi_print(op, active_names[e], delta, NTESTS);
  }
}

#define PAPI_BEFORE()        papi_read_before()
#define PAPI_AFTER(op)       papi_bench_report(op)

#else /* !WITH_PAPI */

static inline void papi_init(void) {}
#define PAPI_BEFORE()        ((void)0)
#define PAPI_AFTER(op)       ((void)0)

#endif /* WITH_PAPI */

/* ── Benchmark helpers ───────────────────────────────────────────────────── */

/*
 * BENCH(label, body) — time NTESTS executions of body, print results, then
 * emit PAPI counter lines if enabled.
 */
#define BENCH(label, body)                     \
  do {                                         \
    PAPI_BEFORE();                             \
    for (unsigned int _i = 0; _i < NTESTS; _i++) { \
      t[_i] = cpucycles();                     \
      body;                                    \
    }                                          \
    print_results(label ": ", t, NTESTS);      \
    PAPI_AFTER(label);                         \
  } while (0)

/* ── Main ────────────────────────────────────────────────────────────────── */

static uint64_t t[NTESTS];
static uint8_t  seed[KYBER_SYMBYTES] = {0};

int main(int argc, char *argv[])
{
  int nspins = 1;
  if (argc > 1) {
    nspins = atoi(argv[1]);
    if (nspins <= 0) {
      fprintf(stderr, "usage: %s [nspins]\n", argv[0]);
      return 1;
    }
  }

  papi_init();

  uint8_t pk[CRYPTO_PUBLICKEYBYTES];
  uint8_t sk[CRYPTO_SECRETKEYBYTES];
  uint8_t ct[CRYPTO_CIPHERTEXTBYTES];
  uint8_t key[CRYPTO_BYTES];
  uint8_t coins32[KYBER_SYMBYTES];
  uint8_t coins64[2*KYBER_SYMBYTES];
  polyvec matrix[KYBER_K];
  poly    ap;

  randombytes(coins32, KYBER_SYMBYTES);
  randombytes(coins64, 2*KYBER_SYMBYTES);

  for (int spin = 1; spin <= nspins; spin++) {
    printf("Loop spin: %d\n", spin);

    BENCH("gen_a",
      gen_matrix(matrix, seed, 0));

    BENCH("poly_getnoise_eta1",
      poly_getnoise_eta1(&ap, seed, 0));

    BENCH("poly_getnoise_eta2",
      poly_getnoise_eta2(&ap, seed, 0));

    BENCH("NTT",
      poly_ntt(&ap));

    BENCH("INVNTT",
      poly_invntt_tomont(&ap));

    BENCH("polyvec_basemul_acc_montgomery",
      polyvec_basemul_acc_montgomery(&ap, &matrix[0], &matrix[1]));

    BENCH("poly_tomsg",
      poly_tomsg(ct, &ap));

    BENCH("poly_frommsg",
      poly_frommsg(&ap, ct));

    BENCH("poly_compress",
      poly_compress(ct, &ap));

    BENCH("poly_decompress",
      poly_decompress(&ap, ct));

    BENCH("polyvec_compress",
      polyvec_compress(ct, &matrix[0]));

    BENCH("polyvec_decompress",
      polyvec_decompress(&matrix[0], ct));

    BENCH("indcpa_keypair",
      indcpa_keypair_derand(pk, sk, coins32));

    BENCH("indcpa_enc",
      indcpa_enc(ct, key, pk, seed));

    BENCH("indcpa_dec",
      indcpa_dec(key, ct, sk));

    BENCH("kyber_keypair_derand",
      crypto_kem_keypair_derand(pk, sk, coins64));

    BENCH("kyber_keypair",
      crypto_kem_keypair(pk, sk));

    BENCH("kyber_encaps_derand",
      crypto_kem_enc_derand(ct, key, pk, coins32));

    BENCH("kyber_encaps",
      crypto_kem_enc(ct, key, pk));

    BENCH("kyber_decaps",
      crypto_kem_dec(key, ct, sk));
  }

  return 0;
}
