from Jade require "crypto_scalarmult/curve25519/amd64/mulx/scalarmult.jazz"

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_key(reg u64 shared_mem sync_mem private_mem, stack u64 key_id) -> stack u64 {
    inline int i;
    stack u8[32] key;
    reg u64 addr idptr pkptr;

    key = #randombytes(key);
    key[0] &= 248;
    key[31] &= 127;
    key[31] |= 64;

    key_id += 1; 

    addr = key_id * 32;
    addr += private_mem;

    for i=0 to 4 {
        (u64)[private_mem + i * 8] = key[u64 i];
    }

    idptr = (u64)[sync_mem + 16];
    idptr += shared_mem;

    (u64)[idptr] = key_id;

    pkptr = (u64)[sync_mem + 24];
    pkptr += shared_mem;

    jade_scalarmult_curve25519_amd64_mulx_base(pkptr, private_mem);

    return key_id;
}

inline fn x25519(reg u64 shared_mem sync_mem private_mem) {
    reg u64 outptr idptr pkptr keyptr;
    reg u64 keyid;

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

    keyid = (u64)[idptr];
    keyid *= 8;
    keyptr = private_mem + keyid;

    jade_scalarmult_curve25519_amd64_mulx(outptr, keyptr, pkptr);
}

export fn agent_start(reg u64 shared_mem sync_mem private_mem) {
    stack u8[8] unused;
    stack u64 key_id;

    key_id = 0;

    () = #spill(shared_mem, sync_mem, private_mem);

     while (true) {
        unused = unused;
        () = #unspill(sync_mem);
        unused = #read(unused, sync_mem);

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

        if (id == 0) {
            key_id = generate_key(shared_mem, sync_mem, private_mem, key_id);
        } else {
            if (id == 1) {
                x25519(shared_mem, sync_mem, private_mem);
            }
        }

        () = #unspill(sync_mem);

        unused = #write(unused, sync_mem);
    }
}