import { genKeys, sharedKey } from "./ecdh.js";
import { keccakAEAD } from "./aead.js";
import { keccakKDF } from "./kdf.js";

const socket = io();
let secret = null;
let sharedsecret = {};
let dh_ratchets = {};
let sending_ratchets = {};
let receiving_ratchets = {};

function init_ratchets(order, user_pubkey) {
    let dh_ratchet = new keccakKDF()
    let key_1 =  dh_ratchet.init(sharedsecret[user_pubkey], new Uint8Array(0));
    let key_2 = dh_ratchet.next(new Uint8Array(0));
    dh_ratchets[user_pubkey] = dh_ratchet;
    let sending_ratchet = new keccakKDF();
    let receiving_ratchet = new keccakKDF();
    switch (order) {
        case 0:
            sending_ratchet.init(key_1, new Uint8Array(0));
            receiving_ratchet.init(key_2, new Uint8Array(0));
            break;
        case 1:
            sending_ratchet.init(key_2, new Uint8Array(0));
            receiving_ratchet.init(key_1, new Uint8Array(0));
            break;
    }
    sending_ratchets[user_pubkey] = sending_ratchet;
    receiving_ratchets[user_pubkey] = receiving_ratchet;

}

socket.on('chat message', (msg, room, tag_received, iv, nonce) => {
    console.log(`received: ${msg}`);
    const messages = document.getElementById(`messages-${room}`);
    const associated_data = fromHexString(Array.from((document.getElementById('pubkey')).classList).find(className => className.startsWith('key-')).replace('key-', ''));
    const pubkey = Array.from(messages.classList).find(className => className.startsWith('key-')).replace('key-', '');
    const item = document.createElement('li');
    let {plaintext, tag} = decrypt_message(msg, pubkey, fromHexString(iv), fromHexString(nonce), associated_data);
    item.textContent = plaintext
    console.log(tag);
    console.log(tag_received);
    messages.appendChild(item);
    if (document.getElementById(`room-${room}`).style.display === 'block') {
        window.scrollTo(0, document.body.scrollHeight);
    }
});

socket.on('key exchange', (user_pubkey, pubkey, part) => {
    let keys = null;
    switch (part) {
        case 0:
            keys = genKeys();
            secret = keys.privkey;
            socket.emit('key exchange', user_pubkey, toHexString(keys.pubkey), 1);
            break;
        case 1:
            keys = genKeys();
            secret = keys.privkey
            sharedsecret[user_pubkey] = sharedKey(secret, fromHexString(pubkey));
            socket.emit('key exchange', user_pubkey, toHexString(keys.pubkey), 2);
            console.log(`shared secret: ${toHexString(sharedsecret[user_pubkey])}`);
            init_ratchets(0, user_pubkey)
            break;
        case 2:
            sharedsecret[user_pubkey] = sharedKey(secret, fromHexString(pubkey));
            console.log(`shared secret: ${toHexString(sharedsecret[user_pubkey])}`);
            init_ratchets(1, user_pubkey)
            break;
    }
});

export function create_listener(form, input) {
    form.addEventListener('submit', function(e) {
        e.preventDefault();
        if (input.value) {
            const pubkey = Array.from(form.classList).find(className => className.startsWith('key-')).replace('key-', '');
            let {cipher, tag, iv, nonce} = encrypt_message(input.value, pubkey)
            socket.emit('chat message', cipher, form.id, tag, iv, nonce);
            input.value = '';
        }
    });
}

function encrypt_message(message, user_pubkey) {
    let encryption_key = sending_ratchets[user_pubkey].next();
    let encoded_msg = (new TextEncoder()).encode(message);
    let iv = generateRandomUint8Array();
    let nonce = generateRandomUint8Array();
    let associated_data = fromHexString(user_pubkey);
    let {cipher, tag} = keccakAEAD.encrypt(encryption_key, encoded_msg, iv, associated_data, nonce);
    console.log(iv)
    console.log(nonce)
    console.log(associated_data)
    console.log(encryption_key)
    console.log(encoded_msg)
    console.log(cipher)
    return {cipher: toHexString(cipher),
            tag: toHexString(tag),
            iv: toHexString(iv),
            nonce: toHexString(nonce)
            };
}

function decrypt_message(cipher, user_pubkey, iv, nonce, associated_data) {
    let decryption_key = receiving_ratchets[user_pubkey].next();
    let cipher_array = fromHexString(cipher);
    let {plaintext, tag} = keccakAEAD.decrypt(decryption_key, cipher_array, iv, associated_data, nonce);
    console.log(iv)
    console.log(nonce)
    console.log(associated_data)
    console.log(decryption_key)
    console.log(plaintext)
    console.log(cipher_array)
    return {plaintext: (new TextDecoder('utf-8')).decode(plaintext),
            tag: toHexString(tag)
            };
}

function generateRandomUint8Array(length = 16) {
    const randomArray = new Uint8Array(length);
    window.crypto.getRandomValues(randomArray);
    return randomArray;
}

const fromHexString = (hexString) =>
    Uint8Array.from(hexString.match(/.{1,2}/g).map((byte) => parseInt(byte, 16)));

const toHexString = (bytes) =>
    bytes.reduce((str, byte) => str + byte.toString(16).padStart(2, '0'), '');