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

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

    encrypt(m) {
        // Compute g^m r^n mod n^2
        let r = random2048();

        // Resample to avoid modulo bias.
        while (r >= this.n) {
            r = random2048();
        }

        // Compute g^m by binomial theorem.
        let gm = (1n + this.n * m) % this.n ** 2n;
        // Compute g^m r^n from crt
        return (gm * mod_exp(r, this.n, this.n ** 2n)) % this.n ** 2n;
    }

    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 };
}