import { cryptoRandom } from "../crypto/random_primes.js";

function cryptoRange(upper) {
    // This is ridiculous: why implement a BigInt primitive, have it behave like a number, and then _not_ offer
    // mathematical operations like ilog2? Thankfully JavaScript is, for some reason, relatively fast at processing
    // strings (that is relative to it processing anything else)
    //
    // Actual comment: subtract 1 is important as otherwise this overestimates the logarithm for powers of two.
    let bitLength = BigInt((upper - 1n).toString(2).length + 1);
    let mask = 2n ** bitLength - 1n;

    let r = cryptoRandom(1024);

    while ((r & mask) >= BigInt(upper)) {
        r >>= bitLength;
    }

    return r & mask;
}

/**
 * CSPRNG Fisher-Yates shuffle.
 *
 * Only works on lists up to 255 elements.
 */
function cryptoShuffle(l) {
    let out = [];
    for (let i = l.length - 1; i > 0; i--) {
        let value = new Uint8Array([0]);
        crypto.getRandomValues(value);
        while (value[0] > i) {
            crypto.getRandomValues(value);
        }

        let v = l.splice(value[0], 1);
        out.push(v[0]);
    }

    out.push(l[0]);

    return out;
}

const ROUNDS = 24;

function getCoins(text) {
    // Construct verifier coins
    let hasher = new jsSHA("SHA3-256", "TEXT");
    hasher.update(text);
    let hash = hasher.getHash("UINT8ARRAY");

    let verifierCoins = [];
    for (let i = 0; i < ROUNDS / 8; i++) {
        let v = hash[i];
        for (let j = 0; j < 8; j++) {
            verifierCoins.push(v & 1);
            v >>= 1;
        }
    }

    return verifierCoins;
}

window.cryptoShuffle = cryptoShuffle;

export function proveRegions(regions) {
    // Construct prover coins
    let proofs = [];
    let privateInputs = [];

    let regionNames = Object.keys(regions).sort();
    for (let x = 0; x < ROUNDS; x++) {
        let psi = cryptoShuffle(structuredClone(regionNames)).join("");
        let newRegions = structuredClone(regions);

        // Rearrange keys
        for (let index = 0; index < regionNames.length; index++) {
            newRegions[regionNames[index]] = regions[psi[index]].pubKey.encrypt(
                regions[psi[index]].plainText * -1n
            );
        }

        proofs.push(newRegions);
        privateInputs.push(psi);
    }

    let verifierCoins = getCoins(JSON.stringify(proofs));

    let verifications = [];

    // Construct prover proofs
    for (let i = 0; i < ROUNDS; i++) {
        let coin = verifierCoins[i];
        let proof = proofs[i];
        let privateInput = privateInputs[i];

        if (coin === 1) {
            // Reveal bijection and proof for zero
            let zeroProofs = {};

            for (let i = 0; i < regionNames.length; i++) {
                let name = regionNames[i];
                let psiName = privateInput[i];

                let c = proof[name].clone();
                c.update(regions[psiName]);
                zeroProofs[name] = c.proveNI();
            }

            let ver = {
                psi: privateInput,
                zeroProofs: zeroProofs,
            };
            verifications.push(ver);
        } else {
            // Reveal proof for plaintext
            let valueProofs = {};

            for (let name of regionNames) {
                valueProofs[name] = proof[name].proveNI();
            }

            verifications.push({ valueProofs: valueProofs });
        }
    }

    return {
        regions: regions,
        proofs: proofs,
        verifications: verifications,
    };
}

window.proveRegions = proveRegions;

export function verifyRegions(obj, key) {
    let verifierCoins = getCoins(JSON.stringify(obj.proofs));
    let regions = obj.regions;
    let regionNames = Object.keys(regions).sort();

    for (let i = 0; i < ROUNDS; i++) {
        let proof = obj.proofs[i];
        let verification = obj.verifications[i];

        if (verifierCoins[i] === 1) {
            for (let regionName of regionNames) {
                // Undo psi
                let originalRegion = proof[regionName];

                // Compute product
                let c = new ReadOnlyCiphertext(key, BigInt(regions[regionName]));
                c.update(new ReadOnlyCiphertext(key, BigInt(originalRegion)));

                // Check ciphertext is zero
                let plaintext = c.verifyNI(verification.zeroProofs[regionName]);

                if (plaintext !== 0n) {
                    return false;
                }
            }
        } else {
            let foundOne = false;

            for (let name of Object.keys(verification.valueProofs)) {
                let ciphertext = new ReadOnlyCiphertext(key, BigInt(proof[name]));
                let plaintext = ciphertext.verifyNI(verification.valueProofs[name]);
                if (plaintext === null) {
                    return false;
                } else if (plaintext === 1n) {
                    if (foundOne) {
                        return false;
                    } else {
                        foundOne = true;
                    }
                }
            }
        }
    }

    return true;
}

window.verifyRegions = verifyRegions;

// verifyRegions(proveRegions({A:paillier.pubKey.encrypt(0n),B:paillier.pubKey.encrypt(1n),C:paillier.pubKey.encrypt(0n),D:paillier.pubKey.encrypt(0n),E:paillier.pubKey.encrypt(0n)}), paillier.pubKey)

export function proveRange(cipherText, rangeUpper) {
    if (cipherText.readOnly) {
        throw "Cannot prove range of ReadOnlyCiphertext";
    }

    let key = cipherText.pubKey;

    let proofs = [];

    // Construct \omega_1, \omega_2, ciphertexts
    for (let i = 0; i < ROUNDS; i++) {
        let o1 = cryptoRange(rangeUpper);
        let o2 = o1 - rangeUpper;

        let cs = cryptoShuffle([key.encrypt(o1), key.encrypt(o2)]);

        proofs.push({
            cs: cs,
        });
    }

    let coins = getCoins(JSON.stringify(proofs));
    let verifications = [];

    for (let i = 0; i < ROUNDS; i++) {
        let coin = coins[i];
        let proof = proofs[i];

        if (coin === 1) {
            // Prove that ciphertexts are valid
            verifications.push({
                c1: proof.cs[0].proveNI(),
                c2: proof.cs[1].proveNI(),
            });
        } else {
            // Prove that one of the sums is valid
            if ((cipherText.plainText + proof.cs[0].plainText) % key.n2 <= rangeUpper) {
                let ct = cipherText.clone();
                ct.update(proof.cs[0]);

                verifications.push({
                    csIndex: 0,
                    proof: ct.proveNI(),
                });
            } else {
                let ct = cipherText.clone();
                ct.update(proof.cs[1]);

                verifications.push({
                    csIndex: 1,
                    proof: ct.proveNI(),
                });
            }
        }
    }

    return {
        cipherText: cipherText,
        rangeUpper: "0x" + rangeUpper.toString(16),
        proofs: proofs,
        verifications: verifications,
    };
}

window.proveRange = proveRange;

export function verifyRange(obj, key) {
    let coins = getCoins(JSON.stringify(obj.proofs));
    let rangeUpper = BigInt(obj.rangeUpper);

    for (let i = 0; i < ROUNDS; i++) {
        let coin = coins[i];
        let proof = obj.proofs[i];
        let verification = obj.verifications[i];

        if (coin === 1) {
            let c1 = new ReadOnlyCiphertext(key, BigInt(proof.cs[0]));
            let c2 = new ReadOnlyCiphertext(key, BigInt(proof.cs[1]));
            let o1 = c1.verifyNI(verification.c1);
            let o2 = c2.verifyNI(verification.c2);
            let diff;

            if (o1 < o2) {
                o2 -= key.n2;
                diff = o1 - o2;
            } else {
                o1 -= key.n2;
                diff = o2 - o1;
            }

            if (diff !== rangeUpper) {
                return false;
            }
        } else {
            let c = new ReadOnlyCiphertext(key, BigInt(proof.cs[verification.csIndex]));
            c.update(new ReadOnlyCiphertext(key, BigInt(obj.cipherText)));
            let value = c.verifyNI(verification.proof);

            if (value === null || value > rangeUpper) {
                return false;
            }
        }
    }

    return true;
}

window.verifyRange = verifyRange;

/**
 *  - We prove that the set contains |S| - 2 zeros, with the final pair summing to zero
 *  - We also attach some form of adjacency guarantee: that is, we prove the sums on all adjacent pairs are zero
 *  - We also attach a range proof for the new region values
 */
export function proveFortify() {}