import { cryptoRandom, generate_prime, KEY_SIZE } from "./random_primes.js";
import { gcd, mod_exp } from "./math.js";

const PAILLIER = 0;
const JURIK = 1;

class Ciphertext {
    constructor(key, plainText, r, set) {
        if (set !== undefined) {
            this.pubKey = key;
            this.plainText = plainText;

            this.readOnly = false;

            return;
        }

        if (r === undefined) {
            // Use the optimised form using Jacobi classes
            r = cryptoRandom();

            // Compute g^m by binomial theorem.
            let gm = (1n + key.n * plainText) % key.n2;

            // Compute g^m h^r.
            this.cipherText = (gm * mod_exp(key.hn, r, key.n2)) % key.n2;

            // Force into range.
            while (this.cipherText < 0n) {
                this.cipherText += key.n2;
            }

            this.mode = JURIK;
            this.r = mod_exp(key.h, r, key.n);
        } else {
            // Use the standard form
            // Compute g^m by binomial theorem.
            let gm = (1n + key.n * plainText) % key.n2;

            // Compute g^m r^n.
            this.cipherText = (gm * mod_exp(r, key.n, key.n2)) % key.n2;

            // Force into range.
            while (this.cipherText < 0n) {
                this.cipherText += key.n2;
            }

            this.mode = PAILLIER;
            this.r = r;
        }

        this.pubKey = key;
        this.plainText = plainText;

        this.readOnly = false;
    }

    update(c) {
        this.cipherText = (this.cipherText * c.cipherText) % this.pubKey.n2;
        this.r = (this.r * c.r) % this.pubKey.n2;
        this.plainText += c.plainText;

        // Force into range
        while (this.cipherText < 0n) {
            this.cipherText += this.pubKey.n2;
        }
    }

    toString() {
        return "0x" + this.cipherText.toString(16);
    }

    prove() {
        return new ValueProofSessionProver(this);
    }

    // Construct a non-interactive proof
    proveNI() {
        let rp = cryptoRandom(KEY_SIZE * 2);
        while (rp >= this.pubKey.n) {
            rp = cryptoRandom(KEY_SIZE * 2);
        }

        let a = mod_exp(rp, this.pubKey.n, this.pubKey.n2);
        let hasher = new jsSHA("SHAKE256", "HEX");

        let plainText = this.plainText.toString(16);
        if (plainText.length % 2 !== 0) {
            plainText = "0" + plainText;
        }

        let aStr = a.toString(16);
        if (aStr.length % 2 !== 0) {
            aStr = "0" + aStr;
        }

        hasher.update(this.pubKey.g.toString(16));
        hasher.update(plainText);
        hasher.update(aStr);

        let challenge = BigInt("0x" + hasher.getHash("HEX", { outputLen: 2048 }));

        return {
            plainText: "0x" + this.plainText.toString(16),
            a: "0x" + a.toString(16),
            challenge: "0x" + challenge.toString(16),
            proof:
                "0x" +
                (
                    ((rp % this.pubKey.n) * mod_exp(this.r, challenge, this.pubKey.n)) %
                    this.pubKey.n
                ).toString(16),
        };
    }

    asReadOnlyCiphertext() {
        return new ReadOnlyCiphertext(this.pubKey, this.cipherText);
    }

    clone() {
        let c = new Ciphertext(this.pubKey, this.plainText, 0, true);
        c.cipherText = this.cipherText;
        c.r = this.r;
        c.mode = this.mode;

        return c;
    }
}

class ValueProofSessionProver {
    constructor(cipherText) {
        this.cipherText = cipherText;

        this.rp = cryptoRandom(KEY_SIZE * 2);
        while (this.rp >= this.cipherText.pubKey.n) {
            this.rp = cryptoRandom(KEY_SIZE * 2);
        }
    }

    get a() {
        return mod_exp(this.rp, this.cipherText.pubKey.n, this.cipherText.pubKey.n2);
    }

    noise() {
        return mod_exp(this.rp, this.cipherText.pubKey.n, this.cipherText.pubKey.n2);
    }

    prove(challenge) {
        return (
            ((this.rp % this.cipherText.pubKey.n) *
                mod_exp(this.cipherText.r, challenge, this.cipherText.pubKey.n)) %
            this.cipherText.pubKey.n
        );
    }

    asVerifier() {
        return this.cipherText
            .asReadOnlyCiphertext()
            .prove(this.cipherText.plainText, this.noise());
    }
}

window.Ciphertext = Ciphertext;

export class ReadOnlyCiphertext {
    constructor(key, cipherText) {
        this.cipherText = cipherText;
        this.pubKey = key;

        this.readOnly = true;
    }

    update(c) {
        this.cipherText = (this.cipherText * c.cipherText) % this.pubKey.n2;

        // Force into range
        while (this.cipherText < 0n) {
            this.cipherText += this.pubKey.n2;
        }
    }

    prove(plainText, a) {
        return new ValueProofSessionVerifier(this, plainText, a);
    }

    verifyNI(statement) {
        let verifier = new ValueProofSessionVerifier(
            this,
            BigInt(statement.plainText),
            BigInt(statement.a),
            BigInt(statement.challenge)
        );

        return verifier.verify(BigInt(statement.proof));
    }

    clone() {
        return new ReadOnlyCiphertext(this.pubKey, this.cipherText);
    }
}

class ValueProofSessionVerifier {
    constructor(cipherText, plainText, a, challenge) {
        // Clone, otherwise the update below will mutate the original value
        this.cipherText = cipherText.clone();
        this.cipherText.update(this.cipherText.pubKey.encrypt(-1n * plainText, 1n));

        if (challenge === undefined) {
            // Shift the challenge down by 1 to ensure it is smaller than either prime factor.
            this.challenge = cryptoRandom(KEY_SIZE) << 1n;
        } else {
            this.challenge = challenge;
        }

        this.a = a;

        this.plainText = plainText;
    }

    verify(proof) {
        // check coprimality
        if (gcd(proof, this.cipherText.pubKey.n) !== 1n) return -1;
        if (gcd(this.cipherText.cipherText, this.cipherText.pubKey.n) !== 1n) return -2;
        if (gcd(this.a, this.cipherText.pubKey.n) !== 1n) return -3;

        // check exp
        return mod_exp(proof, this.cipherText.pubKey.n, this.cipherText.pubKey.n2) ===
            (this.a *
                mod_exp(
                    this.cipherText.cipherText,
                    this.challenge,
                    this.cipherText.pubKey.n2
                )) %
                this.cipherText.pubKey.n2
            ? 1
            : -4;
    }
}

window.ReadOnlyCiphertext = ReadOnlyCiphertext;

export class PaillierPubKey {
    constructor(n, h) {
        this.n = n;

        if (h === undefined) {
            let x = cryptoRandom(KEY_SIZE * 2);
            while (x >= this.n) {
                x = cryptoRandom(KEY_SIZE * 2);
            }

            this.h = ((-1n * x ** 2n) % this.n) + this.n;
        } else {
            this.h = h;
        }

        this.g = this.n + 1n;

        this.n2 = this.n ** 2n;
        this.hn = mod_exp(this.h, this.n, this.n2);
    }

    encrypt(m, r) {
        return new Ciphertext(this, m, r);
    }

    toJSON() {
        return {
            n: "0x" + this.n.toString(16),
            h: "0x" + this.h.toString(16),
        };
    }

    static fromJSON(data) {
        return new PaillierPubKey(BigInt(data.n), BigInt(data.h));
    }
}

class PaillierPrivKey {
    constructor(p, q) {
        this.n = p * q;
        // precompute square of n
        this.n2 = this.n ** 2n;
        this.lambda = (p - 1n) * (q - 1n);
        this.mu = mod_exp(this.lambda, this.lambda - 1n, this.n);
    }

    decrypt(c) {
        return (((mod_exp(c, this.lambda, this.n2) - 1n) / this.n) * this.mu) % this.n;
    }
}

function check_gcd(primes, new_prime) {
    for (let prime of primes) {
        if (gcd(prime - 1n, new_prime - 1n) === 2n) {
            return prime;
        }
    }
    return null;
}

export function generate_keypair() {
    let p, q, pubKey, privKey;

    if (
        window.sessionStorage.getItem("p") !== null &&
        window.sessionStorage.getItem("q") !== null
    ) {
        p = BigInt(window.sessionStorage.getItem("p"));
        q = BigInt(window.sessionStorage.getItem("q"));
    } else {
        let p1 = generate_prime();
        while (p1 % 4n !== 3n) {
            p1 = generate_prime();
        }

        let primes = [p1];

        while (
            check_gcd(primes.slice(0, primes.length - 1), primes[primes.length - 1]) ===
            null
        ) {
            q = generate_prime();
            while (q % 4n !== 3n) {
                q = generate_prime();
            }
            primes.push(q);
        }

        p = check_gcd(primes.slice(0, primes.length - 1), primes[primes.length - 1]);
    }

    window.sessionStorage.setItem("p", p);
    window.sessionStorage.setItem("q", q);

    pubKey = new PaillierPubKey(p * q);
    privKey = new PaillierPrivKey(p, q);

    return { pubKey, privKey };
}

// p = a.prove(); v = p.asVerifier(); v.verify(p.prove(v.challenge));