import { cryptoRandom } from "../crypto/random_primes.js";
import { Region } from "./map.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 and sums with the original
 * set are 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(fortify) {
    let proofs = [];
    let privateInputs = [];

    let regionNames = Object.keys(fortify).sort();
    for (let x = 0; x < ROUNDS; x++) {
        let psi = cryptoShuffle(structuredClone(regionNames)).join("");
        let psiMap = {};
        for (let i = 0; i < regionNames.length; i++) {
            psiMap[regionNames[i]] = psi[i];
        }

        let newRegions = { regions: {} };

        // Rearrange keys
        for (let r of regionNames) {
            newRegions.regions[psiMap[r]] = fortify[r].pubKey.encrypt(
                -fortify[r].plainText
            );
        }

        let edges = [];
        let proofEdges = [];

        // Attach edges
        for (let i = 0; i < regionNames.length; i++) {
            let region = regionNames[i];
            let psiRegion = psi[i];

            for (let n of Region.getRegion(region).neighbours) {
                if (regionNames.includes(n.name)) {
                    let psiNeighbour = psiMap[n.name];

                    if (psiNeighbour > psiRegion) {
                        let salt = cryptoRandom(128);
                        let hasher = new jsSHA("SHA3-256", "TEXT");
                        hasher.update(psiRegion);
                        hasher.update(psiNeighbour);
                        hasher.update(salt.toString(16));
                        let hash = hasher.getHash("HEX");

                        edges.push({
                            hash: hash,
                            salt: salt,
                            edge: [psiRegion, psiNeighbour],
                        });
                        proofEdges.push(hash);
                    }
                }
            }
        }

        newRegions.edges = cryptoShuffle(proofEdges);

        proofs.push(newRegions);
        privateInputs.push({
            psi: psi,
            edges: edges,
        });
    }

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

    for (let i = 0; i < ROUNDS; i++) {
        let coin = coins[i];
        let proof = proofs[i].regions;
        let input = privateInputs[i];
        let psiMap = {};
        for (let i = 0; i < regionNames.length; i++) {
            psiMap[regionNames[i]] = input.psi[i];
        }

        if (coin === 1) {
            // Show |S| - 2 zeroes
            let verification = {
                regions: {},
            };
            let pairCipherText = null;
            let pair = [];

            for (let r of regionNames) {
                if (proof[r].plainText === 0n) {
                    verification.regions[r] = proof[r].proveNI();
                } else if (pairCipherText === null) {
                    pairCipherText = proof[r].clone();
                    pair.push(r);
                } else {
                    pairCipherText.update(proof[r]);
                    verification.pairCipherText = pairCipherText.proveNI();
                    pair.push(r);
                }
            }

            // Show pair is joined by edge
            let pairName = pair.sort();
            verification.pairEdgeSalt = input.edges.find(
                (e) => e.edge[0] === pairName[0] && e.edge[1] === pairName[1]
            ).salt;

            verifications.push(verification);
        } else {
            // Show isomorphism
            let edges = {};
            for (let e of input.edges) {
                edges[e.edge[0] + e.edge[1]] = e.salt;
            }

            let zeroProofs = {};
            for (let r of regionNames) {
                let c = proof[psiMap[r]].clone();
                c.update(fortify[r]);
                zeroProofs[r] = c.proveNI();
            }

            verifications.push({
                psi: input.psi,
                salts: edges,
                zeroProofs: zeroProofs,
            });
        }
    }

    // TODO range proof for regions

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

window.proveFortify = proveFortify;

export function verifyFortify(obj, key) {
    let coins = getCoins(JSON.stringify(obj.proofs));
    let fortify = obj.fortify;

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

        if (coin === 1) {
            // Check |S| - 2 zeroes
            if (regionNames.length - 2 !== Object.keys(verification.regions).length) {
                return false;
            }

            let regions = new Set(regionNames);

            for (let r of Object.keys(verification.regions)) {
                let c = new ReadOnlyCiphertext(key, BigInt(proof.regions[r]));
                let p = c.verifyNI(verification.regions[r]);
                if (p !== 0n) {
                    return false;
                }

                regions.delete(r);
            }

            // Check remaining pair sums to zero
            let pair = [...regions.values()].sort();
            let c = new ReadOnlyCiphertext(key, BigInt(proof.regions[pair[0]]));
            c.update(proof.regions[pair[1]]);
            let p = c.verifyNI(verification.pairCipherText);
            if (p !== 0n) {
                return false;
            }

            // Check edge is okay
            let hasher = new jsSHA("SHA3-256", "TEXT");
            hasher.update(pair[0]);
            hasher.update(pair[1]);
            hasher.update(verification.pairEdgeSalt.toString(16));
            let hash = hasher.getHash("HEX");

            if (!proof.edges.includes(hash)) {
                return false;
            }
        } else {
            // Check isomorphism
            let psi = verification.psi;
            let psiMap = {};
            for (let i = 0; i < regionNames.length; i++) {
                psiMap[regionNames[i]] = psi[i];
            }

            let salts = verification.salts;

            // Psi matches?
            if (psi.split("").sort().join("") !== regionNames.join("")) {
                return false;
            }

            // Ciphertexts match?
            for (let r of regionNames) {
                let c = new ReadOnlyCiphertext(key, BigInt(proof.regions[r]));
                c.update(fortify[psiMap[r]]);
                let p = c.verifyNI(verification.zeroProofs[r]);

                if (p !== 0n) {
                    return false;
                }
            }

            // Edges match?
            let proverEdges = new Set(proof.edges);

            for (let r of regionNames) {
                let psiRegion = psiMap[r];

                for (let n of Region.getRegion(r).neighbours) {
                    let psiNeighbour = psiMap[n.name];

                    if (psiNeighbour > psiRegion) {
                        let edgeSalt = salts[psiRegion + psiNeighbour];
                        let hasher = new jsSHA("SHA3-256", "TEXT");
                        hasher.update(psiRegion);
                        hasher.update(psiNeighbour);
                        hasher.update(edgeSalt.toString(16));
                        let hash = hasher.getHash("HEX");

                        if (proverEdges.has(hash)) {
                            proverEdges.delete(hash);
                        } else {
                            return false;
                        }
                    }
                }
            }

            // Check it is not over-specified
            if (proverEdges.size !== 0) {
                return false;
            }
        }
    }

    return true;
}

window.verifyFortify = verifyFortify;

// proveFortify({A:paillier.pubKey.encrypt(0n),B:paillier.pubKey.encrypt(3n),C:paillier.pubKey.encrypt(-3n),D:paillier.pubKey.encrypt(0n),E:paillier.pubKey.encrypt(0n)})