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

const PAILLIER = 0;
const JURIK = 1;

function RSTransform(g, a, p) {
    let plainText = p.toString(16);
    if (plainText.length % 2 !== 0) {
        plainText = "0" + plainText;
    }

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

    let hasher = new jsSHA("SHAKE256", "HEX");
    hasher.update(g.toString(16));
    hasher.update(plainText);
    hasher.update(aStr);

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

class Ciphertext {
    constructor(key, plainText, r, set) {
        while (plainText < 0n) {
            plainText += key.n2;
        }

        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 * key.hn_exp(r)) % key.n2;

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

            this.mode = JURIK;
            this.r = key.h_exp(r);
        } 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 = (this.plainText + c.plainText) % this.pubKey.n2;

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

    mul(e) {
        this.cipherText = mod_exp(this.cipherText, e, this.pubKey.n2);
        this.r = mod_exp(this.r, e, this.pubKey.n2);
        this.plainText = (this.plainText * e) % this.pubKey.n2;

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

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

    toJSON() {
        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 challenge = RSTransform(this.pubKey.g, a, this.plainText);

        return {
            plainText: "0x" + this.plainText.toString(16),
            a: "0x" + a.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) {
        if (typeof cipherText !== "bigint") {
            throw "ReadOnlyCiphertext must take BigInt parameter";
        }

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

    mul(e) {
        this.cipherText = mod_exp(this.cipherText, e, 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 challenge = RSTransform(
            this.pubKey.g,
            BigInt(statement.a),
            BigInt(statement.plainText)
        );

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

        if (verifier.verify(BigInt(statement.proof))) {
            return BigInt(statement.plainText);
        } else {
            return null;
        }
    }

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

        this._h_cache = [];
        this._hn_cache = [];

        // Browser dies on higher key sizes :P
        if (KEY_SIZE <= 1024) {
            for (let i = 0n; i < BigInt(KEY_SIZE); i++) {
                this._h_cache.push(mod_exp(this.h, 2n ** i, this.n));
                this._hn_cache.push(mod_exp(this.hn, 2n ** i, this.n2));
            }
        }
    }

    h_exp(b) {
        if (KEY_SIZE > 1024) {
            return mod_exp(this.h, b, this.n);
        }

        let ctr = 1n;
        let i = 0;
        while (b !== 0n) {
            if (b % 2n === 1n) {
                ctr *= this._h_cache[i];
                ctr %= this.n;
            }
            i++;
            b >>= 1n;
        }

        return ctr;
    }

    hn_exp(b) {
        if (KEY_SIZE > 1024) {
            return mod_exp(this.hn, b, this.n2);
        }

        let ctr = 1n;
        let i = 0;
        while (b !== 0n) {
            if (b % 2n === 1n) {
                ctr *= this._hn_cache[i];
                ctr %= this.n2;
            }
            i++;
            b >>= 1n;
        }

        return ctr;
    }

    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 {
        p = generate_safe_prime();
        q = generate_safe_prime();
    }

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