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

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

render_rooms_wrapper();

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, pubkey_received) => {
    console.log(`received: ${msg}`);
    let messages = document.getElementById(`messages-${room}`);
    if (!messages) {
        render_room(room, pubkey_received);
        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-', '');
    let {plaintext, tag} = decrypt_message(msg, pubkey, fromHexString(iv), fromHexString(nonce), associated_data);
    if (tag === tag_received && pubkey == pubkey_received) {
        append_message(1, plaintext, messages);
    } else {
        append_message(2, "Corrupted message.", messages);
    }
    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, messages) {
    form.addEventListener('submit', function(e) {
        e.preventDefault();
        if (input.value) {
            append_message(0, input.value, messages);
            window.scrollTo(0, document.body.scrollHeight);
            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);
    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);
    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;
}

function append_message(type, text, messageUl) {
    const item = document.createElement('li');
    const bubble = document.createElement('div');
    bubble.textContent = text;
    switch (type) {
        case 0:
            bubble.className = "sent-msg";
            item.className = "sent";
            break;
        case 1:
            bubble.className = "received-msg";
            item.className = "received";
            break;
        case 2:
            bubble.className = "corrupted-msg";
            item.className = "corrupted";
            break;
    }
    item.appendChild(bubble);
    messageUl.appendChild(item);
}

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'), '');

export async function reconnectSocket() {
    socket.disconnect();
    console.log("Socket disconnected.");
    setTimeout(() => {
        socket.connect();
        console.log("Socket reconnected.");
    }, 100);
}