#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <time.h>
#include <oqs/oqs.h>

static inline uint64_t get_ns(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC_RAW, &ts);
    return (uint64_t)ts.tv_sec * 1000000000ULL + (uint64_t)ts.tv_nsec;
}

static void* safe_malloc(size_t n) {
    void* p = malloc(n);
    if (!p) { fprintf(stderr, "malloc failed\n"); exit(1); }
    return p;
}

int main(int argc, char *argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <num_traces> <decaps|encaps>\n", argv[0]);
        return 1;
    }

    int num_traces = atoi(argv[1]);
    int do_decaps  = strcmp(argv[2], "decaps") == 0;

    OQS_KEM *kem = OQS_KEM_new(OQS_KEM_alg_ml_kem_768);
    if (!kem) {
        fprintf(stderr, "Failed to initialize ML-KEM-768\n");
        return 1;
    }

    uint8_t *pk = safe_malloc(kem->length_public_key);
    uint8_t *sk = safe_malloc(kem->length_secret_key);
    uint8_t *ct = safe_malloc(kem->length_ciphertext);
    uint8_t *ss = safe_malloc(kem->length_shared_secret);

    fprintf(stderr, "[*] Generating ML-KEM-768 key pair...\n");
    OQS_KEM_keypair(kem, pk, sk);

    /* Fixed ciphertext -- constant test vector */
    uint8_t *fixed_ct = safe_malloc(kem->length_ciphertext);
    uint8_t *fixed_ss = safe_malloc(kem->length_shared_secret);
    OQS_KEM_encaps(kem, fixed_ct, fixed_ss, pk);
    fprintf(stderr, "[*] Fixed test vector generated.\n");

    /*
     * Pre-generate random ciphertexts (for decaps test)
     * and random public keys (for encaps test).
     *
     * This is the critical fix -- previously encaps random was
     * measuring keygen + encaps vs encaps only.
     * Now both measure encaps only with different inputs.
     */
    int NUM_RAND = 1000;

    /* For decaps: pre-generate random ciphertexts */
    uint8_t **rand_cts = malloc(NUM_RAND * sizeof(uint8_t*));
    uint8_t **rand_ss_arr = malloc(NUM_RAND * sizeof(uint8_t*));
    fprintf(stderr, "[*] Pre-generating %d random ciphertexts...\n", NUM_RAND);
    for (int k = 0; k < NUM_RAND; k++) {
        rand_cts[k]    = safe_malloc(kem->length_ciphertext);
        rand_ss_arr[k] = safe_malloc(kem->length_shared_secret);
        OQS_KEM_encaps(kem, rand_cts[k], rand_ss_arr[k], pk);
    }

    /* For encaps: pre-generate random public keys */
    uint8_t **rand_pks = malloc(NUM_RAND * sizeof(uint8_t*));
    fprintf(stderr, "[*] Pre-generating %d random public keys...\n", NUM_RAND);
    for (int k = 0; k < NUM_RAND; k++) {
        rand_pks[k] = safe_malloc(kem->length_public_key);
        uint8_t *tmp_sk = safe_malloc(kem->length_secret_key);
        OQS_KEM_keypair(kem, rand_pks[k], tmp_sk);
        free(tmp_sk);
    }

    /* Warmup -- 500 discarded traces */
    fprintf(stderr, "[*] Warming up (500 discarded traces)...\n");
    for (int i = 0; i < 500; i++) {
        if (do_decaps) {
            OQS_KEM_decaps(kem, ss, fixed_ct, sk);
        } else {
            OQS_KEM_encaps(kem, ct, ss, pk);
        }
    }

    printf("fixed_ns,random_ns\n");
    fflush(stdout);

    fprintf(stderr, "[*] Collecting %d traces...\n", num_traces);

    for (int i = 0; i < num_traces; i++) {
        uint64_t t_fixed, t_random, t_start, t_end;

        /* Fixed measurement -- same input every time */
        t_start = get_ns();
        if (do_decaps) {
            OQS_KEM_decaps(kem, ss, fixed_ct, sk);
        } else {
            OQS_KEM_encaps(kem, ct, ss, pk);  /* fixed public key */
        }
        t_end   = get_ns();
        t_fixed = t_end - t_start;

        /* Random measurement -- pre-generated input, no keygen overhead */
        int idx = i % NUM_RAND;
        t_start = get_ns();
        if (do_decaps) {
            OQS_KEM_decaps(kem, ss, rand_cts[idx], sk);
        } else {
            OQS_KEM_encaps(kem, ct, ss, rand_pks[idx]);  /* pre-generated key */
        }
        t_end    = get_ns();
        t_random = t_end - t_start;

        printf("%llu,%llu\n",
               (unsigned long long)t_fixed,
               (unsigned long long)t_random);

        if (i % 1000 == 0) {
            fprintf(stderr, "    traces: %d/%d\n", i, num_traces);
        }
    }

    /* Cleanup */
    for (int k = 0; k < NUM_RAND; k++) {
        free(rand_cts[k]);
        free(rand_ss_arr[k]);
        free(rand_pks[k]);
    }
    free(rand_cts);
    free(rand_ss_arr);
    free(rand_pks);
    free(pk); free(sk); free(ct); free(ss);
    free(fixed_ct); free(fixed_ss);
    OQS_KEM_free(kem);

    fprintf(stderr, "[*] Done.\n");
    return 0;
}
