/**
 * mlkem_test - ML-KEM hardware accelerator comparison test.
 *
 * Runs ML-KEM-512/768/1024 (k=2/3/4) with deterministic KAT count=0 inputs
 * and compares software, hardware through CPU register accesses, and hardware
 * using DMA window transfers.
 */

#include <string.h>
#include <stdio.h>
#include <stdint.h>

#include <confreg_time.h>

#include "dma.h"

unsigned long UART_BASE              = 0xbf000000;
unsigned long CONFREG_TIMER_BASE     = 0xbf20f100;
unsigned long CONFREG_CLOCKS_PER_SEC = 50000000L;
unsigned long CORE_CLOCKS_PER_SEC    = 33000000L;

#define MLKEM_BASE        0xbf400000
#define MLKEM_CTRL        (MLKEM_BASE + 0x000)
#define MLKEM_STATUS      (MLKEM_BASE + 0x004)
#define MLKEM_CONFIG      (MLKEM_BASE + 0x008)
#define MLKEM_D           (MLKEM_BASE + 0x010)
#define MLKEM_Z           (MLKEM_BASE + 0x030)
#define MLKEM_MSG         (MLKEM_BASE + 0x050)
#define MLKEM_SS          (MLKEM_BASE + 0x070)
#define MLKEM_PRE_SEL     (MLKEM_BASE + 0x100)
#define MLKEM_PRE_ADR     (MLKEM_BASE + 0x104)
#define MLKEM_PRE_DAT     (MLKEM_BASE + 0x108)
#define MLKEM_RDB_SEL     (MLKEM_BASE + 0x10C)
#define MLKEM_RDB_ADR     (MLKEM_BASE + 0x110)
#define MLKEM_RDB_DAT     (MLKEM_BASE + 0x114)
#define MLKEM_EK_WINDOW   (MLKEM_BASE + 0x1000)
#define MLKEM_DK_WINDOW   (MLKEM_BASE + 0x2000)
#define MLKEM_CT_WINDOW   (MLKEM_BASE + 0x3000)

#define CTRL_START        0x1
#define STATUS_DONE       0x2

#define OP_KEYGEN         0
#define OP_ENCAPS         1
#define OP_DECAPS         2

#define PRE_EK            0
#define PRE_DK            1
#define PRE_C             2
#define RDB_EK            0
#define RDB_DK            2
#define RDB_CT            3

#define MAX_K             4
#define MAX_EK_BYTES      (384 * MAX_K + 32)
#define MAX_DK_BYTES      (768 * MAX_K + 96)
#define MAX_CT_BYTES      1568
#define SYM_BYTES         32
#define RUN_SW_DECAPS     1
#define DIAG_ITERS        100
#define TRACE_STAGES      0

#define REG_WR(addr, val) (*(volatile uint32_t *)(addr) = (uint32_t)(val))
#define REG_RD(addr)      (*(volatile uint32_t *)(addr))
#define PHYS_ADDR(addr)   ((uint32_t)(addr) & 0x1FFFFFFFu)
#define UNCACHED_ADDR(addr) ((PHYS_ADDR(addr)) | 0xA0000000u)

typedef int (*mlkem_keypair_derand_fn)(uint8_t *pk, uint8_t *sk,
                                       const uint8_t *coins);
typedef int (*mlkem_enc_derand_fn)(uint8_t *ct, uint8_t *ss,
                                   const uint8_t *pk, const uint8_t *msg);
typedef int (*mlkem_dec_fn)(uint8_t *ss, const uint8_t *ct,
                            const uint8_t *sk);

extern int swmlkem512_keypair_derand(uint8_t *pk, uint8_t *sk,
                                     const uint8_t *coins);
extern int swmlkem512_enc_derand(uint8_t *ct, uint8_t *ss,
                                 const uint8_t *pk, const uint8_t *msg);
extern int swmlkem512_dec(uint8_t *ss, const uint8_t *ct,
                          const uint8_t *sk);
extern int swmlkem768_keypair_derand(uint8_t *pk, uint8_t *sk,
                                     const uint8_t *coins);
extern int swmlkem768_enc_derand(uint8_t *ct, uint8_t *ss,
                                 const uint8_t *pk, const uint8_t *msg);
extern int swmlkem768_dec(uint8_t *ss, const uint8_t *ct,
                          const uint8_t *sk);
extern int swmlkem1024_keypair_derand(uint8_t *pk, uint8_t *sk,
                                      const uint8_t *coins);
extern int swmlkem1024_enc_derand(uint8_t *ct, uint8_t *ss,
                                  const uint8_t *pk, const uint8_t *msg);
extern int swmlkem1024_dec(uint8_t *ss, const uint8_t *ct,
                           const uint8_t *sk);

typedef struct {
    int k;
    const char *name;
    mlkem_keypair_derand_fn keypair_derand;
    mlkem_enc_derand_fn enc_derand;
    mlkem_dec_fn dec;
    const uint8_t *expect_ek32;
    const uint8_t *expect_ct32;
    const uint8_t *expect_ss;
} mlkem_level_t;

typedef struct {
    int ek_bytes;
    int dk_bytes;
    int ct_bytes;
} mlkem_sizes_t;

typedef struct {
    uint8_t ek[MAX_EK_BYTES] __attribute__((aligned(64)));
    uint8_t dk[MAX_DK_BYTES] __attribute__((aligned(64)));
    uint8_t ct[MAX_CT_BYTES] __attribute__((aligned(64)));
    uint8_t ss_enc[SYM_BYTES];
    uint8_t ss_dec[SYM_BYTES];
} mlkem_result_t;

typedef struct {
    unsigned long keygen;
    unsigned long encaps;
    unsigned long decaps;
} stage_cycles_t;

static mlkem_result_t sw_res;
static mlkem_result_t cpu_res;
static mlkem_result_t dma_res;

static const uint8_t EXPECT_512_EK32[SYM_BYTES] = {
    0x80, 0xb7, 0x01, 0x00, 0x83, 0x1d, 0x7e, 0xe1,
    0x68, 0xf4, 0x73, 0x92, 0x48, 0x31, 0x19, 0xf9,
    0x38, 0x63, 0xec, 0xb0, 0xbc, 0x60, 0x53, 0x12,
    0xd3, 0x57, 0x28, 0x87, 0x6b, 0x00, 0xf5, 0xf5
};

static const uint8_t EXPECT_512_CT32[SYM_BYTES] = {
    0x91, 0x7e, 0x59, 0xd1, 0xe1, 0x82, 0xa2, 0x96,
    0x48, 0xa6, 0x19, 0x55, 0xa1, 0x54, 0xb3, 0x04,
    0xa8, 0x42, 0xf4, 0x58, 0xeb, 0xa5, 0x93, 0xf3,
    0x06, 0x5f, 0xb7, 0xd7, 0x44, 0xfb, 0xa2, 0x4f
};

static const uint8_t EXPECT_512_SS[SYM_BYTES] = {
    0x2b, 0x5c, 0x52, 0xee, 0x72, 0x94, 0x63, 0x31,
    0x98, 0x3b, 0xa0, 0x50, 0xbe, 0x0f, 0x43, 0x50,
    0x55, 0xc0, 0x54, 0x79, 0x01, 0xe0, 0x35, 0x59,
    0xb3, 0x56, 0x51, 0x78, 0x89, 0xea, 0x27, 0xc5
};

static const uint8_t EXPECT_768_EK32[SYM_BYTES] = {
    0x01, 0xf6, 0x0a, 0xf1, 0xdc, 0x8e, 0x63, 0x60,
    0xae, 0x78, 0xb5, 0x9d, 0x4a, 0x50, 0x42, 0xeb,
    0x91, 0x45, 0xa2, 0x69, 0x04, 0x6d, 0x62, 0x36,
    0xb8, 0x30, 0x4f, 0x30, 0x5c, 0x2d, 0x9d, 0xcb
};

static const uint8_t EXPECT_768_CT32[SYM_BYTES] = {
    0x16, 0xa6, 0x1f, 0xf8, 0x47, 0x87, 0xfd, 0x4a,
    0x5f, 0x19, 0xca, 0x59, 0xb3, 0x65, 0x7d, 0xb3,
    0xa1, 0x06, 0xa7, 0x32, 0x9e, 0x2d, 0x62, 0x74,
    0x7a, 0x2e, 0xf8, 0x51, 0x49, 0x16, 0x31, 0x09
};

static const uint8_t EXPECT_768_SS[SYM_BYTES] = {
    0xb4, 0x08, 0xd5, 0xd1, 0x15, 0x71, 0x3f, 0x0a,
    0x93, 0x04, 0x7d, 0xbb, 0xea, 0x83, 0x2e, 0x43,
    0x40, 0x78, 0x76, 0x86, 0xd5, 0x9a, 0x9a, 0x2d,
    0x10, 0x6b, 0xd6, 0x62, 0xba, 0x0a, 0xa0, 0x35
};

static const uint8_t EXPECT_1024_EK32[SYM_BYTES] = {
    0xa8, 0x12, 0x2f, 0x37, 0x6b, 0x3f, 0x5d, 0x35,
    0x52, 0x63, 0xeb, 0xa5, 0x22, 0xc4, 0x39, 0x95,
    0x04, 0x4b, 0xec, 0xa7, 0x8b, 0x9b, 0xa9, 0x92,
    0x4b, 0x87, 0x03, 0x54, 0x90, 0xae, 0x78, 0xa9
};

static const uint8_t EXPECT_1024_CT32[SYM_BYTES] = {
    0x70, 0x7d, 0x18, 0xca, 0xbc, 0xf8, 0x96, 0x70,
    0xc8, 0x00, 0x03, 0xb4, 0x7d, 0x2b, 0x86, 0x78,
    0xad, 0x0d, 0xa4, 0xaa, 0xad, 0x78, 0x1a, 0x33,
    0x51, 0xe8, 0x2a, 0xc3, 0xe4, 0x47, 0xe7, 0x01
};

static const uint8_t EXPECT_1024_SS[SYM_BYTES] = {
    0x23, 0xf2, 0x11, 0xb8, 0x4a, 0x6e, 0xe2, 0x0c,
    0x8c, 0x29, 0xf6, 0xe5, 0x31, 0x4c, 0x91, 0xb4,
    0x14, 0xe9, 0x40, 0x51, 0x3d, 0x38, 0x0a, 0xdd,
    0x17, 0xbd, 0x72, 0x4a, 0xb3, 0xa1, 0x3a, 0x52
};

static const mlkem_level_t LEVELS[] = {
    // {2, "ML-KEM-512",  swmlkem512_keypair_derand,  swmlkem512_enc_derand,  swmlkem512_dec,
    //  EXPECT_512_EK32, EXPECT_512_CT32, EXPECT_512_SS},
    // {3, "ML-KEM-768",  swmlkem768_keypair_derand,  swmlkem768_enc_derand,  swmlkem768_dec,
    //  EXPECT_768_EK32, EXPECT_768_CT32, EXPECT_768_SS},
    {4, "ML-KEM-1024", swmlkem1024_keypair_derand, swmlkem1024_enc_derand, swmlkem1024_dec,
     EXPECT_1024_EK32, EXPECT_1024_CT32, EXPECT_1024_SS},
};

/* Deterministic ML-KEM KAT count=0 inputs. */
static const uint8_t SEED_D[SYM_BYTES] = {
    0x6d, 0xbb, 0xc4, 0x37, 0x51, 0x36, 0xdf, 0x3b,
    0x07, 0xf7, 0xc7, 0x0e, 0x63, 0x9e, 0x22, 0x3e,
    0x17, 0x7e, 0x7f, 0xd5, 0x3b, 0x16, 0x1b, 0x3f,
    0x4d, 0x57, 0x79, 0x17, 0x94, 0xf1, 0x26, 0x24
};

static const uint8_t SEED_Z[SYM_BYTES] = {
    0xf6, 0x96, 0x48, 0x40, 0x48, 0xec, 0x21, 0xf9,
    0x6c, 0xf5, 0x0a, 0x56, 0xd0, 0x75, 0x9c, 0x44,
    0x8f, 0x37, 0x79, 0x75, 0x2f, 0x03, 0x83, 0xd3,
    0x74, 0x49, 0x69, 0x06, 0x94, 0xcf, 0x7a, 0x68
};

static const uint8_t MSG_M[SYM_BYTES] = {
    0x20, 0xa7, 0xb7, 0xe1, 0x0f, 0x70, 0x49, 0x6c,
    0xc3, 0x82, 0x20, 0xb9, 0x44, 0xde, 0xf6, 0x99,
    0xbf, 0x14, 0xd1, 0x4e, 0x55, 0xcf, 0x4c, 0x90,
    0xa1, 0x2c, 0x1b, 0x33, 0xfc, 0x80, 0xff, 0xff
};

static mlkem_sizes_t sizes_for_k(int k)
{
    mlkem_sizes_t s;
    s.ek_bytes = 384 * k + 32;
    s.dk_bytes = 768 * k + 96;
    s.ct_bytes = (k == 4) ? 1568 : (32 * 10 * k + 32 * 4);
    return s;
}

static void clear_done(void)
{
    REG_WR(MLKEM_STATUS, STATUS_DONE);
}

static void wait_done(void)
{
    while ((REG_RD(MLKEM_STATUS) & STATUS_DONE) == 0)
        ;
}

static void run_op(int k, uint32_t op)
{
    REG_WR(MLKEM_CONFIG, ((op & 0x3) << 4) | (k & 0x7));
    clear_done();
    REG_WR(MLKEM_CTRL, CTRL_START);
    wait_done();
}

static void write_block32(uint32_t base, const uint8_t *bytes)
{
    int w;
    for (w = 0; w < 8; w++) {
        uint32_t v = (uint32_t)bytes[w * 4 + 0]
                   | ((uint32_t)bytes[w * 4 + 1] << 8)
                   | ((uint32_t)bytes[w * 4 + 2] << 16)
                   | ((uint32_t)bytes[w * 4 + 3] << 24);
        REG_WR(base + w * 4, v);
    }
}

static void read_block32(uint32_t base, uint8_t *bytes)
{
    int w;
    for (w = 0; w < 8; w++) {
        uint32_t v = REG_RD(base + w * 4);
        bytes[w * 4 + 0] = (uint8_t)(v & 0xFF);
        bytes[w * 4 + 1] = (uint8_t)((v >> 8) & 0xFF);
        bytes[w * 4 + 2] = (uint8_t)((v >> 16) & 0xFF);
        bytes[w * 4 + 3] = (uint8_t)((v >> 24) & 0xFF);
    }
}

static void preload(uint32_t sel, const uint8_t *buf, int n)
{
    int i;
    REG_WR(MLKEM_PRE_SEL, sel);
    REG_WR(MLKEM_PRE_ADR, 0);
    for (i = 0; i < n; i++)
        REG_WR(MLKEM_PRE_DAT, buf[i]);
}

static void readback(uint32_t sel, uint8_t *buf, int n)
{
    int i;
    REG_WR(MLKEM_RDB_SEL, sel);
    for (i = 0; i < n; i++) {
        REG_WR(MLKEM_RDB_ADR, i);
        buf[i] = (uint8_t)(REG_RD(MLKEM_RDB_DAT) & 0xFF);
    }
}

static volatile uint8_t *uncached_u8(uint8_t *buf)
{
    return (volatile uint8_t *)UNCACHED_ADDR(buf);
}

static void clear_uc(volatile uint8_t *b, int n)
{
    int i;
    for (i = 0; i < n; i++)
        b[i] = 0;
}

static void copy_from_uc(uint8_t *dst, volatile uint8_t *src, int n)
{
    int i;
    for (i = 0; i < n; i++)
        dst[i] = src[i];
}

static void dma_to_mlkem(uint32_t dst_window, volatile uint8_t *src, int n)
{
    dma_start_transfer(0, PHYS_ADDR(src), PHYS_ADDR(dst_window), (uint32_t)n, 0);
    dma_wait_polling(0);
}

static void dma_from_mlkem(volatile uint8_t *dst, uint32_t src_window, int n)
{
    dma_start_transfer(0, PHYS_ADDR(src_window), PHYS_ADDR(dst), (uint32_t)n, 0);
    dma_wait_polling(0);
}

static void clear_result_cpu(mlkem_result_t *res)
{
    memset(res, 0, sizeof(*res));
}

static void clear_result_dma(mlkem_result_t *res)
{
    volatile uint8_t *p = (volatile uint8_t *)UNCACHED_ADDR(res);
    clear_uc(p, sizeof(*res));
}

static unsigned long run_roundtrip_sw(const mlkem_level_t *level,
                                      const uint8_t *seed_d,
                                      const uint8_t *seed_z,
                                      const uint8_t *msg,
                                      mlkem_result_t *res,
                                      int *status,
                                      stage_cycles_t *stages)
{
    unsigned long t0, t1;
    uint8_t coins[2 * SYM_BYTES];
    int rc;

    clear_result_cpu(res);
    memcpy(coins, seed_d, SYM_BYTES);
    memcpy(coins + SYM_BYTES, seed_z, SYM_BYTES);

    if (TRACE_STAGES) printf("SW_KEYGEN_BEGIN\n");
    t0 = get_clock_count();
    rc = level->keypair_derand(res->ek, res->dk, coins);
    t1 = get_clock_count();
    stages->keygen = t1 - t0;
    if (TRACE_STAGES) printf("SW_KEYGEN_END rc=%d\n", rc);

    if (rc == 0) {
        if (TRACE_STAGES) printf("SW_ENCAPS_BEGIN\n");
        t0 = get_clock_count();
        rc = level->enc_derand(res->ct, res->ss_enc, res->ek, msg);
        t1 = get_clock_count();
        stages->encaps = t1 - t0;
        if (TRACE_STAGES) printf("SW_ENCAPS_END rc=%d\n", rc);
    } else {
        stages->encaps = 0;
    }

    if (rc == 0 && RUN_SW_DECAPS) {
        if (TRACE_STAGES) printf("SW_DECAPS_BEGIN\n");
        t0 = get_clock_count();
        rc = level->dec(res->ss_dec, res->ct, res->dk);
        t1 = get_clock_count();
        stages->decaps = t1 - t0;
        if (TRACE_STAGES) printf("SW_DECAPS_END rc=%d\n", rc);
    } else {
        if (rc == 0) {
            memcpy(res->ss_dec, res->ss_enc, SYM_BYTES);
            if (TRACE_STAGES) printf("SW_DECAPS_SKIPPED reason=target_sw_dec_hangs\n");
        }
        stages->decaps = 0;
    }

    *status = rc;
    return stages->keygen + stages->encaps + stages->decaps;
}

static unsigned long run_roundtrip_cpu(const mlkem_level_t *level,
                                       mlkem_sizes_t sizes,
                                       const uint8_t *seed_d,
                                       const uint8_t *seed_z,
                                       const uint8_t *msg,
                                       mlkem_result_t *res,
                                       stage_cycles_t *stages)
{
    unsigned long t0, t1;
    int k = level->k;

    clear_result_cpu(res);
    if (TRACE_STAGES) printf("CPU_PATH_BEGIN\n");
    t0 = get_clock_count();

    if (TRACE_STAGES) printf("CPU_KEYGEN_BEGIN\n");
    write_block32(MLKEM_D, seed_d);
    write_block32(MLKEM_Z, seed_z);
    run_op(k, OP_KEYGEN);
    readback(RDB_EK, res->ek, sizes.ek_bytes);
    readback(RDB_DK, res->dk, sizes.dk_bytes);

    t1 = get_clock_count();
    stages->keygen = t1 - t0;
    if (TRACE_STAGES) printf("CPU_KEYGEN_END\n");

    t0 = get_clock_count();
    if (TRACE_STAGES) printf("CPU_ENCAPS_BEGIN\n");
    preload(PRE_EK, res->ek, sizes.ek_bytes);
    write_block32(MLKEM_MSG, msg);
    run_op(k, OP_ENCAPS);
    read_block32(MLKEM_SS, res->ss_enc);
    readback(RDB_CT, res->ct, sizes.ct_bytes);

    t1 = get_clock_count();
    stages->encaps = t1 - t0;
    if (TRACE_STAGES) printf("CPU_ENCAPS_END\n");

    t0 = get_clock_count();
    if (TRACE_STAGES) printf("CPU_DECAPS_BEGIN\n");
    preload(PRE_DK, res->dk, sizes.dk_bytes);
    preload(PRE_C,  res->ct, sizes.ct_bytes);
    run_op(k, OP_DECAPS);
    read_block32(MLKEM_SS, res->ss_dec);

    t1 = get_clock_count();
    stages->decaps = t1 - t0;
    if (TRACE_STAGES) printf("CPU_DECAPS_END\n");
    if (TRACE_STAGES) printf("CPU_PATH_END\n");

    return stages->keygen + stages->encaps + stages->decaps;
}

static unsigned long run_roundtrip_dma(const mlkem_level_t *level,
                                       mlkem_sizes_t sizes,
                                       const uint8_t *seed_d,
                                       const uint8_t *seed_z,
                                       const uint8_t *msg,
                                       mlkem_result_t *res,
                                       stage_cycles_t *stages)
{
    unsigned long t0, t1;
    int k = level->k;
    volatile uint8_t *ek_uc = uncached_u8(res->ek);
    volatile uint8_t *dk_uc = uncached_u8(res->dk);
    volatile uint8_t *ct_uc = uncached_u8(res->ct);

    clear_result_dma(res);
    if (TRACE_STAGES) printf("DMA_PATH_BEGIN\n");
    t0 = get_clock_count();

    if (TRACE_STAGES) printf("DMA_KEYGEN_BEGIN\n");
    write_block32(MLKEM_D, seed_d);
    write_block32(MLKEM_Z, seed_z);
    run_op(k, OP_KEYGEN);
    dma_from_mlkem(ek_uc, MLKEM_EK_WINDOW, sizes.ek_bytes);
    dma_from_mlkem(dk_uc, MLKEM_DK_WINDOW, sizes.dk_bytes);

    t1 = get_clock_count();
    stages->keygen = t1 - t0;
    if (TRACE_STAGES) printf("DMA_KEYGEN_END\n");

    t0 = get_clock_count();
    if (TRACE_STAGES) printf("DMA_ENCAPS_BEGIN\n");
    dma_to_mlkem(MLKEM_EK_WINDOW, ek_uc, sizes.ek_bytes);
    write_block32(MLKEM_MSG, msg);
    run_op(k, OP_ENCAPS);
    read_block32(MLKEM_SS, res->ss_enc);
    dma_from_mlkem(ct_uc, MLKEM_CT_WINDOW, sizes.ct_bytes);

    t1 = get_clock_count();
    stages->encaps = t1 - t0;
    if (TRACE_STAGES) printf("DMA_ENCAPS_END\n");

    t0 = get_clock_count();
    if (TRACE_STAGES) printf("DMA_DECAPS_BEGIN\n");
    dma_to_mlkem(MLKEM_DK_WINDOW, dk_uc, sizes.dk_bytes);
    dma_to_mlkem(MLKEM_CT_WINDOW, ct_uc, sizes.ct_bytes);
    run_op(k, OP_DECAPS);
    read_block32(MLKEM_SS, res->ss_dec);

    t1 = get_clock_count();
    stages->decaps = t1 - t0;
    if (TRACE_STAGES) printf("DMA_DECAPS_END\n");

    copy_from_uc(res->ek, ek_uc, sizes.ek_bytes);
    copy_from_uc(res->dk, dk_uc, sizes.dk_bytes);
    copy_from_uc(res->ct, ct_uc, sizes.ct_bytes);
    if (TRACE_STAGES) printf("DMA_PATH_END\n");

    return stages->keygen + stages->encaps + stages->decaps;
}

static int verify_result(const mlkem_result_t *res)
{
    int i, mismatch = 0;
    for (i = 0; i < SYM_BYTES; i++)
        if (res->ss_enc[i] != res->ss_dec[i])
            mismatch++;
    return mismatch;
}

static int compare_results(const mlkem_result_t *a,
                           const mlkem_result_t *b,
                           mlkem_sizes_t sizes)
{
    return memcmp(a->ek, b->ek, sizes.ek_bytes) ||
           memcmp(a->dk, b->dk, sizes.dk_bytes) ||
           memcmp(a->ct, b->ct, sizes.ct_bytes) ||
           memcmp(a->ss_enc, b->ss_enc, SYM_BYTES) ||
           memcmp(a->ss_dec, b->ss_dec, SYM_BYTES);
}

static int golden_mismatches(const mlkem_level_t *level,
                             const mlkem_result_t *res)
{
    int mismatches = 0;

    if (memcmp(res->ek, level->expect_ek32, SYM_BYTES) != 0)
        mismatches++;
    if (memcmp(res->ct, level->expect_ct32, SYM_BYTES) != 0)
        mismatches++;
    if (memcmp(res->ss_enc, level->expect_ss, SYM_BYTES) != 0)
        mismatches++;
    if (memcmp(res->ss_dec, level->expect_ss, SYM_BYTES) != 0)
        mismatches++;

    return mismatches;
}

static void print_hex(const char *label, const uint8_t *b, int n)
{
    int i;
    printf("%s", label);
    for (i = 0; i < n; i++)
        printf("%02x", b[i]);
    printf("\n");
}

static unsigned long cycles_to_us(unsigned long cycles)
{
    return cycles / (CORE_CLOCKS_PER_SEC / USEC_PER_SEC);
}

static unsigned long speedup_x100(unsigned long base, unsigned long accel)
{
    if (accel == 0)
        return 0;
    return (base * 100u) / accel;
}

static void print_stage_perf(int lvl, const char *stage,
                             unsigned long sw,
                             unsigned long cpu,
                             unsigned long dma)
{
    printf("MLKEM%d_%s_SW_CYCLES=%lu\n", lvl, stage, sw);
    printf("MLKEM%d_%s_CPU_CYCLES=%lu\n", lvl, stage, cpu);
    printf("MLKEM%d_%s_DMA_CYCLES=%lu\n", lvl, stage, dma);
    printf("MLKEM%d_%s_SW_US=%lu\n", lvl, stage, cycles_to_us(sw));
    printf("MLKEM%d_%s_CPU_US=%lu\n", lvl, stage, cycles_to_us(cpu));
    printf("MLKEM%d_%s_DMA_US=%lu\n", lvl, stage, cycles_to_us(dma));
    printf("MLKEM%d_%s_CPU_SPEEDUP_X100=%lu\n", lvl, stage,
           speedup_x100(sw, cpu));
    printf("MLKEM%d_%s_DMA_SPEEDUP_X100=%lu\n", lvl, stage,
           speedup_x100(sw, dma));
    printf("MLKEM%d_%s_DMA_VS_CPU_X100=%lu\n", lvl, stage,
           speedup_x100(cpu, dma));
}

static void print_diag_failure(const char *tag, const mlkem_result_t *res)
{
    printf("%s_FIRST_FAIL_DETAIL_BEGIN\n", tag);
    print_hex("EK32=", res->ek, 32);
    print_hex("CT32=", res->ct, 32);
    print_hex("SS_ENC=", res->ss_enc, 32);
    print_hex("SS_DEC=", res->ss_dec, 32);
    printf("%s_FIRST_FAIL_DETAIL_END\n", tag);
}

static int run_sw_diag(const mlkem_level_t *level)
{
    int lvl = level->k * 256;
    int i;
    int fail_count = 0;
    int first_fail = -1;
    int sw_status = 0;
    unsigned long total_cycles = 0;
    stage_cycles_t stages;

    printf("MLKEM%d_SW_ONLY_BEGIN iters=%d\n", lvl, DIAG_ITERS);
    for (i = 0; i < DIAG_ITERS; i++) {
        int self_mismatch;
        int golden;
        unsigned long cycles;

        cycles = run_roundtrip_sw(level, SEED_D, SEED_Z, MSG_M,
                                  &sw_res, &sw_status, &stages);
        total_cycles += cycles;
        self_mismatch = verify_result(&sw_res);
        golden = golden_mismatches(level, &sw_res);

        if (sw_status != 0 || self_mismatch != 0 || golden != 0) {
            fail_count++;
            if (first_fail < 0) {
                first_fail = i;
                printf("MLKEM%d_SW_ONLY_FIRST_FAIL iter=%d status=%d self_mismatch=%d golden_mismatch=%d\n",
                       lvl, i, sw_status, self_mismatch, golden);
                print_diag_failure("SW_ONLY", &sw_res);
            }
        }
    }

    printf("MLKEM%d_SW_ONLY_ITERS=%d\n", lvl, DIAG_ITERS);
    printf("MLKEM%d_SW_ONLY_FAILS=%d\n", lvl, fail_count);
    printf("MLKEM%d_SW_ONLY_AVG_CYCLES=%lu\n", lvl,
           total_cycles / DIAG_ITERS);
    printf("MLKEM%d_SW_ONLY_AVG_US=%lu\n", lvl,
           cycles_to_us(total_cycles / DIAG_ITERS));
    printf("MLKEM%d_SW_ONLY_RESULT=%s\n", lvl,
           fail_count == 0 ? "PASS" : "FAIL");
    printf("MLKEM%d_SW_ONLY_END\n", lvl);

    return fail_count;
}

static int run_cpu_diag(const mlkem_level_t *level)
{
    mlkem_sizes_t sizes = sizes_for_k(level->k);
    int lvl = level->k * 256;
    int i;
    int fail_count = 0;
    int first_fail = -1;
    unsigned long total_cycles = 0;
    stage_cycles_t stages;

    printf("MLKEM%d_CPU_ONLY_BEGIN iters=%d\n", lvl, DIAG_ITERS);
    for (i = 0; i < DIAG_ITERS; i++) {
        int self_mismatch;
        int golden;
        unsigned long cycles;

        cycles = run_roundtrip_cpu(level, sizes, SEED_D, SEED_Z, MSG_M,
                                   &cpu_res, &stages);
        total_cycles += cycles;
        self_mismatch = verify_result(&cpu_res);
        golden = golden_mismatches(level, &cpu_res);

        if (self_mismatch != 0 || golden != 0) {
            fail_count++;
            if (first_fail < 0) {
                first_fail = i;
                printf("MLKEM%d_CPU_ONLY_FIRST_FAIL iter=%d self_mismatch=%d golden_mismatch=%d\n",
                       lvl, i, self_mismatch, golden);
                print_diag_failure("CPU_ONLY", &cpu_res);
            }
        }
    }

    printf("MLKEM%d_CPU_ONLY_ITERS=%d\n", lvl, DIAG_ITERS);
    printf("MLKEM%d_CPU_ONLY_FAILS=%d\n", lvl, fail_count);
    printf("MLKEM%d_CPU_ONLY_AVG_CYCLES=%lu\n", lvl,
           total_cycles / DIAG_ITERS);
    printf("MLKEM%d_CPU_ONLY_AVG_US=%lu\n", lvl,
           cycles_to_us(total_cycles / DIAG_ITERS));
    printf("MLKEM%d_CPU_ONLY_RESULT=%s\n", lvl,
           fail_count == 0 ? "PASS" : "FAIL");
    printf("MLKEM%d_CPU_ONLY_END\n", lvl);

    return fail_count;
}

static int run_dma_diag(const mlkem_level_t *level)
{
    mlkem_sizes_t sizes = sizes_for_k(level->k);
    int lvl = level->k * 256;
    int i;
    int fail_count = 0;
    int first_fail = -1;
    unsigned long total_cycles = 0;
    stage_cycles_t stages;

    printf("MLKEM%d_DMA_ONLY_BEGIN iters=%d\n", lvl, DIAG_ITERS);
    for (i = 0; i < DIAG_ITERS; i++) {
        int self_mismatch;
        int golden;
        unsigned long cycles;

        cycles = run_roundtrip_dma(level, sizes, SEED_D, SEED_Z, MSG_M,
                                   &dma_res, &stages);
        total_cycles += cycles;
        self_mismatch = verify_result(&dma_res);
        golden = golden_mismatches(level, &dma_res);

        if (self_mismatch != 0 || golden != 0) {
            fail_count++;
            if (first_fail < 0) {
                first_fail = i;
                printf("MLKEM%d_DMA_ONLY_FIRST_FAIL iter=%d self_mismatch=%d golden_mismatch=%d\n",
                       lvl, i, self_mismatch, golden);
                print_diag_failure("DMA_ONLY", &dma_res);
            }
        }
    }

    printf("MLKEM%d_DMA_ONLY_ITERS=%d\n", lvl, DIAG_ITERS);
    printf("MLKEM%d_DMA_ONLY_FAILS=%d\n", lvl, fail_count);
    printf("MLKEM%d_DMA_ONLY_AVG_CYCLES=%lu\n", lvl,
           total_cycles / DIAG_ITERS);
    printf("MLKEM%d_DMA_ONLY_AVG_US=%lu\n", lvl,
           cycles_to_us(total_cycles / DIAG_ITERS));
    printf("MLKEM%d_DMA_ONLY_RESULT=%s\n", lvl,
           fail_count == 0 ? "PASS" : "FAIL");
    printf("MLKEM%d_DMA_ONLY_END\n", lvl);

    return fail_count;
}

static int run_level_diag(const mlkem_level_t *level)
{
    mlkem_sizes_t sizes = sizes_for_k(level->k);
    int lvl = level->k * 256;
    int fails = 0;

    printf("\n============================================================\n");
    printf("%s isolated diagnostic (k=%d)\n", level->name, level->k);
    printf("EK=%d bytes, DK=%d bytes, CT=%d bytes\n",
           sizes.ek_bytes, sizes.dk_bytes, sizes.ct_bytes);
    print_hex("EXPECT_EK32=", level->expect_ek32, 32);
    print_hex("EXPECT_CT32=", level->expect_ct32, 32);
    print_hex("EXPECT_SS=", level->expect_ss, 32);

    fails += run_sw_diag(level);
    fails += run_cpu_diag(level);
    fails += run_dma_diag(level);

    printf("MLKEM%d_ISOLATED_FAILS=%d\n", lvl, fails);
    printf("MLKEM%d_RESULT=%s\n", lvl, fails == 0 ? "PASS" : "FAIL");

    return fails;
}

int main(void)
{
    int i;
    int failures = 0;

    printf("MLKEM_START\n");
    printf("MLKEM_INPUTS=STATIC_KAT_COUNT_0\n");
    printf("MLKEM_LEVELS=k2,k3,k4\n");
    printf("MLKEM_DIAG_MODE=isolated_sw_cpu_dma\n");
    printf("MLKEM_DIAG_ITERS=%d\n", DIAG_ITERS);

    for (i = 0; i < (int)(sizeof(LEVELS) / sizeof(LEVELS[0])); i++)
        failures += run_level_diag(&LEVELS[i]);

    if (failures == 0)
        printf("MLKEM_RESULT=PASS\n");
    else
        printf("MLKEM_RESULT=FAIL failures=%d\n", failures);

    printf("MLKEM_DONE\n");
    return failures == 0 ? 0 : 1;
}