import { NO_WAIT, getContractInstanceFromInstantiationParams } from '@aztec/aztec.js/contracts';
import { Fr } from '@aztec/aztec.js/fields';
import { EpochNumber, SlotNumber } from '@aztec/foundation/branded-types';
import { timesAsync, unique } from '@aztec/foundation/collection';
import { retryUntil } from '@aztec/foundation/retry';
import { pluralize } from '@aztec/foundation/string';
import { TestContract, TestContractArtifact } from '@aztec/noir-test-contracts.js/Test';
import { getPXEConfig, getPXEConfig as getRpcConfig } from '@aztec/pxe/server';
import { getRoundForOffense } from '@aztec/slasher';
import { submitTxsTo } from '../shared/submit-transactions.js';
import { TestWallet } from '../test-wallet/test_wallet.js';
import { proveInteraction } from '../test-wallet/utils.js';
// submits a set of transactions to the provided Private eXecution Environment (PXE)
export const submitComplexTxsTo = async (logger, from, spamContract, numTxs, opts = {})=>{
    const txs = [];
    const seed = 1234n;
    const spamCount = 15;
    for(let i = 0; i < numTxs; i++){
        const method = spamContract.methods.spam(seed + BigInt(i * spamCount), spamCount, !!opts.callPublic);
        const { txHash } = await method.send({
            from,
            wait: NO_WAIT
        });
        logger.info(`Tx sent with hash ${txHash.toString()}`);
        txs.push(txHash);
    }
    return txs;
};
// creates a wallet and submit a given number of transactions through it.
export const submitTransactions = async (logger, node, numTxs, fundedAccount)=>{
    const rpcConfig = getRpcConfig();
    rpcConfig.proverEnabled = false;
    const wallet = await TestWallet.create(node, {
        ...getPXEConfig(),
        proverEnabled: false
    }, {
        loggerActorLabel: 'pxe-tx'
    });
    const fundedAccountManager = await wallet.createSchnorrAccount(fundedAccount.secret, fundedAccount.salt);
    return submitTxsTo(wallet, fundedAccountManager.address, numTxs, logger);
};
export async function prepareTransactions(logger, node, numTxs, fundedAccount) {
    const rpcConfig = getRpcConfig();
    rpcConfig.proverEnabled = false;
    const wallet = await TestWallet.create(node, {
        ...getPXEConfig(),
        proverEnabled: false
    }, {
        loggerActorLabel: 'pxe-tx'
    });
    const fundedAccountManager = await wallet.createSchnorrAccount(fundedAccount.secret, fundedAccount.salt);
    const testContractInstance = await getContractInstanceFromInstantiationParams(TestContractArtifact, {
        salt: Fr.random()
    });
    await wallet.registerContract(testContractInstance, TestContractArtifact);
    const contract = TestContract.at(testContractInstance.address, wallet);
    return timesAsync(numTxs, async ()=>{
        const tx = await proveInteraction(wallet, contract.methods.emit_nullifier(Fr.random()), {
            from: fundedAccountManager.address
        });
        logger.info(`Tx prepared with hash ${tx.getTxHash()}`);
        return tx;
    });
}
export function awaitProposalExecution(slashingProposer, timeoutSeconds, logger) {
    return new Promise((resolve, reject)=>{
        const timeout = setTimeout(()=>{
            logger.warn(`Timed out waiting for proposal execution`);
            reject(new Error(`Timeout waiting for proposal execution after ${timeoutSeconds}s`));
        }, timeoutSeconds * 1000);
        if (slashingProposer.type === 'empire') {
            const unwatch = slashingProposer.listenToPayloadSubmitted((args)=>{
                logger.warn(`Proposal ${args.payload} from round ${args.round} executed`);
                clearTimeout(timeout);
                unwatch();
                resolve(args.round);
            });
        } else if (slashingProposer.type === 'tally') {
            const unwatch = slashingProposer.listenToRoundExecuted((args)=>{
                logger.warn(`Slash from round ${args.round} executed`);
                clearTimeout(timeout);
                unwatch();
                resolve(args.round);
            });
        } else {
            clearTimeout(timeout);
            reject(new Error(`Unknown slashing proposer type: ${slashingProposer.type}`));
        }
    });
}
export async function awaitCommitteeExists({ rollup, logger }) {
    logger.info(`Waiting for committee to be set`);
    let committee;
    await retryUntil(async ()=>{
        committee = await rollup.getCurrentEpochCommittee();
        return committee && committee.length > 0;
    }, 'non-empty committee', 60);
    logger.warn(`Committee has been formed`, {
        committee: committee.map((c)=>c.toString())
    });
    return committee.map((c)=>c.toString());
}
/**
 * Advance epochs until we find one where the target proposer is selected for at least one slot,
 * then stop one epoch before it. This leaves time for the caller to start sequencers before
 * warping to the target epoch, avoiding the race where the target epoch passes before sequencers
 * are ready.
 *
 * Returns the target epoch number so the caller can warp to it after starting sequencers.
 */ export async function advanceToEpochBeforeProposer({ epochCache, cheatCodes, targetProposer, logger, maxAttempts = 20 }) {
    const { epochDuration } = await cheatCodes.getConfig();
    for(let attempt = 0; attempt < maxAttempts; attempt++){
        const currentEpoch = await cheatCodes.getEpoch();
        // Check the NEXT epoch's slots so we stay one epoch before the target,
        // giving the caller time to start sequencers before the target epoch arrives.
        const nextEpoch = Number(currentEpoch) + 1;
        const startSlot = nextEpoch * Number(epochDuration);
        const endSlot = startSlot + Number(epochDuration);
        logger.info(`Checking next epoch ${nextEpoch} (slots ${startSlot}-${endSlot - 1}) for proposer ${targetProposer} (current epoch: ${currentEpoch})`);
        for(let s = startSlot; s < endSlot; s++){
            const proposer = await epochCache.getProposerAttesterAddressInSlot(SlotNumber(s));
            if (proposer && proposer.equals(targetProposer)) {
                logger.warn(`Found target proposer ${targetProposer} in slot ${s} of epoch ${nextEpoch}. Staying at epoch ${currentEpoch} to allow sequencer startup.`);
                return {
                    targetEpoch: EpochNumber(nextEpoch)
                };
            }
        }
        logger.info(`Target proposer not found in epoch ${nextEpoch}, advancing to next epoch`);
        await cheatCodes.advanceToNextEpoch();
    }
    throw new Error(`Target proposer ${targetProposer} not found in any slot after ${maxAttempts} epoch attempts`);
}
/**
 * Advance epochs until we find one where the target proposer is selected for at least one slot.
 * With N validators and M slots per epoch, a specific proposer may not be selected in any given epoch.
 * For example, with 4 validators and 2 slots/epoch, there is about a 44% chance per epoch.
 */ export async function awaitEpochWithProposer({ epochCache, cheatCodes, targetProposer, logger, maxAttempts = 20 }) {
    const { epochDuration } = await cheatCodes.getConfig();
    for(let attempt = 0; attempt < maxAttempts; attempt++){
        const currentEpoch = await cheatCodes.getEpoch();
        const startSlot = Number(currentEpoch) * Number(epochDuration);
        const endSlot = startSlot + Number(epochDuration);
        logger.info(`Checking epoch ${currentEpoch} (slots ${startSlot}-${endSlot - 1}) for proposer ${targetProposer}`);
        for(let s = startSlot; s < endSlot; s++){
            const proposer = await epochCache.getProposerAttesterAddressInSlot(SlotNumber(s));
            if (proposer && proposer.equals(targetProposer)) {
                logger.warn(`Found target proposer ${targetProposer} in slot ${s} of epoch ${currentEpoch}`);
                return;
            }
        }
        logger.info(`Target proposer not found in epoch ${currentEpoch}, advancing to next epoch`);
        await cheatCodes.advanceToNextEpoch();
    }
    throw new Error(`Target proposer ${targetProposer} not found in any slot after ${maxAttempts} epoch attempts`);
}
export async function awaitOffenseDetected({ logger, nodeAdmin, slashingRoundSize, epochDuration, waitUntilOffenseCount, timeoutSeconds = 120 }) {
    const targetOffenseCount = waitUntilOffenseCount ?? 1;
    logger.warn(`Waiting for ${pluralize('offense', targetOffenseCount)} to be detected`);
    const offenses = await retryUntil(async ()=>{
        const offenses = await nodeAdmin.getSlashOffenses('all');
        if (offenses.length >= targetOffenseCount) {
            return offenses;
        }
    }, 'non-empty offenses', timeoutSeconds);
    logger.info(`Hit ${offenses.length} offenses on rounds ${unique(offenses.map((o)=>getRoundForOffense(o, {
            slashingRoundSize,
            epochDuration
        })))}`, {
        offenses
    });
    return offenses;
}
/**
 * Await the committee to be slashed out of the validator set.
 * Currently assumes that the committee is the same size as the validator set.
 */ export async function awaitCommitteeKicked({ rollup, cheatCodes, committee, slashFactory, slashingProposer, slashingRoundSize, aztecSlotDuration, aztecEpochDuration, logger, offenseEpoch }) {
    if (!slashingProposer) {
        throw new Error('No slashing proposer configured. Cannot test slashing.');
    }
    await cheatCodes.debugRollup();
    if (slashingProposer.type === 'empire') {
        // Await for the slash payload to be created if empire (no payload is created on tally until execution time)
        const targetEpoch = EpochNumber(await cheatCodes.getEpoch() + await rollup.getLagInEpochsForValidatorSet() + 1);
        logger.info(`Advancing to epoch ${targetEpoch} so we start slashing`);
        await cheatCodes.advanceToEpoch(targetEpoch);
        const slashPayloadEvents = await retryUntil(async ()=>{
            const events = await slashFactory.getSlashPayloadCreatedEvents();
            return events.length > 0 ? events : undefined;
        }, 'slash payload created', 120, 1);
        expect(slashPayloadEvents.length).toBe(1);
        // The uniqueness check is needed since a validator may be slashed more than once on the same round (eg because they let two epochs be pruned)
        expect(unique(slashPayloadEvents[0].slashes.map((slash)=>slash.validator.toString()))).toHaveLength(committee.length);
    } else {
        // Use the slash offset to ensure we are in the right epoch for tally
        const slashOffsetInRounds = await slashingProposer.getSlashOffsetInRounds();
        const slashingRoundSizeInEpochs = slashingRoundSize / aztecEpochDuration;
        const slashingOffsetInEpochs = Number(slashOffsetInRounds) * slashingRoundSizeInEpochs;
        const firstEpochInOffenseRound = offenseEpoch - offenseEpoch % slashingRoundSizeInEpochs;
        const targetEpoch = firstEpochInOffenseRound + slashingOffsetInEpochs;
        logger.info(`Advancing to epoch ${targetEpoch} so we start slashing`);
        await cheatCodes.advanceToEpoch(EpochNumber(targetEpoch), {
            offset: -aztecSlotDuration / 2
        });
    }
    const attestersPre = await rollup.getAttesters();
    expect(attestersPre.length).toBe(committee.length);
    for (const attester of attestersPre){
        const attesterInfo = await rollup.getAttesterView(attester);
        expect(attesterInfo.status).toEqual(1); // Validating
    }
    const timeout = slashingRoundSize * 2 * aztecSlotDuration + 30;
    logger.info(`Waiting for slash to be executed (timeout ${timeout}s)`);
    await awaitProposalExecution(slashingProposer, timeout, logger);
    // The attesters should still form the committee but they should be reduced to the "living" status
    await cheatCodes.debugRollup();
    const committeePostSlashing = await rollup.getCurrentEpochCommittee();
    expect(committeePostSlashing?.length).toBe(attestersPre.length);
    const attestersPostSlashing = await rollup.getAttesters();
    expect(attestersPostSlashing.length).toBe(0);
    for (const attester of attestersPre){
        const attesterInfo = await rollup.getAttesterView(attester);
        expect(attesterInfo.status).toEqual(2); // Living
    }
    logger.info(`Advancing to check current committee`);
    await cheatCodes.debugRollup();
    await cheatCodes.advanceToEpoch(EpochNumber(await cheatCodes.getEpoch() + await rollup.getLagInEpochsForValidatorSet() + 1));
    await cheatCodes.debugRollup();
    const committeeNextEpoch = await rollup.getCurrentEpochCommittee();
    // The committee should be undefined, since the validator set is empty
    // and the tests currently using this helper always set a target committee size.
    expect(committeeNextEpoch).toBeUndefined();
    const attestersNextEpoch = await rollup.getAttesters();
    expect(attestersNextEpoch.length).toBe(0);
}
