import { socket, game } from "./main.js";
import { Packet } from "./packet.js";
import { cryptoRandom } from "../crypto/random_primes.js";

class RandomSession {
    constructor(range) {
        this.range = range;
        this.hmacs = {};
        this.keys = {};
        this.noises = {};
        this.ourKey = cryptoRandom(1024).toString(16);
        this.ourNoise = cryptoRandom(64);
        this.finalValue = null;
        this.resolvers = [];
    }

    hmac() {
        let hasher = new jsSHA("SHA3-256", "HEX");
        hasher.update(this.ourKey);
        hasher.update(this.ourNoise.toString(16));
        return hasher.getHash("HEX");
    }
}

export class Random {
    constructor() {
        this.sessions = {};
    }

    async get(n, sessionId) {
        if (this.sessions[sessionId] === undefined) {
            this.initialiseSession(n, sessionId);
        }

        let promise;
        await navigator.locks.request(`random-${sessionId}`, () => {
            if (this.sessions[sessionId].finalValue === null) {
                let session = this.sessions[sessionId];
                let resolver;
                promise = new Promise((resolve) => {
                    resolver = resolve;
                });
                session.resolvers.push(resolver);
            } else {
                promise = new Promise((resolve) => {
                    resolve(this.sessions[sessionId].finalValue);
                });
            }
        });
        return promise;
    }

    /**
     * Start a cooperative random session.
     */
    initialiseSession(n, sessionId) {
        let session = new RandomSession(n);
        if (sessionId === undefined) {
            sessionId = window.crypto.randomUUID();
        }

        this.sessions[sessionId] = session;

        socket.emit(
            "message",
            Packet.createRandomHMAC(sessionId, n, session.hmac(), session.ourKey)
        );

        return session;
    }

    /**
     * Process cooperative random protocol.
     *
     * @param data Packet received
     */
    async processCooperativeRandom(data) {
        // Step 0: extract relevant information from data
        let session = this.sessions[data.session];
        const stage = data.stage;

        if (session === undefined) {
            session = this.initialiseSession(data.range, data.session);
        }

        if (stage === "COMMIT") {
            session.hmacs[data.author] = data.hmac;
            session.keys[data.author] = data.key;

            if (
                Object.keys(session.hmacs).length ===
                Object.keys(game.players).length - 1
            ) {
                // Step 3: release our key once all players have sent a ciphertext
                socket.emit(
                    "message",
                    Packet.createRandomNoise(
                        data.session,
                        "0x" + session.ourNoise.toString(16)
                    )
                );
            }
        } else if (stage === "REVEAL") {
            // Check HMAC
            let noise = BigInt(data.noise) % 2n ** 64n;
            let hasher = new jsSHA("SHA3-256", "HEX");
            hasher.update(session.keys[data.author]);
            hasher.update(noise.toString(16));
            let hash = hasher.getHash("HEX");

            if (hash === session.hmacs[data.author]) {
                session.noises[data.author] = noise;
            }

            // Step 4: get final random value
            if (
                Object.keys(session.noises).length ===
                Object.keys(game.players).length - 1
            ) {
                // Lock out wait calls as they may resolve to never-ending promises.
                await navigator.locks.request(`random-${data.session}`, () => {
                    let total = session.ourNoise;

                    for (let noise of Object.values(session.noises)) {
                        total += noise;
                    }

                    // Find first good block of bits
                    let blockSize = BigInt(Math.ceil(Math.log2(session.range)));
                    let blockMask = 2n ** blockSize - 1n;

                    while ((total & blockMask) >= BigInt(session.range)) {
                        total >>= blockSize;
                    }

                    session.finalValue = total & blockMask;

                    this.resolve(data.session);
                });
            }
        }
    }

    /**
     * Resolve a session by calling any callbacks associated with the session.
     *
     * @param sessionId
     */
    resolve(sessionId) {
        const session = this.sessions[sessionId];
        for (let resolve of session.resolvers) {
            resolve(session.finalValue);
        }
    }
}