import KBucket from 'k-bucket';
import { MessageInterface, MessageListener } from '../communication-interface/MessageInterface';
import { getRandomUUID, getRandomNumber, getTimeBasedAddress, sha256, getRandomSet } from '../crypto';
import { timeout } from '../helper';
import { RsaOaepAesMessage } from '../crypto/rsa-oaep-aes-message';
import { LongAsyncRemoteFunction } from '../objectTransport/LongAsyncRemoteFunction';
import { WebRTCTransport } from '../webrtc-transport/WebRTCTransport';
import { EcdheAesGcmMessage } from '../crypto/ecdhe-aes-gcm-message';
import { ContactManager } from '../contact-management/ContactManager';

const PING_TIMEOUT = 10000;
const MINIMUM_NUMBER_OF_KNOWN_PEERS = 20;
const TTL = 20;

class DhtContact {
    id: Uint8Array;
    vectorClock: number;
    publicKey: string;
    signature: string;


    constructor(id: Uint8Array, publicKey: string, signature: string, vectorClock = 0) {
        this.id = id;
        this.publicKey = publicKey;
        this.signature = signature;
        this.vectorClock = vectorClock;
    }
}

type DhtContactDto = {
    id: string,
    publicKey: string,
    signature: string,
    vectorClock: number
}

interface DhtParameterMap {
    'ping': { pingId: string },
    'pong': { pingId: string },
    'findNode': { id: string, respondTo?: string },
    'foundNode': { contacts: DhtContactDto[] },
    'handshakeWith': { from: string, fromPK: string, pkSignature: string, webRTCAddressBundle: string },
    'handshakeResponse': { from: string, fromPK: string, pkSignature: string, webRTCAddressBundle: string, respondVia: string },
    'handshakeCompleted': { from: string, fromPK: string, pkSignature: string, respondVia: string },
    'send': { to: string, message: string, ttl: number },
    'broadcast': { message: string, type: string },
}

export class WebRtcDht implements MessageInterface {
    private kBucket: KBucket<DhtContact>;
    private _initPromise: Promise<void>;
    private webRTCTransport: WebRTCTransport;
    private ecdhAesMessage: EcdheAesGcmMessage;
    private longAsyncRF: LongAsyncRemoteFunction<DhtParameterMap>;
    private contactManager: ContactManager;
    private pkSignatures = new Map<string, string>();

    private readonly textEncoder = new TextEncoder();
    private readonly textDecoder = new TextDecoder();

    private timeAddressToUserId = new Map<string, string>();
    private rsaOaepeAesMessage: RsaOaepAesMessage;

    /**
     * 
     * @param userId The local userId
     * @param webRTCTransport an existing WebRTCTransport with at least one connection
     * @param ecdheMessageWebRTCTransport the EcdheAesGcmMessage used by the WebRTCTransport
     */
    constructor(userId: string, webRTCTransport: WebRTCTransport, ecdheMessageWebRTCTransport: EcdheAesGcmMessage, rsaOaepAesMessage: RsaOaepAesMessage, contactManager: ContactManager) {
        //link transport
        this.webRTCTransport = webRTCTransport;
        //link encryption, already connected to transport (given to us because it shares keys with a different transport)
        this.ecdhAesMessage = ecdheMessageWebRTCTransport;
        //end-to-end encryption based on dht content
        this.rsaOaepeAesMessage = rsaOaepAesMessage;
        //integrity of dht content
        this.contactManager = contactManager;
        this.longAsyncRF = new LongAsyncRemoteFunction(this.ecdhAesMessage, 'dhtTransport');
        this.longAsyncRF.register('ping', this._r_ping.bind(this));
        this.longAsyncRF.register('pong', this._r_pong.bind(this));
        this.longAsyncRF.register('findNode', this._r_findNode.bind(this));
        this.longAsyncRF.register('foundNode', this._r_foundNode.bind(this));
        this.longAsyncRF.register('handshakeWith', this._r_handshakeWith.bind(this));
        this.longAsyncRF.register('handshakeResponse', this._r_handshakeResponse.bind(this));
        this.longAsyncRF.register('handshakeCompleted', this._r_handshakeCompleted.bind(this));
        this.longAsyncRF.register('send', this._r_onMessage.bind(this));
        this.longAsyncRF.register('broadcast', this._r_onBroadcast.bind(this));
        this.kBucket = new KBucket<DhtContact>({
            numberOfNodesPerKBucket: 20,
            localNodeId: this.textEncoder.encode(userId),
        });

        this._initPromise = this.init(userId);
    }

    async addRTCConnection(userId: string, fromPK: string, signature: string) {
        await this._initPromise;
        await this.addContact(userId, fromPK);
        this.kBucket.add(new DhtContact(this.textEncoder.encode(userId), fromPK, signature));
    }

    async connectToMorePeers() {
        try {
            const knownNodesWeAreNotConnectedTo = this.getKnownNodes().filter(node => !this.webRTCTransport?.getAllContacts().includes(node) && !this.outstandingConnections.has(node));
            const numberOfMissingConnections = this.kBucket.numberOfNodesPerKBucket - this.kBucket.count() - this.outstandingConnections.size;
            const randomNodes = getRandomSet(knownNodesWeAreNotConnectedTo, numberOfMissingConnections);
            for (const node of randomNodes) {
                void this.connectTo(node);
            }
        } catch (err) {
            console.error(err);
        }
    }

    async getNodeAddress(userId: string, prev?: string, depthLimit = 20): Promise<void> {
        if (depthLimit <= 0) {
            throw new Error('depth limit reached');
        }
        await this._initPromise;
        const bestContact = await this._getClosestContact(userId, prev);
        if (!bestContact) {
            throw new Error('no contact to forward the message to');
        }
        if (this.textDecoder.decode(bestContact.id) === userId) {
            return
        }
        await this.findNodeOn(userId, this.textDecoder.decode(bestContact.id));
        return await this.getNodeAddress(userId, this.textDecoder.decode(bestContact.id), depthLimit - 1);
    }

    async connectTo(userId: string) {
        const fromAddressBundle = await this.webRTCTransport.getOwnAddress(userId);
        const fromPK = await this.getKeyInformation();
        const handshake = this.longAsyncRF!.getCallCode('handshakeWith', userId, { from: await this.getOwnUserId(), fromPK, pkSignature: await this.contactManager.sign(fromPK), webRTCAddressBundle: fromAddressBundle });
        await this.send(userId, handshake);
        await this.longAsyncRF!.expectCallInvocation({ functionName: 'handshakeResponse' }, userId);
    }

    private outstandingConnections = new Map<string, { resolve: () => void, reject: () => void }>();

    padMessage(message: string) {
        const padlength = Math.ceil(message.length / 1000);
        const padding = ' '.repeat(padlength);
        return message + padding.slice(message.length);
    }

    async sendOnion(userId: string, message: string, onion: string[]) {
        await this._initPromise;
        if (onion.length === 0) {
            throw new Error('onion must have at least one layer');
        }

        message = this.longAsyncRF.getCallCode('send', userId, { to: userId, message, ttl: TTL });
        message = await this.rsaOaepeAesMessage.encryptTo(userId, message);

        for (const layer of onion.slice(0, -1)) {
            message = this.longAsyncRF.getCallCode('send', layer, { to: layer, message, ttl: TTL });
            message = await this.rsaOaepeAesMessage.encryptTo(layer, message);
        }
        await this.send(onion[onion.length - 1], message);
    }

    async _getClosestContact(userId: string, ignoreUserId?: string) {
        await this._initPromise;
        const closestContacts = this.kBucket?.closest(this.textEncoder.encode(userId), 10);
        for (const contact of closestContacts) {
            if (this.textDecoder.decode(contact.id) === await this.getOwnUserId()) {
                continue;
            }
            if (ignoreUserId && this.textDecoder.decode(contact.id) === ignoreUserId) {
                continue;
            }
            //if the contact is added to the kbucket, but connectionestablished has not been called on the webrtc transport yet
            if (!this.webRTCTransport.getAllContacts().includes(this.textDecoder.decode(contact.id))) {
                continue;
            }
            return contact;
        }
    }

    /**
     * Sends a message via the DHT network
     * @param userId the time based address of the recipient
     * @param message the message to send
     * @param isRetry do not set, for internal use only
     * @param from do not set, for internal use only
     */
    async send(userId: string, message: string, isRetry = false, from?: string, ttl?: number) {
        await this._initPromise;
        ttl = Math.min((ttl ?? TTL), TTL) - 1;
        if (!from) {
            //encrypt message end-to-end
            try {
                message = await this.rsaOaepeAesMessage.encryptTo(userId, message);
            } catch (err: any) {
                if (err?.message === 'Contact not found') {
                    await this.getNodeAddress(userId);
                    await this.send(userId, message, isRetry, from);
                } else {
                    throw err;
                }
            }
        }

        const directConnected = this.webRTCTransport.getAllContacts().includes(userId);
        if (directConnected) {
            try {
                await this.longAsyncRF!.call('send', userId, { to: userId, message, ttl });
                return;
            } catch (err) {
                console.error(err);
                this.webRTCTransport.removeContact(userId);
                await this.send(userId, message, isRetry, from);
            }
        }
        const bestContact = await this._getClosestContact(userId, from);
        if (!bestContact) {
            if (!isRetry) {
                console.warn('closest contact is self, retrying in 15 seconds...');
                await timeout(15000 * (Math.random() + 0.5));
                await this.send(userId, message, true, from);
                return;
            }
            throw new Error('no contact to forward the message to');
        }
        const closestContact = bestContact;

        console.info(`send to ${userId} via ${this.textDecoder.decode(closestContact.id)}`);

        console.info(`dht num contacts: ${this.kBucket.count() - 1}`);


        try {
            await this.longAsyncRF!.call('send', this.textDecoder.decode(closestContact.id), { to: userId, message, ttl: TTL });
        } catch (e) {
            console.warn(e);
            console.info(`removing disconnected contact ${this.textDecoder.decode(closestContact.id)}`);
            this.kBucket.remove(closestContact.id);
            if (this.kBucket.count() < this.kBucket.numberOfNodesPerKBucket) {
                void this.connectToMorePeers();
            }
            await this.send(userId, message, isRetry, from);
        }
    }

    private listeners: MessageListener[] = [];

    async _r_onMessage(from: string, { message, to, ttl }: DhtParameterMap['send']) {
        if (to === await this.contactManager.getOwnTimeBasedAddress()) {
            console.log(`message for us from ${from}`);
            message = (await this.rsaOaepeAesMessage.decryptMessage(message)).payload;
            //if this message can be interpreted by the longAsyncRF, do so
            if (!(await this.longAsyncRF.messageRecieved(from, message))) {
                //otherwise, pass it to the listeners

                //resolve time based address (only works for contacts, will throw if not a contact)
                from = await this.contactManager.getAddressFromTimeBasedAddress(from);
                this.listeners.forEach(listener => listener(from, message));
            }
        } else {
            //if the message is not for us, forward it
            if (ttl > 0) {
                await this.send(to, message, false, from, ttl);
                console.log(`forwarded message from ${from} to ${to}`);
            }
        }
    }

    private knownBroadcasts = new Set<string>();

    async _r_onBroadcast(from: string, { message, type }: { message: string, type: string }) {
        const hash = await sha256(message);
        if (this.knownBroadcasts.has(hash)) {
            return;
        }
        this.knownBroadcasts.add(hash);

        console.log(`broadcast from ${from}: ${message}`);
        switch (type) {
            case 'violation':
                this.kBucket.remove(this.textEncoder.encode(from));
                break;
            default:
                console.error(`unknown broadcast type ${type}`);
                return;
        }
        const randGUID = await getRandomUUID();
        for (const contact of this.kBucket.closest(this.textEncoder.encode(randGUID), 3)) {
            await this.longAsyncRF!.call('broadcast', this.textDecoder.decode(contact.id), { message, type });
        }
    }


    addMessageListener(listener: MessageListener): void {
        this.listeners.push(listener);
    }

    private knownNodes = new Set<string>();

    async addContact(userId: string, keyInformation: any, replaceExisting = false): Promise<string | void> {
        await this._initPromise;
        const { ecdhKeyInfo, rsaKeyInfo } = JSON.parse(keyInformation);
        await this.ecdhAesMessage!.addContact(userId, ecdhKeyInfo, replaceExisting);
        await this.rsaOaepeAesMessage.addContact(userId, rsaKeyInfo, replaceExisting);
        this.knownNodes.add(userId);
    }

    getKnownNodes(): string[] {
        return [...this.knownNodes];
    }

    async getKeyInformation(): Promise<any> {
        await this._initPromise;
        const ecdhKeyInfo = await this.ecdhAesMessage.getKeyInformation();
        const rsaKeyInfo = await this.rsaOaepeAesMessage.getKeyInformation();
        return JSON.stringify({ ecdhKeyInfo, rsaKeyInfo });
    }

    /**
     * identical to getTimeBasedUserId
     * @returns the time based address of the user
     */
    async getOwnUserId(): Promise<string> {
        await this._initPromise;
        return await this.ecdhAesMessage!.getOwnUserId();
    }

    /**
     * identical to getOwnUserId
     * @returns the time based address of the user
     */
    async getTimeBasedUserId(): Promise<string> {
        return await this.getOwnUserId();
    }

    async findNodeOn(search: string, userId: string) {
        await this._initPromise;
        const findNode = this.longAsyncRF!.getCallCode('findNode', userId, { id: search });
        await this.send(userId, findNode);
        await this.longAsyncRF!.expectCallInvocation({ functionName: 'foundNode' }, userId);
    }


    async _r_findNode(from: string, findNode: DhtParameterMap['findNode']) {
        await this._initPromise;
        const ownUserId = await this.getOwnUserId();
        const contacts = this.kBucket?.closest(this.textEncoder.encode(findNode.id), 3).map(contact => ({ id: this.textDecoder.decode(contact.id), publicKey: contact.publicKey, vectorClock: contact.vectorClock, signature: contact.signature })).filter(c => c.id != ownUserId) || [];
        if (!findNode.respondTo) {
            await this.longAsyncRF!.call('foundNode', from, { contacts });
        } else {
            const foundNode = this.longAsyncRF!.getCallCode('foundNode', findNode.respondTo, { contacts });
            await this.send(from, foundNode);
        }
    }

    async _r_foundNode(from: string, foundNode: { contacts: DhtContactDto[] }) {
        for (const contact of foundNode.contacts) {
            if (contact.id === await this.getOwnUserId()) {
                continue;
            }
            if (contact.id === from) {
                continue;
            }
            const existing = this.kBucket.get(this.textEncoder.encode(contact.id));
            if (existing?.id === this.textEncoder.encode(contact.id)) {
                continue;
            }
            console.info(`found contact ${contact.id} via ${from}`);
            try {
                const verified = await this.contactManager.verify(contact.id, contact.publicKey, contact.signature);
                this.pkSignatures.set(contact.id, contact.signature);
                await this.addContact(contact.id, contact.publicKey, verified);
                try {
                    await this.connectTo(contact.id);
                } catch (err) {
                    console.error(`contactAdded ${contact.id}, establishing connection failed`);
                    console.error(err);
                }
            } catch (err) {
                console.error(`contactAdded ${contact.id}, signature verification failed: ${err}`);
                await this._r_onBroadcast(await this.getOwnUserId(), { message: 'contact signature verification failed', type: 'violation' });
            }

        }
    }

    async _r_handshakeWith(fromTransportAddress: string, { from, fromPK, pkSignature, webRTCAddressBundle }: DhtParameterMap['handshakeWith']) {
        await this._initPromise;
        const verified = await this.contactManager.verify(from, fromPK, pkSignature);
        this.pkSignatures.set(from, pkSignature);
        const myAddressBundle = await this.webRTCTransport.addContact(from, webRTCAddressBundle);
        await this.addContact(from, fromPK, verified);
        if (!myAddressBundle) {
            throw new Error('webRTC addContact did not return a value!');
        }
        const myFromPK = await this.getKeyInformation();
        const respondVia = getRandomSet(this.webRTCTransport.getAllContacts(), 1)[0];
        const response = this.longAsyncRF!.getCallCode('handshakeResponse', from, { from: await this.getOwnUserId(), fromPK: myFromPK, pkSignature: await this.contactManager.sign(myFromPK), webRTCAddressBundle: myAddressBundle, respondVia });
        await this.send(from, response);

    }

    async _r_handshakeResponse(fromTransportAddress: string, { from, fromPK, pkSignature, webRTCAddressBundle, respondVia }: DhtParameterMap['handshakeResponse']) {
        await this._initPromise;
        const verified = await this.contactManager.verify(from, fromPK, pkSignature);
        await this.webRTCTransport.addContact(from, webRTCAddressBundle);
        await this.addContact(from, fromPK, verified);
        this.kBucket.add(new DhtContact(this.textEncoder.encode(from), fromPK, pkSignature));
        this.outstandingConnections.get(from)?.resolve();
        //if we need more contacts, ask new contact for his contacts close to us
        const myFromPk = await this.getKeyInformation();
        const myRespondVia = getRandomSet(this.webRTCTransport.getAllContacts(), 1)[0];
        const response = this.longAsyncRF!.getCallCode('handshakeCompleted', from, { from: await this.getOwnUserId(), fromPK: myFromPk, pkSignature: await this.contactManager.sign(myFromPk), respondVia: myRespondVia });
        await this.sendOnion(from, response, [respondVia]);
    }

    async _r_handshakeCompleted(fromTransport: string, { from, fromPK, pkSignature }: DhtParameterMap['handshakeCompleted']) {
        if (this.outstandingConnections.has(from)) {
            this.outstandingConnections.get(from)?.resolve();
        }
        const signature = this.pkSignatures.get(from);
        if (!signature) {
            throw new Error('no signature found for contact');
        }
        await this.contactManager.verify(from, fromPK, signature);
        this.kBucket.add(new DhtContact(this.textEncoder.encode(from), fromPK, signature));
    }

    async _r_ping(from: string, { pingId }: { pingId: string }) {
        await this._initPromise;
        await this.longAsyncRF!.call('pong', from, { pingId });
    }

    _r_pong(from: string, { pingId }: { pingId: string }) {
        const newOutstandingPings = this.outstandingPings.get(pingId)?.filter((contact) => contact.id !== this.textEncoder.encode(from));
        if (newOutstandingPings) {
            this.outstandingPings.set(pingId, newOutstandingPings);
        }
    }

    async init(userId: string) {
        const ownKeyInfo = await this.getKeyInformation();
        this.kBucket.add(new DhtContact(this.textEncoder.encode(userId), ownKeyInfo, await this.contactManager.sign(ownKeyInfo)));
        this.kBucket.on('ping', this.sendPing.bind(this));
        this.kBucket.on('added', this.contactAdded.bind(this));
    }

    private outstandingPings = new Map<string, DhtContact[]>();

    async sendPing(oldContacts: DhtContact[], newContact: DhtContact) {
        await this._initPromise;
        const id = await getRandomUUID();
        this.outstandingPings.set(id, oldContacts);
        for (const contact of oldContacts) {
            await this.longAsyncRF!.call('ping', this.textDecoder.decode(contact.id), { pingId: id });
        }
        await timeout(PING_TIMEOUT);
        this.outstandingPings.get(id)?.forEach(contact => this.kBucket?.remove(contact.id));
        if ((this.outstandingPings.get(id)?.length || 0) > 0) {
            this.kBucket?.add(newContact);
        }
    }

    async contactAdded(contact: DhtContact) {
        await this._initPromise;
        console.info(`contact added: ${this.textDecoder.decode(contact.id)}`);
        if (this.kBucket.count() || 0 < this.kBucket.numberOfNodesPerKBucket || this.knownNodes.size < MINIMUM_NUMBER_OF_KNOWN_PEERS) {
            try {
                await this.longAsyncRF.call('findNode', this.textDecoder.decode(contact.id), { id: await this.getOwnUserId() }, { functionName: 'foundNode' });
            } catch (e) {
                console.error(e);
            }
        }
    }
}
