import blake2bReference from 'blake2b';
import _sodium from 'libsodium-wrappers';
import tweetNaCl from 'tweetnacl';

import {
    argon2id,
    blake2bMac256,
    hmacSha256,
    scrypt,
    sha256,
    x25519DerivePublicKey,
    x25519HSalsa20DeriveSharedSecret,
    xChaCha20Poly1305Decrypt,
    xChaCha20Poly1305Encrypt,
    xSalsa20Poly1305Decrypt,
    xSalsa20Poly1305Encrypt,
} from '../../build/wasm/nodejs/libthreema.js';

import {assert, evaluateTestResults, init, runTest} from './utils.js';

// Test SHA-256 against test vectors
function testSha256() {
    // Test vectors taken from https://www.di-mgt.com.au/sha_testvectors.html
    const testVectors = [
        {
            data: '',
            expected:
                'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855',
        },
        {
            data: 'abc',
            expected:
                'ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad',
        },
        {
            data: 'a'.repeat(1_000_000),
            expected:
                'cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0',
        },
    ];

    for (const testVector of testVectors) {
        const computed = Buffer.from(
            sha256(Buffer.from(testVector.data, 'utf8')),
        ).toString('hex');

        assert(
            testVector.expected === computed,
            `Computed SHA-256 should match expected one. Input: ${testVector.data}`,
        );
    }
}

// Test HMAC-SHA-256 against test vectors
function testHmacSha256() {
    // Test vectors taken from https://datatracker.ietf.org/doc/html/rfc4231#section-4
    const testVectors = [
        {
            key: '0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b',
            data: 'Hi There',
            expected:
                'b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7',
        },
        {
            key: '4a656665',
            data: 'what do ya want for nothing?',
            expected:
                '5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843',
        },
        {
            key: 'a'.repeat(262),
            data: 'Test Using Larger Than Block-Size Key - Hash Key First',
            expected:
                '60e431591ee0b67f0d8a26aacbf5b77f8e0bc6213728c5140546040f0ee37f54',
        },
    ];

    for (const testVector of testVectors) {
        const computed = Buffer.from(
            hmacSha256(
                Buffer.from(testVector.key, 'hex'),
                Buffer.from(testVector.data, 'utf8'),
            ),
        ).toString('hex');

        assert(
            testVector.expected === computed,
            `Computed HMAC-SHA-256 should match expected one. Input: ${testVector.data}, Key: ${testVector.key}`,
        );
    }
}

// Test Argon2id against test vectors
function testArgon2id() {
    // Test vector taken from
    // https://github.com/RustCrypto/password-hashes/blob/master/argon2/tests/kat.rs

    const argon2idParameters = {
        memoryCost: 32,
        timeCost: 3,
        parallelism: 4,
        outputLength: 32,
    };

    const computed = Buffer.from(
        argon2id(
            Buffer.from('01'.repeat(32), 'hex'),
            Buffer.from('02'.repeat(16), 'hex'),
            argon2idParameters,
        ),
    ).toString('hex');

    assert(
        '03aab965c12001c9d7d0d2de33192c0494b684bb148196d73c1df1acaf6d0c2e' ===
            computed,
        'Computed Argon2id hash should match expected one. ',
    );
}

// Test Scrypt against test vectors
function testScrypt() {
    // Test vector taken from the original paper (https://www.tarsnap.com/scrypt/scrypt.pdf)

    const scryptParameters = {
        logMemoryCost: 10,
        blockSize: 8,
        parallelism: 16,
        outputLength: 64,
    };

    const computed = Buffer.from(
        scrypt(
            Buffer.from('password', 'utf8'),
            Buffer.from('NaCl', 'utf8'),
            scryptParameters,
        ),
    ).toString('hex');

    assert(
        'fdbabe1c9d3472007856e7190d01e9fe7c6ad7cbc8237830e77376634b3731622eaf30d92e22a3886ff109279d9830dac727afb94a83ee6d8360cbdfa2cc0640' ===
            computed,
        'Computed Scrypt hash should match expected one',
    );
}

// Test BLake2b against test vectors
function testBlake2b() {
    const outputLength = 32;

    const testVectors = [
        {key: undefined, salt: undefined, personal: undefined},
        {key: undefined, salt: '01'.repeat(16), personal: undefined},
        {key: undefined, salt: '02'.repeat(16), personal: 'cd'.repeat(16)},
        {key: 'hello my fancy key', salt: undefined, personal: undefined},
        {
            key: 'hello dear personal',
            salt: '03'.repeat(16),
            personal: 'ab'.repeat(16),
        },
        {
            key: 'hello my dear salt',
            salt: '04'.repeat(16),
            personal: 'ef'.repeat(16),
        },
    ];

    for (const testVector of testVectors) {
        const key =
            testVector.key === undefined
                ? undefined
                : Buffer.from(testVector.key, 'utf8');
        const salt =
            testVector.salt === undefined
                ? undefined
                : Buffer.from(testVector.salt, 'hex');
        const personal =
            testVector.personal === undefined
                ? undefined
                : Buffer.from(testVector.personal, 'hex');

        const computed = Buffer.from(
            blake2bMac256(
                key,
                personal ?? new Uint8Array(0),
                salt ?? new Uint8Array(0),
            ),
        ).toString('hex');

        // Dispatch to the correct function call depending on the existance of
        // salt and personal
        const expected = blake2bReference(
            outputLength,
            key,
            salt,
            personal,
            true,
        ).digest('hex');

        assert(
            expected === computed,
            `Computed Blake2b should match expected one. Key: ${testVector.key}, Salt: ${testVector.salt}, Personal: ${testVector.personal}`,
        );
    }
}

// Test XSalsa20Poly1305 by encrypting/decrypting between libthreema and TweetNacl
function testXSalsa20Poly1305() {
    // Compute Keys
    // Alice uses libthreema
    const aliceSecretKey = Buffer.from('aa'.repeat(32), 'hex');
    const alicePublicKey = x25519DerivePublicKey(aliceSecretKey);

    // Bob uses tweetNaCl
    const bobSecretKey = Buffer.from('bb'.repeat(32), 'hex');
    const bobPublicKey =
        tweetNaCl.box.keyPair.fromSecretKey(bobSecretKey).publicKey;

    // Derive shared keys between Alice and Bob
    const sharedKeyAliceBob = x25519HSalsa20DeriveSharedSecret(
        bobPublicKey,
        aliceSecretKey,
    );
    const sharedKeyBobAlice = tweetNaCl.box.before(
        alicePublicKey,
        bobSecretKey,
    );

    // Encrypt from Alice to Bob
    const nonceAliceBob = Buffer.from('01'.repeat(24), 'hex');
    const messageAliceBob = 'Hi Bob';
    const encryptedMessageAliceBob = xSalsa20Poly1305Encrypt(
        sharedKeyAliceBob,
        nonceAliceBob,
        Buffer.from(messageAliceBob, 'utf8'),
        [],
    );

    // Decrypt
    const decryptedMessageAliceBob = Buffer.from(
        tweetNaCl.box.open.after(
            encryptedMessageAliceBob,
            nonceAliceBob,
            sharedKeyBobAlice,
        ),
    ).toString('utf8');
    assert(
        decryptedMessageAliceBob === messageAliceBob,
        'Decrypted message from Alice to Bob does not match original one',
    );

    // Encrypt from Bob to Alice
    const nonceBobAlice = Buffer.from('02'.repeat(24), 'hex');
    const messageBobAlice = 'Hi Alice from Bob';
    const encryptedMessageBobAlice = tweetNaCl.box.after(
        Buffer.from(messageBobAlice, 'utf8'),
        nonceBobAlice,
        sharedKeyBobAlice,
    );

    // Decrypt
    const decryptedMessageBobAlice = Buffer.from(
        xSalsa20Poly1305Decrypt(
            sharedKeyAliceBob,
            nonceBobAlice,
            encryptedMessageBobAlice,
        ),
    ).toString('utf8');
    assert(
        decryptedMessageBobAlice === messageBobAlice,
        'Decrypted message from Bob to Alice does not match original one',
    );
}

async function testXChaCha20Poly1305() {
    await _sodium.ready;
    const sodium = _sodium;

    // Compute Keys
    // Alice uses libthreema
    const aliceSecretKey = Buffer.from('aa'.repeat(32, 'hex'));
    const alicePublicKey = x25519DerivePublicKey(aliceSecretKey);

    // Bob uses libsodium
    const bobKeys = sodium.crypto_box_seed_keypair(
        Buffer.from('bb'.repeat(32, 'hex')),
    );
    const bobSecretKey = bobKeys.privateKey;
    const bobPublicKey = bobKeys.publicKey;

    // Derive shared keys between Alice and Bob
    const sharedKeyAliceBob = x25519HSalsa20DeriveSharedSecret(
        bobPublicKey,
        aliceSecretKey,
    );
    const sharedKeyBobAlice = sodium.crypto_box_beforenm(
        alicePublicKey,
        bobSecretKey,
    );
    assert(
        Buffer.from(sharedKeyAliceBob).toString('hex') ===
            Buffer.from(sharedKeyBobAlice).toString('hex'),
        'Derived Keys do not match.',
    );

    // Encrypt from Alice to Bob
    const nonceAliceBob = Buffer.from('01'.repeat(24, 'hex'));
    const messageAliceBob = 'Hi Bob';
    const encryptedMessageAliceBob = xChaCha20Poly1305Encrypt(
        sharedKeyAliceBob,
        nonceAliceBob,
        Buffer.from(messageAliceBob, 'utf8'),
        [],
    );

    // Decrypt
    const decryptedMessageAliceBob = Buffer.from(
        sodium.crypto_aead_xchacha20poly1305_ietf_decrypt(
            null,
            encryptedMessageAliceBob,
            null,
            nonceAliceBob,
            sharedKeyBobAlice,
        ),
    ).toString('utf8');
    assert(
        decryptedMessageAliceBob === messageAliceBob,
        'Decrypted message from Alice to Bob does not match original one',
    );

    // Encrypt from Bob to Alice
    const nonceBobAlice = Buffer.from('02'.repeat(24, 'hex'));
    const messageBobAlice = 'Hi Alice from Bob';
    const encryptedMessageBobAlice =
        sodium.crypto_aead_xchacha20poly1305_ietf_encrypt(
            Buffer.from(messageBobAlice, 'utf8'),
            null,
            null,
            nonceBobAlice,
            sharedKeyBobAlice,
        );

    // Decrypt
    const decryptedMessageBobAlice = Buffer.from(
        xChaCha20Poly1305Decrypt(
            sharedKeyAliceBob,
            nonceBobAlice,
            encryptedMessageBobAlice,
            [],
        ),
    ).toString('utf8');

    assert(
        decryptedMessageBobAlice === messageBobAlice,
        'Decrypted message from Bob to Alice does not match original one',
    );
}

// Initialize libthreema
init();

// Run tests
evaluateTestResults([
    runTest(testSha256, 'SHA-256'),
    runTest(testHmacSha256, 'HMAC-SHA-256'),
    runTest(testArgon2id, 'Argon2id'),
    runTest(testScrypt, 'Scrypt'),
    runTest(testBlake2b, 'Blake2b'),
    runTest(testXSalsa20Poly1305, 'XSalsa20Poly1305'),
    runTest(testXChaCha20Poly1305, 'XChaCha20Poly1305'),
]);
