import DilithiumModule from "./wasm.cjs";
import {Buffer} from "buffer/index.js";

async function dilithiumInit({randomBytes}) {
    let mod = await DilithiumModule();

    if (!randomBytes) {
        randomBytes = () => {
            throw new Error("Random generator implementation not provided but required.");
        };
    }

    global.dilithium_randombytes = (function (out, outlen) {
        let buf = this.HEAPU8.subarray(out, out + outlen);
        let randomBytesBuffer = randomBytes(outlen);
        buf.set(new Uint8Array(randomBytesBuffer), 0);
    }).bind(mod);

    return mod;
}

async function dilithiumGenKeyPair({randomBytes}) {
    const mod = await dilithiumInit({randomBytes});

    if (mod._dilithium_gen_key_pair() === 0) {
        let pk_ptr = mod._get_pk_buffer();
        let pk_len = mod._get_pk_length();
        const pk = new Uint8Array(mod.HEAP8.buffer, pk_ptr, pk_len);

        let sk_ptr = mod._get_sk_buffer();
        let sk_len = mod._get_sk_length();
        const sk = new Uint8Array(mod.HEAP8.buffer, sk_ptr, sk_len);

        return { publicKey: Buffer.from(pk), secretKey: Buffer.from(sk) };
    } else {
        throw new Error('Failed to generate key pair.');
    }
}

async function dilithiumSign({secretKey, challenge}) {
    const mod = await dilithiumInit({});

    if (challenge.length !== 32) {
        throw new Error('Challenge buffer must be 32 bytes long.');
    }

    let m_ptr = mod._get_m_buffer();
    let m_len = mod._get_m_length();
    const m = new Uint8Array(mod.HEAP8.buffer, m_ptr, m_len);

    let sk_ptr = mod._get_sk_buffer();
    let sk_len = mod._get_sk_length();
    const sk = new Uint8Array(mod.HEAP8.buffer, sk_ptr, sk_len);

    let sig_ptr = mod._get_sig_buffer();
    let sig_len = mod._get_sig_length();
    const sig = new Uint8Array(mod.HEAP8.buffer, sig_ptr, sig_len);

    sk.set(Buffer.from(secretKey, 'hex'), 0);
    m.set(Buffer.from(challenge, 'hex'), 0);

    if (mod._dilithium_sign() === 0) {
        return Buffer.from(sig);
    } else {
        throw new Error('Failed to sign message.');
    }
}

async function dilithiumVerifySig({publicKey, challenge, signature}) {
    const mod = await dilithiumInit({});

    if (challenge.length !== 32) {
        throw new Error('Challenge buffer must be 32 bytes long.');
    }

    let m_ptr = mod._get_m_buffer();
    let m_len = mod._get_m_length();
    const m = new Uint8Array(mod.HEAP8.buffer, m_ptr, m_len);

    let sig_ptr = mod._get_sig_buffer();
    let sig_len = mod._get_sig_length();
    const sig = new Uint8Array(mod.HEAP8.buffer, sig_ptr, sig_len);

    let pk_ptr = mod._get_pk_buffer();
    let pk_len = mod._get_pk_length();
    const pk = new Uint8Array(mod.HEAP8.buffer, pk_ptr, pk_len);

    pk.set(publicKey, 0);
    sig.set(signature, 0);
    m.set(challenge, 0);

    return mod._dilithium_verify_signature() === 0;
}

export { dilithiumGenKeyPair, dilithiumSign, dilithiumVerifySig };