import { getSchnorrWalletWithSecretKey } from '@aztec/accounts/schnorr';
import { EthAddress, Fr } from '@aztec/aztec.js';
import { GSEContract, MultiAdderArtifact, RollupContract, createL1TxUtilsFromViemWallet, deployL1Contract, getL1ContractsConfigEnvVars } from '@aztec/ethereum';
import { ChainMonitor } from '@aztec/ethereum/test';
import { SecretValue } from '@aztec/foundation/config';
import { createLogger } from '@aztec/foundation/log';
import { RollupAbi, SlasherAbi, TestERC20Abi } from '@aztec/l1-artifacts';
import { SpamContract } from '@aztec/noir-test-contracts.js/Spam';
import { createBootstrapNodeFromPrivateKey, getBootstrapNodeEnr } from '@aztec/p2p/test-helpers';
import { tryStop } from '@aztec/stdlib/interfaces/server';
import { SlashFactoryContract } from '@aztec/stdlib/l1-contracts';
import { ZkPassportProofParams } from '@aztec/stdlib/zkpassport';
import { getGenesisValues } from '@aztec/world-state/testing';
import getPort from 'get-port';
import { getAddress, getContract } from 'viem';
import { privateKeyToAccount } from 'viem/accounts';
import { ATTESTER_PRIVATE_KEYS_START_INDEX, createValidatorConfig, generatePrivateKeys } from '../fixtures/setup_p2p_test.js';
import { createSnapshotManager, deployAccounts } from '../fixtures/snapshot_manager.js';
import { getPrivateKeyFromIndex, getSponsoredFPCAddress } from '../fixtures/utils.js';
import { getEndToEndTestTelemetryClient } from '../fixtures/with_telemetry_utils.js';
// Use a fixed bootstrap node private key so that we can re-use the same snapshot and the nodes can find each other
const BOOTSTRAP_NODE_PRIVATE_KEY = '080212208f988fc0899e4a73a5aee4d271a5f20670603a756ad8d84f2c94263a6427c591';
const l1ContractsConfig = getL1ContractsConfigEnvVars();
export const WAIT_FOR_TX_TIMEOUT = l1ContractsConfig.aztecSlotDuration * 3;
export const SHORTENED_BLOCK_TIME_CONFIG_NO_PRUNES = {
    aztecSlotDuration: 12,
    ethereumSlotDuration: 4,
    aztecProofSubmissionWindow: 640
};
export class P2PNetworkTest {
    testName;
    bootstrapNodeEnr;
    bootNodePort;
    numberOfValidators;
    numberOfNodes;
    metricsPort;
    snapshotManager;
    baseAccountPrivateKey;
    baseAccount;
    logger;
    monitor;
    ctx;
    attesterPrivateKeys;
    attesterPublicKeys;
    peerIdPrivateKeys;
    validators;
    deployedAccounts;
    prefilledPublicData;
    // The re-execution test needs a wallet and a spam contract
    wallet;
    defaultAccountAddress;
    spamContract;
    bootstrapNode;
    constructor(testName, bootstrapNodeEnr, bootNodePort, numberOfValidators, initialValidatorConfig, numberOfNodes = 0, // If set enable metrics collection
    metricsPort, startProverNode, mockZkPassportVerifier){
        this.testName = testName;
        this.bootstrapNodeEnr = bootstrapNodeEnr;
        this.bootNodePort = bootNodePort;
        this.numberOfValidators = numberOfValidators;
        this.numberOfNodes = numberOfNodes;
        this.metricsPort = metricsPort;
        this.attesterPrivateKeys = [];
        this.attesterPublicKeys = [];
        this.peerIdPrivateKeys = [];
        this.validators = [];
        this.deployedAccounts = [];
        this.prefilledPublicData = [];
        this.logger = createLogger(`e2e:e2e_p2p:${testName}`);
        // Set up the base account and node private keys for the initial network deployment
        this.baseAccountPrivateKey = `0x${getPrivateKeyFromIndex(1).toString('hex')}`;
        this.baseAccount = privateKeyToAccount(this.baseAccountPrivateKey);
        this.attesterPrivateKeys = generatePrivateKeys(ATTESTER_PRIVATE_KEYS_START_INDEX + numberOfNodes, numberOfValidators);
        this.attesterPublicKeys = this.attesterPrivateKeys.map((privateKey)=>privateKeyToAccount(privateKey).address);
        const zkPassportParams = ZkPassportProofParams.random();
        this.snapshotManager = createSnapshotManager(`e2e_p2p_network/${testName}`, process.env.E2E_DATA_PATH, {
            ...initialValidatorConfig,
            ethereumSlotDuration: initialValidatorConfig.ethereumSlotDuration ?? l1ContractsConfig.ethereumSlotDuration,
            aztecEpochDuration: initialValidatorConfig.aztecEpochDuration ?? l1ContractsConfig.aztecEpochDuration,
            aztecSlotDuration: initialValidatorConfig.aztecSlotDuration ?? l1ContractsConfig.aztecSlotDuration,
            aztecProofSubmissionEpochs: initialValidatorConfig.aztecProofSubmissionEpochs ?? l1ContractsConfig.aztecProofSubmissionEpochs,
            slashingRoundSizeInEpochs: initialValidatorConfig.slashingRoundSizeInEpochs ?? l1ContractsConfig.slashingRoundSizeInEpochs,
            slasherFlavor: initialValidatorConfig.slasherFlavor ?? 'tally',
            aztecTargetCommitteeSize: numberOfValidators,
            salt: 420,
            metricsPort: metricsPort,
            numberOfInitialFundedAccounts: 2,
            startProverNode
        }, {
            ...initialValidatorConfig,
            aztecEpochDuration: initialValidatorConfig.aztecEpochDuration ?? l1ContractsConfig.aztecEpochDuration,
            slashingRoundSizeInEpochs: initialValidatorConfig.slashingRoundSizeInEpochs ?? l1ContractsConfig.slashingRoundSizeInEpochs,
            slasherFlavor: initialValidatorConfig.slasherFlavor ?? 'tally',
            ethereumSlotDuration: initialValidatorConfig.ethereumSlotDuration ?? l1ContractsConfig.ethereumSlotDuration,
            aztecSlotDuration: initialValidatorConfig.aztecSlotDuration ?? l1ContractsConfig.aztecSlotDuration,
            aztecProofSubmissionEpochs: initialValidatorConfig.aztecProofSubmissionEpochs ?? l1ContractsConfig.aztecProofSubmissionEpochs,
            aztecTargetCommitteeSize: numberOfValidators,
            initialValidators: [],
            zkPassportArgs: {
                mockZkPassportVerifier,
                zkPassportDomain: zkPassportParams.domain,
                zkPassportScope: zkPassportParams.scope
            }
        });
    }
    static async create({ testName, numberOfNodes, numberOfValidators, basePort, metricsPort, initialConfig, startProverNode, mockZkPassportVerifier }) {
        const port = basePort || await getPort();
        const bootstrapNodeENR = await getBootstrapNodeEnr(BOOTSTRAP_NODE_PRIVATE_KEY, port);
        const bootstrapNodeEnr = bootstrapNodeENR.encodeTxt();
        const initialValidatorConfig = await createValidatorConfig(initialConfig ?? {}, bootstrapNodeEnr);
        return new P2PNetworkTest(testName, bootstrapNodeEnr, port, numberOfValidators, initialValidatorConfig, numberOfNodes, metricsPort, startProverNode, mockZkPassportVerifier);
    }
    get fundedAccount() {
        if (!this.deployedAccounts[0]) {
            throw new Error('Call snapshot t.setupAccount to create a funded account.');
        }
        return this.deployedAccounts[0];
    }
    async addBootstrapNode() {
        await this.snapshotManager.snapshot('add-bootstrap-node', async ({ aztecNodeConfig })=>{
            const telemetry = getEndToEndTestTelemetryClient(this.metricsPort);
            this.bootstrapNode = await createBootstrapNodeFromPrivateKey(BOOTSTRAP_NODE_PRIVATE_KEY, this.bootNodePort, telemetry, aztecNodeConfig);
            // Overwrite enr with updated info
            this.bootstrapNodeEnr = this.bootstrapNode.getENR().encodeTxt();
        });
    }
    getValidators() {
        const validators = [];
        for(let i = 0; i < this.numberOfValidators; i++){
            const keyIndex = i;
            const attester = privateKeyToAccount(this.attesterPrivateKeys[keyIndex]);
            validators.push({
                attester: EthAddress.fromString(attester.address),
                withdrawer: EthAddress.fromString(attester.address),
                bn254SecretKey: new SecretValue(Fr.random().toBigInt())
            });
            this.logger.info(`Adding attester ${attester.address} as validator`);
        }
        return {
            validators
        };
    }
    async applyBaseSnapshots() {
        await this.addBootstrapNode();
        await this.snapshotManager.snapshot('add-validators', async ({ deployL1ContractsValues, dateProvider, cheatCodes })=>{
            const rollup = getContract({
                address: deployL1ContractsValues.l1ContractAddresses.rollupAddress.toString(),
                abi: RollupAbi,
                client: deployL1ContractsValues.l1Client
            });
            this.logger.info(`Adding ${this.numberOfValidators} validators`);
            const stakingAsset = getContract({
                address: deployL1ContractsValues.l1ContractAddresses.stakingAssetAddress.toString(),
                abi: TestERC20Abi,
                client: deployL1ContractsValues.l1Client
            });
            const { address: multiAdderAddress } = await deployL1Contract(deployL1ContractsValues.l1Client, MultiAdderArtifact.contractAbi, MultiAdderArtifact.contractBytecode, [
                rollup.address,
                deployL1ContractsValues.l1Client.account.address
            ]);
            const multiAdder = getContract({
                address: multiAdderAddress.toString(),
                abi: MultiAdderArtifact.contractAbi,
                client: deployL1ContractsValues.l1Client
            });
            const stakeNeeded = await rollup.read.getActivationThreshold() * BigInt(this.numberOfValidators);
            await Promise.all([
                await stakingAsset.write.mint([
                    multiAdder.address,
                    stakeNeeded
                ], {})
            ].map((txHash)=>deployL1ContractsValues.l1Client.waitForTransactionReceipt({
                    hash: txHash
                })));
            const { validators } = this.getValidators();
            this.validators = validators;
            const gseAddress = deployL1ContractsValues.l1ContractAddresses.gseAddress;
            if (!gseAddress) {
                throw new Error('GSE contract not deployed');
            }
            const gseContract = new GSEContract(deployL1ContractsValues.l1Client, gseAddress.toString());
            const makeValidatorTuples = async (validator)=>{
                const registrationTuple = await gseContract.makeRegistrationTuple(validator.bn254SecretKey.getValue());
                return {
                    attester: validator.attester.toString(),
                    withdrawer: validator.withdrawer.toString(),
                    ...registrationTuple
                };
            };
            const validatorTuples = await Promise.all(validators.map(makeValidatorTuples));
            await deployL1ContractsValues.l1Client.waitForTransactionReceipt({
                hash: await multiAdder.write.addValidators([
                    validatorTuples
                ])
            });
            await cheatCodes.rollup.advanceToEpoch(await cheatCodes.rollup.getEpoch() + await rollup.read.getLagInEpochs() + 1n, {
                updateDateProvider: dateProvider
            });
            // Send and await a tx to make sure we mine a block for the warp to correctly progress.
            await this._sendDummyTx(deployL1ContractsValues.l1Client);
        });
    }
    async setupAccount() {
        await this.snapshotManager.snapshot('setup-account', deployAccounts(1, this.logger, false), async ({ deployedAccounts }, { pxe })=>{
            this.deployedAccounts = deployedAccounts;
            const [account] = deployedAccounts;
            this.wallet = await getSchnorrWalletWithSecretKey(pxe, account.secret, account.signingKey, account.salt);
            this.defaultAccountAddress = this.wallet.getAddress();
        });
    }
    async deploySpamContract() {
        await this.snapshotManager.snapshot('add-spam-contract', async ()=>{
            if (!this.wallet) {
                throw new Error('Call snapshot t.setupAccount before deploying account contract');
            }
            const spamContract = await SpamContract.deploy(this.wallet).send({
                from: this.defaultAccountAddress
            }).deployed();
            return {
                contractAddress: spamContract.address
            };
        }, async ({ contractAddress })=>{
            if (!this.wallet) {
                throw new Error('Call snapshot t.setupAccount before deploying account contract');
            }
            this.spamContract = await SpamContract.at(contractAddress, this.wallet);
        });
    }
    async removeInitialNode() {
        await this.snapshotManager.snapshot('remove-initial-validator', async ({ deployL1ContractsValues, aztecNode, dateProvider })=>{
            // Send and await a tx to make sure we mine a block for the warp to correctly progress.
            const { receipt } = await this._sendDummyTx(deployL1ContractsValues.l1Client);
            const block = await deployL1ContractsValues.l1Client.getBlock({
                blockNumber: receipt.blockNumber
            });
            dateProvider.setTime(Number(block.timestamp) * 1000);
            await aztecNode.stop();
        });
    }
    async sendDummyTx() {
        return await this._sendDummyTx(this.ctx.deployL1ContractsValues.l1Client);
    }
    async _sendDummyTx(l1Client) {
        const l1TxUtils = createL1TxUtilsFromViemWallet(l1Client);
        return await l1TxUtils.sendAndMonitorTransaction({
            to: l1Client.account.address,
            value: 1n
        });
    }
    async setup() {
        this.ctx = await this.snapshotManager.setup();
        const sponsoredFPCAddress = await getSponsoredFPCAddress();
        const initialFundedAccounts = [
            ...this.ctx.initialFundedAccounts.map((a)=>a.address),
            sponsoredFPCAddress
        ];
        const { prefilledPublicData } = await getGenesisValues(initialFundedAccounts);
        this.prefilledPublicData = prefilledPublicData;
        const rollupContract = RollupContract.getFromL1ContractsValues(this.ctx.deployL1ContractsValues);
        this.monitor = new ChainMonitor(rollupContract, this.ctx.dateProvider).start();
        this.monitor.on('l1-block', ({ timestamp })=>this.ctx.dateProvider.setTime(Number(timestamp) * 1000));
    }
    async stopNodes(nodes) {
        this.logger.info('Stopping nodes');
        if (!nodes || !nodes.length) {
            this.logger.info('No nodes to stop');
            return;
        }
        await Promise.all(nodes.map((node)=>node.stop()));
        this.logger.info('Nodes stopped');
    }
    async teardown() {
        await this.monitor.stop();
        await tryStop(this.bootstrapNode, this.logger);
        await this.snapshotManager.teardown();
    }
    async getContracts() {
        if (!this.ctx.deployL1ContractsValues) {
            throw new Error('DeployL1ContractsValues not set');
        }
        const rollup = new RollupContract(this.ctx.deployL1ContractsValues.l1Client, this.ctx.deployL1ContractsValues.l1ContractAddresses.rollupAddress);
        const slasherContract = getContract({
            address: getAddress(await rollup.getSlasherAddress()),
            abi: SlasherAbi,
            client: this.ctx.deployL1ContractsValues.l1Client
        });
        // Get the actual slashing proposer from rollup (which handles both empire and tally)
        const slashingProposer = await rollup.getSlashingProposer();
        const slashFactory = new SlashFactoryContract(this.ctx.deployL1ContractsValues.l1Client, getAddress(this.ctx.deployL1ContractsValues.l1ContractAddresses.slashFactoryAddress.toString()));
        return {
            rollup,
            slasherContract,
            slashingProposer,
            slashFactory
        };
    }
}
