from Jade require "crypto_kem/kyber_kyber768_avx2/kyber_kyber768_avx2.jazz"

param int MLKEM_SK_SIZE = 2400;

inline fn extract_nth_ptr(reg u64 sync_mem shared_mem, inline int i) -> reg u64 {
    reg u64 p;

    p = (u64)[sync_mem + (16 + i * 8)];
    p += shared_mem;

    return p;
}

inline fn extract_ipc_id(reg u64 sync_mem) -> reg u64 {
    reg u64 id;

    id = (u64)[sync_mem + 8];

    return id;
}

inline fn generate_keypair(reg u64 shared_mem sync_mem private_mem, stack u64 key_id) -> stack u64 {
    inline int i;
    stack u8[2*KYBER_SYMBYTES] coins;
    reg u64 skptr idptr pkptr flag num_bytes;

    while {
        flag = 0;
        coins, num_bytes = #randombytes(coins, flag);
    } (num_bytes != 2*KYBER_SYMBYTES)

    skptr = key_id * MLKEM_SK_SIZE;
    skptr += private_mem;

    for i=0 to (2*KYBER_SYMBYTES)/8  {
        (u64)[skptr + i * 8] = coins[u64 i];
    }

    idptr = extract_nth_ptr(sync_mem, shared_mem, 0);
    pkptr = extract_nth_ptr(sync_mem, shared_mem, 1);

    jade_kem_kyber_kyber768_amd64_avx2_keypair_derand(pkptr, skptr, skptr);

    (u64)[idptr] = key_id;

    key_id += 1;

    return key_id;
}

inline fn encapsulate(reg u64 shared_mem sync_mem private_mem) {
    inline int i;
    reg u64 ssptr ctptr pkptr flag num_bytes;
    stack u8[KYBER_SYMBYTES] coins;

    while {
        flag = 0;
        coins, num_bytes = #randombytes(coins, flag);
    } (num_bytes != KYBER_SYMBYTES)

    for i=0 to KYBER_SYMBYTES/8 {
        (u64)[ssptr + i * 8] = coins[u64 i];
    }

    ctptr = extract_nth_ptr(sync_mem, shared_mem, 0);
    ssptr = extract_nth_ptr(sync_mem, shared_mem, 1);
    pkptr = extract_nth_ptr(sync_mem, shared_mem, 2);

    jade_kem_kyber_kyber768_amd64_avx2_enc_derand(ctptr, ssptr, pkptr, ssptr);
}

export fn agent_start(reg u64 shared_mem sync_mem private_mem_fd) {
    stack u64 key_id;
    reg u64 private_mem;
    // Futex varibles
    reg u64 futex_op val timeout uaddr2 val3 woken_up;
    // Mmap variables
    reg u64 addr len prot flag off;

    () = #spill(shared_mem, sync_mem);

    private_mem_fd = private_mem_fd;
    addr = 0;
    len = 3200;
    prot = 3; // PROT_READ | PROT_WRITE
    flag = 1; // MAP_SHARED;
    off = 0;
    private_mem = #mmap(addr, len, prot, flag, private_mem_fd, off);

    () = #spill(private_mem);

    key_id = 0;

     while (true) {

        while {
            () = #unspill(sync_mem);

            futex_op = 0;
            val = 0;
            timeout = 0;
            uaddr2 = 0;
            val3 = 0;
            woken_up = #futex(sync_mem, futex_op, val, timeout, uaddr2, val3);
        } (woken_up != 0)

        () = #unspill(shared_mem, sync_mem, private_mem);
        
        reg u64 id;
        id = extract_ipc_id(sync_mem);

        if (id == 0) {
            key_id = generate_keypair(shared_mem, sync_mem, private_mem, key_id);
        } else {
            if (id == 1) {
                encapsulate(shared_mem, sync_mem, private_mem);
            } else {
                if (id == 2) {
                }
            }
        }

        while {
            () = #unspill(sync_mem);
            futex_op = 1;
            val = 1;
            timeout = 0;
            uaddr2 = 0;
            val3 = 0;
            woken_up = #futex(sync_mem, futex_op, val, timeout, uaddr2, val3);
        } (woken_up != 1)
    }
}