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

class Cyphertext {
    constructor(key, plainText, r) {
        if (r === undefined) {
            r = cryptoRandom(4096);

            // Resample to avoid modulo bias.
            while (r >= key.n) {
                r = cryptoRandom(4096);
            }
        }

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

        // Compute g^m r^n from crt
        this.cyphertext = (gm * mod_exp(r, key.n, key.n ** 2n)) % key.n ** 2n;
        this.r = r;
        this.pubKey = key;
        this.plainText = plainText;

        this.readOnly = false;
    }

    update(c) {
        this.cyphertext = (this.cyphertext * c.cyphertext) % this.pubKey.n ** 2n;
        this.r = (this.r * c.r) % this.pubKey.n ** 2n;
        this.plainText += c.plainText;
    }

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

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

    asReadOnlyCyphertext() {
        return new ReadOnlyCyphertext(this.pubKey, this.cyphertext);
    }
}

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

        this.rp = cryptoRandom(4096);
        while (this.rp >= this.cipherText.pubKey.n) {
            this.rp = cryptoRandom(4096);
        }
    }

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

    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
            .asReadOnlyCyphertext()
            .prove(this.cipherText.plainText, this.noise());
    }
}

window.Cyphertext = Cyphertext;

export class ReadOnlyCyphertext {
    constructor(key, cyphertext) {
        this.cyphertext = cyphertext;
        this.pubKey = key;

        this.readOnly = true;
    }

    update(c) {
        this.cyphertext = (this.cyphertext * c.cyphertext) % this.pubKey.n ** 2n;
    }

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

class ProofSessionVerifier {
    constructor(cipherText, plainText, a) {
        this.cipherText = cipherText;
        this.cipherText.update(this.cipherText.pubKey.encrypt(-1n * plainText, 1n));
        // Shift the challenge down by 1 to ensure it is smaller than either prime factor.
        this.challenge = cryptoRandom(2048) << 1n;
        this.a = a;
    }

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

        // check exp
        return (
            mod_exp(proof, this.cipherText.pubKey.n, this.cipherText.pubKey.n ** 2n) ===
            (this.a *
                mod_exp(
                    this.cipherText.cyphertext,
                    this.challenge,
                    this.cipherText.pubKey.n ** 2n
                )) %
                this.cipherText.pubKey.n ** 2n
        );
    }
}

window.ReadOnlyCyphertext = ReadOnlyCyphertext;

export class PaillierPubKey {
    constructor(n) {
        this.n = n;
        this.g = this.n + 1n;
    }

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

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

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

class PaillierPrivKey {
    constructor(p, q) {
        this.n = p * q;
        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.n ** 2n) - 1n) / this.n) * this.mu) % this.n
        );
    }
}

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

    if (window.sessionStorage.getItem("p") === null) {
        p = generate_prime();
        window.sessionStorage.setItem("p", p);
    } else {
        p = BigInt(window.sessionStorage.getItem("p"));
    }

    if (window.sessionStorage.getItem("q") === null) {
        q = generate_prime();
        window.sessionStorage.setItem("q", q);
    } else {
        q = BigInt(window.sessionStorage.getItem("q"));
    }

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

    return { pubKey, privKey };
}