"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    var desc = Object.getOwnPropertyDescriptor(m, k);
    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
      desc = { enumerable: true, get: function() { return m[k]; } };
    }
    Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
    var ownKeys = function(o) {
        ownKeys = Object.getOwnPropertyNames || function (o) {
            var ar = [];
            for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
            return ar;
        };
        return ownKeys(o);
    };
    return function (mod) {
        if (mod && mod.__esModule) return mod;
        var result = {};
        if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
        __setModuleDefault(result, mod);
        return result;
    };
})();
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.ProofGenerator = void 0;
/* eslint-disable no-console */
const core_1 = require("@maci-protocol/core");
const crypto_1 = require("@maci-protocol/crypto");
const fs_1 = __importDefault(require("fs"));
const path_1 = __importDefault(require("path"));
const logger_1 = require("../../ts/logger");
const proofs_1 = require("../../ts/proofs");
const utils_1 = require("../../ts/utils");
/**
 * Proof generator class for message processing and tally.
 */
class ProofGenerator {
    /**
     * Get maci state from local file or from contract
     *
     * @param {IPrepareStateParams} params - params to prepare maci state
     * @returns {MaciState} maci state
     */
    static async prepareState({ maciContract, pollContract, pollId, maciPrivateKey, coordinatorKeypair, signer, ipfsMessageBackupFiles, outputDir, options: { transactionHash, stateFile, startBlock, endBlock, blocksPerBatch }, }) {
        const isOutputDirExists = fs_1.default.existsSync(path_1.default.resolve(outputDir));
        if (!isOutputDirExists) {
            await fs_1.default.promises.mkdir(path_1.default.resolve(outputDir));
        }
        if (stateFile) {
            const content = await fs_1.default.promises
                .readFile(stateFile)
                .then((res) => JSON.parse(res.toString()));
            const serializedPrivateKey = maciPrivateKey.serialize();
            const maciState = core_1.MaciState.fromJSON(content);
            maciState.polls.forEach((poll) => {
                poll.setCoordinatorKeypair(serializedPrivateKey);
            });
            return maciState;
        }
        // build an off-chain representation of the MACI contract using data in the contract storage
        const [defaultStartBlockSignup, defaultStartBlockPoll, stateRoot, totalSignups] = await Promise.all([
            maciContract.queryFilter(maciContract.filters.SignUp(), startBlock).then((events) => events[0]?.blockNumber ?? 0),
            maciContract
                .queryFilter(maciContract.filters.DeployPoll(), startBlock)
                .then((events) => events[0]?.blockNumber ?? 0),
            maciContract.getStateTreeRoot(),
            pollContract.totalSignups(),
        ]);
        const defaultStartBlock = Math.min(defaultStartBlockPoll, defaultStartBlockSignup);
        let fromBlock = startBlock ? Number(startBlock) : defaultStartBlock;
        const defaultEndBlock = await Promise.all([
            pollContract
                .queryFilter(pollContract.filters.MergeState(stateRoot, totalSignups), fromBlock)
                .then((events) => events[events.length - 1]?.blockNumber),
        ]).then((blocks) => Math.max(...blocks));
        if (transactionHash) {
            const tx = await signer.provider.getTransaction(transactionHash);
            fromBlock = tx?.blockNumber ?? defaultStartBlock;
        }
        (0, logger_1.logMagenta)({ text: (0, logger_1.info)(`Starting to fetch logs from block ${fromBlock}`) });
        const maciContractAddress = await maciContract.getAddress();
        return Promise.resolve().then(() => __importStar(require("../../ts/generateMaciState"))).then(({ generateMaciStateFromContract }) => generateMaciStateFromContract({
            provider: signer.provider,
            address: maciContractAddress,
            coordinatorKeypair,
            pollId: BigInt(pollId),
            fromBlock,
            blocksPerRequest: blocksPerBatch,
            endBlock: endBlock || defaultEndBlock,
            ipfsMessageBackupFiles,
        }));
    }
    /**
     * Initialize class properties
     *
     * @param {IProofGeneratorParams} params - initialization params
     */
    constructor({ poll, messageProcessor, tally, rapidsnark, maciContractAddress, tallyContractAddress, outputDir, tallyOutputFile, mode, }) {
        this.poll = poll;
        this.maciContractAddress = maciContractAddress;
        this.tallyContractAddress = tallyContractAddress;
        this.outputDir = outputDir;
        this.tallyOutputFile = tallyOutputFile;
        this.messageProcessor = messageProcessor;
        this.tally = tally;
        this.rapidsnark = rapidsnark;
        this.mode = mode;
    }
    /**
     * Generate message processing proofs
     *
     * @returns message processing proofs
     */
    async generateMpProofs(options) {
        (0, logger_1.logMagenta)({ text: (0, logger_1.info)(`Generating proofs of message processing...`) });
        performance.mark("messageProcessor-proofs-start");
        const { messageBatchSize } = this.poll.batchSizes;
        const numMessages = this.poll.messages.length;
        let totalMessageBatches = numMessages <= messageBatchSize ? 1 : Math.floor(numMessages / messageBatchSize);
        if (numMessages > messageBatchSize && numMessages % messageBatchSize > 0) {
            totalMessageBatches += 1;
        }
        try {
            const inputs = [];
            const proofs = [];
            // while we have unprocessed messages, process them
            while (this.poll.hasUnprocessedMessages()) {
                // process messages in batches
                const circuitInputs = this.poll.processMessages(BigInt(this.poll.pollId));
                // generate the proof for this batch
                inputs.push(circuitInputs);
                (0, logger_1.logMagenta)({ text: (0, logger_1.info)(`Progress: ${this.poll.totalBatchesProcessed} / ${totalMessageBatches}`) });
            }
            (0, logger_1.logMagenta)({ text: (0, logger_1.info)("Wait until proof generation is finished") });
            const messageProcessorZkey = await (0, proofs_1.extractVerifyingKey)(this.messageProcessor.zkey, false);
            for (let index = 0; index < inputs.length; index += 1) {
                // eslint-disable-next-line no-await-in-loop
                const data = await this.generateProofs(inputs[index], this.messageProcessor, `process_${index}.json`, messageProcessorZkey).then((result) => {
                    options?.onBatchComplete?.({ current: index, total: totalMessageBatches, proofs: result });
                    return result;
                });
                proofs.push(...data);
            }
            (0, logger_1.logGreen)({ text: (0, logger_1.success)("Proof generation is finished") });
            // cleanup threads
            await (0, utils_1.cleanThreads)();
            performance.mark("messageProcessor-proofs-end");
            performance.measure("Generate message processor proofs", "messageProcessor-proofs-start", "messageProcessor-proofs-end");
            options?.onComplete?.(proofs);
            return proofs;
        }
        catch (error) {
            options?.onFail?.(error);
            throw error;
        }
    }
    /**
     * Generate tally proofs
     *
     * @param networkName - current network name
     * @param chainId - current chain id
     * @returns tally proofs
     */
    async generateTallyProofs(networkName, chainId, options) {
        (0, logger_1.logMagenta)({ text: (0, logger_1.info)(`Generating proofs of vote tallying...`) });
        performance.mark("tally-proofs-start");
        const { tallyBatchSize } = this.poll.batchSizes;
        const numStateLeaves = this.poll.pollStateLeaves.length;
        let totalTallyBatches = numStateLeaves <= tallyBatchSize ? 1 : Math.floor(numStateLeaves / tallyBatchSize);
        if (numStateLeaves > tallyBatchSize && numStateLeaves % tallyBatchSize > 0) {
            totalTallyBatches += 1;
        }
        try {
            let tallyCircuitInputs;
            const inputs = [];
            const proofs = [];
            while (this.poll.hasUntalliedBallots()) {
                tallyCircuitInputs = this.poll.tallyVotes();
                inputs.push(tallyCircuitInputs);
                (0, logger_1.logMagenta)({ text: (0, logger_1.info)(`Progress: ${this.poll.numBatchesTallied} / ${totalTallyBatches}`) });
            }
            (0, logger_1.logMagenta)({ text: (0, logger_1.info)("Wait until proof generation is finished") });
            const tallyVerifyingKey = await (0, proofs_1.extractVerifyingKey)(this.tally.zkey, false);
            for (let index = 0; index < inputs.length; index += 1) {
                // eslint-disable-next-line no-await-in-loop
                const data = await this.generateProofs(inputs[index], this.tally, `tally_${index}.json`, tallyVerifyingKey).then((result) => {
                    options?.onBatchComplete?.({ current: index, total: totalTallyBatches, proofs: result });
                    return result;
                });
                proofs.push(...data);
            }
            (0, logger_1.logGreen)({ text: (0, logger_1.success)("Proof generation is finished") });
            // cleanup threads
            await (0, utils_1.cleanThreads)();
            // verify the results
            // Compute newResultsCommitment
            const newResultsCommitment = (0, crypto_1.generateTreeCommitment)(this.poll.tallyResult, BigInt((0, utils_1.asHex)(tallyCircuitInputs.newResultsRootSalt)), this.poll.treeDepths.voteOptionTreeDepth);
            // compute newSpentVoiceCreditsCommitment
            const newSpentVoiceCreditsCommitment = (0, crypto_1.hashLeftRight)(this.poll.totalSpentVoiceCredits, BigInt((0, utils_1.asHex)(tallyCircuitInputs.newSpentVoiceCreditSubtotalSalt)));
            let newTallyCommitment;
            // create the tally file data to store for verification later
            const tallyFileData = {
                maci: this.maciContractAddress,
                pollId: this.poll.pollId.toString(),
                network: networkName,
                chainId,
                mode: this.mode,
                tallyAddress: this.tallyContractAddress,
                newTallyCommitment: (0, utils_1.asHex)(tallyCircuitInputs.newTallyCommitment),
                results: {
                    tally: this.poll.tallyResult.map((x) => x.toString()),
                    salt: (0, utils_1.asHex)(tallyCircuitInputs.newResultsRootSalt),
                    commitment: (0, utils_1.asHex)(newResultsCommitment),
                },
                totalSpentVoiceCredits: {
                    spent: this.poll.totalSpentVoiceCredits.toString(),
                    salt: (0, utils_1.asHex)(tallyCircuitInputs.newSpentVoiceCreditSubtotalSalt),
                    commitment: (0, utils_1.asHex)(newSpentVoiceCreditsCommitment),
                },
            };
            if (this.mode === core_1.EMode.QV) {
                // Compute newPerVoteOptionSpentVoiceCreditsCommitment
                const newPerVoteOptionSpentVoiceCreditsCommitment = (0, crypto_1.generateTreeCommitment)(this.poll.perVoteOptionSpentVoiceCredits, BigInt((0, utils_1.asHex)(tallyCircuitInputs.newPerVoteOptionSpentVoiceCreditsRootSalt)), this.poll.treeDepths.voteOptionTreeDepth);
                // update perVoteOptionSpentVoiceCredits in the tally file data
                tallyFileData.perVoteOptionSpentVoiceCredits = {
                    tally: this.poll.perVoteOptionSpentVoiceCredits.map((x) => x.toString()),
                    salt: (0, utils_1.asHex)(tallyCircuitInputs.newPerVoteOptionSpentVoiceCreditsRootSalt),
                    commitment: (0, utils_1.asHex)(newPerVoteOptionSpentVoiceCreditsCommitment),
                };
                // Compute newTallyCommitment
                newTallyCommitment = (0, crypto_1.hash3)([
                    newResultsCommitment,
                    newSpentVoiceCreditsCommitment,
                    newPerVoteOptionSpentVoiceCreditsCommitment,
                ]);
            }
            else {
                newTallyCommitment = (0, crypto_1.hashLeftRight)(newResultsCommitment, newSpentVoiceCreditsCommitment);
            }
            await fs_1.default.promises.writeFile(this.tallyOutputFile, JSON.stringify(tallyFileData, null, 4));
            (0, logger_1.logGreen)({ text: (0, logger_1.success)(`Tally file:\n${JSON.stringify(tallyFileData, null, 4)}\n`) });
            // compare the commitments
            if ((0, utils_1.asHex)(newTallyCommitment) === tallyFileData.newTallyCommitment) {
                (0, logger_1.logGreen)({ text: (0, logger_1.success)("The tally commitment is correct") });
            }
            else {
                throw new Error("Error: the newTallyCommitment is invalid.");
            }
            performance.mark("tally-proofs-end");
            performance.measure("Generate tally proofs", "tally-proofs-start", "tally-proofs-end");
            options?.onComplete?.(proofs, tallyFileData);
            return { proofs, tallyData: tallyFileData };
        }
        catch (error) {
            options?.onFail?.(error);
            throw error;
        }
    }
    /**
     * Generic function for proofs generation
     *
     * @param {TCircuitInputs} circuitInputs - circuit inputs
     * @param {ICircuitFiles} circuitFiles - circuit files (zkey, witnessGenerator, wasm)
     * @param outputFile - output file
     * @returns proofs
     */
    async generateProofs(circuitInputs, circuitFiles, outputFile, verifyingKey) {
        const proofs = [];
        const { proof, publicSignals } = circuitFiles.wasm
            ? await (0, proofs_1.generateProofSnarkjs)({
                inputs: circuitInputs,
                zkeyPath: circuitFiles.zkey,
                silent: true,
                wasmPath: circuitFiles.wasm,
            })
            : await (0, proofs_1.generateProofRapidSnark)({
                inputs: circuitInputs,
                zkeyPath: circuitFiles.zkey,
                rapidsnarkExePath: this.rapidsnark,
                witnessExePath: circuitFiles.witnessGenerator,
            });
        // verify it
        // eslint-disable-next-line no-await-in-loop
        const isValid = await (0, proofs_1.verifyProof)(publicSignals, proof, verifyingKey, false);
        if (!isValid) {
            throw new Error("Error: generated an invalid proof");
        }
        proofs.push({
            circuitInputs,
            proof,
            publicInputs: publicSignals,
        });
        await fs_1.default.promises.writeFile(path_1.default.resolve(this.outputDir, outputFile), JSON.stringify(proofs[proofs.length - 1], null, 4));
        return proofs;
    }
}
exports.ProofGenerator = ProofGenerator;
//# sourceMappingURL=ProofGenerator.js.map