import { SchnorrAccountContractArtifact } from '@aztec/accounts/schnorr';
import { deployFundedSchnorrAccounts, generateSchnorrAccounts } from '@aztec/accounts/testing';
import { AztecNodeService, getConfigEnvVars } from '@aztec/aztec-node';
import { BatchCall, DefaultWaitForProvenOpts, EthAddress, getContractClassFromArtifact, waitForProven } from '@aztec/aztec.js';
import { publishContractClass, publishInstance } from '@aztec/aztec.js/deployment';
import { AnvilTestWatcher, CheatCodes } from '@aztec/aztec/testing';
import { createBlobSinkServer } from '@aztec/blob-sink/server';
import { createExtendedL1Client, deployMulticall3, getL1ContractsConfigEnvVars } from '@aztec/ethereum';
import { EthCheatCodesWithState, startAnvil } from '@aztec/ethereum/test';
import { asyncMap } from '@aztec/foundation/async-map';
import { SecretValue } from '@aztec/foundation/config';
import { randomBytes } from '@aztec/foundation/crypto';
import { tryRmDir } from '@aztec/foundation/fs';
import { createLogger } from '@aztec/foundation/log';
import { resolver, reviver } from '@aztec/foundation/serialize';
import { TestDateProvider } from '@aztec/foundation/timer';
import { createPXEService, getPXEServiceConfig } from '@aztec/pxe/server';
import { tryStop } from '@aztec/stdlib/interfaces/server';
import { getConfigEnvVars as getTelemetryConfig, initTelemetryClient } from '@aztec/telemetry-client';
import { getGenesisValues } from '@aztec/world-state/testing';
import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'fs';
import { copySync, removeSync } from 'fs-extra/esm';
import fs from 'fs/promises';
import getPort from 'get-port';
import { tmpdir } from 'os';
import path, { join } from 'path';
import { mnemonicToAccount } from 'viem/accounts';
import { foundry } from 'viem/chains';
import { MNEMONIC, TEST_MAX_TX_POOL_SIZE, TEST_PEER_CHECK_INTERVAL_MS } from './fixtures.js';
import { getACVMConfig } from './get_acvm_config.js';
import { getBBConfig } from './get_bb_config.js';
import { setupL1Contracts } from './setup_l1_contracts.js';
import { createAndSyncProverNode, getLogger, getPrivateKeyFromIndex, getSponsoredFPCAddress } from './utils.js';
import { getEndToEndTestTelemetryClient } from './with_telemetry_utils.js';
export function createSnapshotManager(testName, dataPath, config = {}, deployL1ContractsArgs = {
    initialValidators: []
}) {
    return dataPath ? new SnapshotManager(testName, dataPath, config, deployL1ContractsArgs) : new MockSnapshotManager(testName, config, deployL1ContractsArgs);
}
/** Snapshot manager that does not perform snapshotting, it just applies transition and restoration functions as it receives them. */ class MockSnapshotManager {
    config;
    deployL1ContractsArgs;
    context;
    logger;
    constructor(testName, config = {}, deployL1ContractsArgs = {}){
        this.config = config;
        this.deployL1ContractsArgs = deployL1ContractsArgs;
        this.logger = createLogger(`e2e:snapshot_manager:${testName}`);
        this.logger.warn(`No data path given, will not persist any snapshots.`);
    }
    async snapshot(name, apply, restore = ()=>Promise.resolve()) {
        // We are running in disabled mode. Just apply the state.
        const context = await this.setup();
        this.logger.verbose(`Applying state transition for ${name}...`);
        const snapshotData = await apply(context);
        this.logger.verbose(`State transition for ${name} complete.`);
        // Execute the restoration function.
        await restore(snapshotData, context);
        return;
    }
    async setup() {
        if (!this.context) {
            this.context = await setupFromFresh(undefined, this.logger, this.config, this.deployL1ContractsArgs);
        }
        return this.context;
    }
    async teardown() {
        await teardown(this.context);
        this.context = undefined;
    }
}
/**
 * Snapshot engine for local e2e tests. Read more:
 * https://github.com/AztecProtocol/aztec-packages/pull/5526
 */ class SnapshotManager {
    dataPath;
    config;
    deployL1ContractsArgs;
    snapshotStack;
    context;
    livePath;
    logger;
    constructor(testName, dataPath, config = {}, deployL1ContractsArgs = {}){
        this.dataPath = dataPath;
        this.config = config;
        this.deployL1ContractsArgs = deployL1ContractsArgs;
        this.snapshotStack = [];
        this.livePath = join(this.dataPath, 'live', testName);
        this.logger = createLogger(`e2e:snapshot_manager:${testName}`);
    }
    async snapshot(name, apply, restore = ()=>Promise.resolve()) {
        const snapshotPath = join(this.dataPath, 'snapshots', ...this.snapshotStack.map((e)=>e.name), name, 'snapshot');
        if (existsSync(snapshotPath)) {
            // Snapshot exists. Record entry on stack but do nothing else as we're probably still descending the tree.
            // It's the tests responsibility to call setup() before a test to ensure subsystems get created.
            this.logger.verbose(`Snapshot exists at ${snapshotPath}. Continuing...`);
            this.snapshotStack.push({
                name,
                apply,
                restore,
                snapshotPath
            });
            return;
        }
        // Snapshot didn't exist at snapshotPath, and by definition none of the child snapshots can exist.
        // If we have no subsystem context yet, create it from the top of the snapshot stack (if it exists).
        const context = await this.setup();
        this.snapshotStack.push({
            name,
            apply,
            restore,
            snapshotPath
        });
        // Apply current state transition.
        this.logger.verbose(`Applying state transition for ${name}...`);
        const snapshotData = await apply(context);
        this.logger.verbose(`State transition for ${name} complete.`);
        // Execute the restoration function.
        await restore(snapshotData, context);
        // Save the snapshot data.
        const ethCheatCodes = new EthCheatCodesWithState(context.aztecNodeConfig.l1RpcUrls);
        const anvilStateFile = `${this.livePath}/anvil.dat`;
        await ethCheatCodes.dumpChainState(anvilStateFile);
        writeFileSync(`${this.livePath}/${name}.json`, JSON.stringify(snapshotData || {}, resolver));
        // Copy everything to snapshot path.
        // We want it to be atomic, in case multiple processes are racing to create the snapshot.
        this.logger.verbose(`Saving snapshot to ${snapshotPath}...`);
        if (mkdirSync(snapshotPath, {
            recursive: true
        })) {
            copySync(this.livePath, snapshotPath);
            this.logger.verbose(`Snapshot copied to ${snapshotPath}.`);
        } else {
            this.logger.verbose(`Snapshot already exists at ${snapshotPath}. Discarding our version.`);
            await this.teardown();
        }
    }
    /**
   * Creates and returns the subsystem context based on the current snapshot stack.
   * If the subsystem context already exists, just return it.
   * If you want to be sure to get a clean snapshot, be sure to call teardown() before calling setup().
   */ async setup() {
        // We have no subsystem context yet.
        // If one exists on the snapshot stack, create one from that snapshot.
        // Otherwise create a fresh one.
        if (!this.context) {
            removeSync(this.livePath);
            mkdirSync(this.livePath, {
                recursive: true
            });
            const previousSnapshotPath = this.snapshotStack[this.snapshotStack.length - 1]?.snapshotPath;
            if (previousSnapshotPath) {
                this.logger.verbose(`Copying snapshot from ${previousSnapshotPath} to ${this.livePath}...`);
                copySync(previousSnapshotPath, this.livePath);
                this.context = await setupFromState(this.livePath, this.logger);
                // Execute each of the previous snapshots restoration functions in turn.
                await asyncMap(this.snapshotStack, async (e)=>{
                    const snapshotData = JSON.parse(readFileSync(`${e.snapshotPath}/${e.name}.json`, 'utf-8'), reviver);
                    this.logger.verbose(`Executing restoration function for ${e.name}...`);
                    await e.restore(snapshotData, this.context);
                    this.logger.verbose(`Restoration of ${e.name} complete.`);
                });
            } else {
                this.context = await setupFromFresh(this.livePath, this.logger, this.config, this.deployL1ContractsArgs);
            }
        }
        return this.context;
    }
    /**
   * Destroys the current subsystem context.
   */ async teardown() {
        await teardown(this.context);
        this.context = undefined;
        removeSync(this.livePath);
    }
}
/**
 * Destroys the current subsystem context.
 */ async function teardown(context) {
    if (!context) {
        return;
    }
    const logger = getLogger();
    try {
        logger.info('Tearing down subsystems');
        await tryStop(context.proverNode);
        await tryStop(context.aztecNode);
        await context.acvmConfig?.cleanup();
        await context.bbConfig?.cleanup();
        await tryStop(context.anvil);
        await tryStop(context.watcher);
        await tryStop(context.blobSink);
        await tryRmDir(context.directoryToCleanup, logger);
    } catch (err) {
        logger.error('Error during teardown', err);
    }
}
/**
 * Initializes a fresh set of subsystems.
 * If given a statePath, the state will be written to the path.
 * If there is no statePath, in-memory and temporary state locations will be used.
 */ async function setupFromFresh(statePath, logger, { numberOfInitialFundedAccounts = 10, ...opts } = {}, deployL1ContractsArgs = {
    initialValidators: []
}) {
    logger.verbose(`Initializing state...`);
    const blobSinkPort = await getPort();
    // Default to no slashing
    opts.slasherFlavor ??= 'none';
    deployL1ContractsArgs.slasherFlavor ??= opts.slasherFlavor;
    // Fetch the AztecNode config.
    // TODO: For some reason this is currently the union of a bunch of subsystems. That needs fixing.
    const aztecNodeConfig = {
        ...getConfigEnvVars(),
        ...opts
    };
    aztecNodeConfig.peerCheckIntervalMS = TEST_PEER_CHECK_INTERVAL_MS;
    aztecNodeConfig.maxTxPoolSize = opts.maxTxPoolSize ?? TEST_MAX_TX_POOL_SIZE;
    // Only enable proving if specifically requested.
    aztecNodeConfig.realProofs = !!opts.realProofs;
    // Only enforce the time table if requested
    aztecNodeConfig.enforceTimeTable = !!opts.enforceTimeTable;
    // Only set the target committee size if it is explicitly set
    aztecNodeConfig.aztecTargetCommitteeSize = opts.aztecTargetCommitteeSize ?? 0;
    aztecNodeConfig.listenAddress = '127.0.0.1';
    deployL1ContractsArgs.aztecTargetCommitteeSize ??= aztecNodeConfig.aztecTargetCommitteeSize;
    // Create a temp directory for all ephemeral state and cleanup afterwards
    const directoryToCleanup = path.join(tmpdir(), randomBytes(8).toString('hex'));
    await fs.mkdir(directoryToCleanup, {
        recursive: true
    });
    if (statePath === undefined) {
        aztecNodeConfig.dataDirectory = directoryToCleanup;
    } else {
        aztecNodeConfig.dataDirectory = statePath;
    }
    aztecNodeConfig.blobSinkUrl = `http://127.0.0.1:${blobSinkPort}`;
    const hdAccount = mnemonicToAccount(MNEMONIC, {
        addressIndex: 0
    });
    const publisherPrivKeyRaw = hdAccount.getHdKey().privateKey;
    const publisherPrivKey = publisherPrivKeyRaw === null ? null : Buffer.from(publisherPrivKeyRaw);
    const l1Client = createExtendedL1Client([
        aztecNodeConfig.l1RpcUrls[0]
    ], hdAccount, foundry);
    const validatorPrivKey = getPrivateKeyFromIndex(0);
    const proverNodePrivateKey = getPrivateKeyFromIndex(0);
    aztecNodeConfig.publisherPrivateKeys = [
        new SecretValue(`0x${publisherPrivKey.toString('hex')}`)
    ];
    aztecNodeConfig.validatorPrivateKeys = new SecretValue([
        `0x${validatorPrivKey.toString('hex')}`
    ]);
    aztecNodeConfig.coinbase = opts.coinbase ?? EthAddress.fromString(`${hdAccount.address}`);
    logger.info(`Setting up environment with config`, aztecNodeConfig);
    // Start anvil. We go via a wrapper script to ensure if the parent dies, anvil dies.
    logger.verbose('Starting anvil...');
    const res = await startAnvil({
        l1BlockTime: opts.ethereumSlotDuration
    });
    const anvil = res.anvil;
    aztecNodeConfig.l1RpcUrls = [
        res.rpcUrl
    ];
    const ethCheatCodes = new EthCheatCodesWithState(aztecNodeConfig.l1RpcUrls);
    // Deploy our L1 contracts.
    logger.verbose('Deploying L1 contracts...');
    if (opts.l1StartTime) {
        await ethCheatCodes.warp(opts.l1StartTime, {
            resetBlockInterval: true
        });
    }
    const initialFundedAccounts = await generateSchnorrAccounts(numberOfInitialFundedAccounts);
    const sponsoredFPCAddress = await getSponsoredFPCAddress();
    const { genesisArchiveRoot, prefilledPublicData, fundingNeeded } = await getGenesisValues(initialFundedAccounts.map((a)=>a.address).concat(sponsoredFPCAddress), opts.initialAccountFeeJuice);
    await deployMulticall3(l1Client, logger);
    const deployL1ContractsValues = await setupL1Contracts(aztecNodeConfig.l1RpcUrls[0], hdAccount, logger, {
        ...getL1ContractsConfigEnvVars(),
        genesisArchiveRoot,
        feeJuicePortalInitialBalance: fundingNeeded,
        salt: opts.salt,
        ...deployL1ContractsArgs,
        initialValidators: opts.initialValidators
    });
    aztecNodeConfig.l1Contracts = deployL1ContractsValues.l1ContractAddresses;
    aztecNodeConfig.rollupVersion = deployL1ContractsValues.rollupVersion;
    const dateProvider = new TestDateProvider();
    const watcher = new AnvilTestWatcher(new EthCheatCodesWithState(aztecNodeConfig.l1RpcUrls), deployL1ContractsValues.l1ContractAddresses.rollupAddress, deployL1ContractsValues.l1Client, dateProvider);
    await watcher.start();
    const acvmConfig = await getACVMConfig(logger);
    if (acvmConfig) {
        aztecNodeConfig.acvmWorkingDirectory = acvmConfig.acvmWorkingDirectory;
        aztecNodeConfig.acvmBinaryPath = acvmConfig.acvmBinaryPath;
    }
    const bbConfig = await getBBConfig(logger);
    if (bbConfig) {
        aztecNodeConfig.bbBinaryPath = bbConfig.bbBinaryPath;
        aztecNodeConfig.bbWorkingDirectory = bbConfig.bbWorkingDirectory;
    }
    const telemetry = getEndToEndTestTelemetryClient(opts.metricsPort);
    // Setup blob sink service
    const blobSink = await createBlobSinkServer({
        l1ChainId: aztecNodeConfig.l1ChainId,
        l1RpcUrls: aztecNodeConfig.l1RpcUrls,
        l1Contracts: aztecNodeConfig.l1Contracts,
        port: blobSinkPort,
        dataDirectory: aztecNodeConfig.dataDirectory,
        dataStoreMapSizeKb: aztecNodeConfig.dataStoreMapSizeKb
    }, telemetry);
    await blobSink.start();
    logger.info('Creating and synching an aztec node...');
    const aztecNode = await AztecNodeService.createAndSync(aztecNodeConfig, {
        telemetry,
        dateProvider
    }, {
        prefilledPublicData
    });
    let proverNode = undefined;
    if (opts.startProverNode) {
        logger.verbose('Creating and syncing a simulated prover node with p2p disabled...');
        proverNode = await createAndSyncProverNode(`0x${proverNodePrivateKey.toString('hex')}`, aztecNodeConfig, {
            ...aztecNodeConfig.proverNodeConfig,
            dataDirectory: path.join(directoryToCleanup, randomBytes(8).toString('hex')),
            p2pEnabled: false
        }, aztecNode, prefilledPublicData);
    }
    logger.verbose('Creating pxe...');
    const pxeConfig = getPXEServiceConfig();
    pxeConfig.dataDirectory = statePath ?? path.join(directoryToCleanup, randomBytes(8).toString('hex'));
    // Only enable proving if specifically requested.
    pxeConfig.proverEnabled = !!opts.realProofs;
    const pxe = await createPXEService(aztecNode, pxeConfig);
    const cheatCodes = await CheatCodes.create(aztecNodeConfig.l1RpcUrls, pxe);
    if (statePath) {
        writeFileSync(`${statePath}/aztec_node_config.json`, JSON.stringify(aztecNodeConfig, resolver));
        writeFileSync(`${statePath}/accounts.json`, JSON.stringify(initialFundedAccounts, resolver));
    }
    return {
        aztecNodeConfig,
        anvil,
        aztecNode,
        pxe,
        sequencer: aztecNode.getSequencer(),
        acvmConfig,
        bbConfig,
        deployL1ContractsValues,
        proverNode,
        watcher,
        cheatCodes,
        dateProvider,
        blobSink,
        initialFundedAccounts,
        directoryToCleanup
    };
}
/**
 * Given a statePath, setup the system starting from that state.
 */ async function setupFromState(statePath, logger) {
    logger.verbose(`Initializing with saved state at ${statePath}...`);
    const directoryToCleanup = path.join(tmpdir(), randomBytes(8).toString('hex'));
    await fs.mkdir(directoryToCleanup, {
        recursive: true
    });
    // Run the blob sink on a random port
    const blobSinkPort = await getPort();
    // TODO: For some reason this is currently the union of a bunch of subsystems. That needs fixing.
    const aztecNodeConfig = JSON.parse(readFileSync(`${statePath}/aztec_node_config.json`, 'utf-8'), reviver);
    aztecNodeConfig.dataDirectory = statePath;
    aztecNodeConfig.blobSinkUrl = `http://127.0.0.1:${blobSinkPort}`;
    aztecNodeConfig.listenAddress = '127.0.0.1';
    const initialFundedAccounts = JSON.parse(readFileSync(`${statePath}/accounts.json`, 'utf-8'), reviver) || [];
    const { prefilledPublicData } = await getGenesisValues(initialFundedAccounts.map((a)=>a.address));
    // Start anvil. We go via a wrapper script to ensure if the parent dies, anvil dies.
    const { anvil, rpcUrl } = await startAnvil();
    aztecNodeConfig.l1RpcUrls = [
        rpcUrl
    ];
    // Load anvil state.
    const anvilStateFile = `${statePath}/anvil.dat`;
    const ethCheatCodes = new EthCheatCodesWithState(aztecNodeConfig.l1RpcUrls);
    await ethCheatCodes.loadChainState(anvilStateFile);
    // TODO: Encapsulate this in a NativeAcvm impl.
    const acvmConfig = await getACVMConfig(logger);
    if (acvmConfig) {
        aztecNodeConfig.acvmWorkingDirectory = acvmConfig.acvmWorkingDirectory;
        aztecNodeConfig.acvmBinaryPath = acvmConfig.acvmBinaryPath;
    }
    const bbConfig = await getBBConfig(logger);
    if (bbConfig) {
        aztecNodeConfig.bbBinaryPath = bbConfig.bbBinaryPath;
        aztecNodeConfig.bbWorkingDirectory = bbConfig.bbWorkingDirectory;
    }
    logger.verbose('Creating ETH clients...');
    const l1Client = createExtendedL1Client(aztecNodeConfig.l1RpcUrls, mnemonicToAccount(MNEMONIC));
    const dateProvider = new TestDateProvider();
    const watcher = new AnvilTestWatcher(new EthCheatCodesWithState(aztecNodeConfig.l1RpcUrls), aztecNodeConfig.l1Contracts.rollupAddress, l1Client, dateProvider);
    await watcher.start();
    const telemetry = initTelemetryClient(getTelemetryConfig());
    const blobSink = await createBlobSinkServer({
        l1ChainId: aztecNodeConfig.l1ChainId,
        l1RpcUrls: aztecNodeConfig.l1RpcUrls,
        l1Contracts: aztecNodeConfig.l1Contracts,
        port: blobSinkPort,
        dataDirectory: statePath,
        dataStoreMapSizeKb: aztecNodeConfig.dataStoreMapSizeKb
    }, telemetry);
    await blobSink.start();
    logger.verbose('Creating aztec node...');
    const aztecNode = await AztecNodeService.createAndSync(aztecNodeConfig, {
        telemetry,
        dateProvider
    }, {
        prefilledPublicData
    });
    let proverNode = undefined;
    if (aztecNodeConfig.startProverNode) {
        logger.verbose('Creating and syncing a simulated prover node...');
        const proverNodePrivateKey = getPrivateKeyFromIndex(2);
        const proverNodePrivateKeyHex = `0x${proverNodePrivateKey.toString('hex')}`;
        proverNode = await createAndSyncProverNode(proverNodePrivateKeyHex, aztecNodeConfig, {
            ...aztecNodeConfig.proverNodeConfig,
            dataDirectory: path.join(directoryToCleanup, randomBytes(8).toString('hex')),
            p2pEnabled: false
        }, aztecNode, prefilledPublicData);
    }
    logger.verbose('Creating pxe...');
    const pxeConfig = getPXEServiceConfig();
    pxeConfig.dataDirectory = statePath;
    const pxe = await createPXEService(aztecNode, pxeConfig);
    const cheatCodes = await CheatCodes.create(aztecNodeConfig.l1RpcUrls, pxe);
    return {
        aztecNodeConfig,
        anvil,
        aztecNode,
        pxe,
        sequencer: aztecNode.getSequencer(),
        acvmConfig,
        bbConfig,
        proverNode,
        deployL1ContractsValues: {
            l1Client,
            l1ContractAddresses: aztecNodeConfig.l1Contracts,
            rollupVersion: aztecNodeConfig.rollupVersion
        },
        watcher,
        cheatCodes,
        dateProvider,
        blobSink,
        initialFundedAccounts,
        directoryToCleanup
    };
}
/**
 * Snapshot 'apply' helper function to add accounts.
 * The 'restore' function is not provided, as it must be a closure within the test context to capture the results.
 */ export const deployAccounts = (numberOfAccounts, logger, waitUntilProven = false)=>async ({ pxe, initialFundedAccounts })=>{
        if (initialFundedAccounts.length < numberOfAccounts) {
            throw new Error(`Cannot deploy more than ${initialFundedAccounts.length} initial accounts.`);
        }
        logger.verbose('Deploying accounts funded with fee juice...');
        const deployedAccounts = initialFundedAccounts.slice(0, numberOfAccounts);
        await deployFundedSchnorrAccounts(pxe, deployedAccounts, undefined, waitUntilProven ? DefaultWaitForProvenOpts : undefined);
        return {
            deployedAccounts
        };
    };
/**
 * Registers the contract class used for test accounts and publicly deploys the instances requested.
 * Use this when you need to make a public call to an account contract, such as for requesting a public authwit.
 * @param sender - Wallet to send the deployment tx.
 * @param accountsToDeploy - Which accounts to publicly deploy.
 * @param waitUntilProven - Whether to wait for the tx to be proven.
 * @param pxeOrNode - PXE or AztecNode to wait for proven.
 */ export async function publicDeployAccounts(sender, accountsToDeploy, waitUntilProven = false, pxeOrNode) {
    const accountAddressesToDeploy = accountsToDeploy.map((a)=>'address' in a ? a.address : a);
    const instances = (await Promise.all(accountAddressesToDeploy.map((account)=>sender.getContractMetadata(account)))).map((metadata)=>metadata.contractInstance);
    const contractClass = await getContractClassFromArtifact(SchnorrAccountContractArtifact);
    const alreadyRegistered = (await sender.getContractClassMetadata(contractClass.id)).isContractClassPubliclyRegistered;
    const calls = await Promise.all([
        ...!alreadyRegistered ? [
            publishContractClass(sender, SchnorrAccountContractArtifact)
        ] : [],
        ...instances.map((instance)=>publishInstance(sender, instance))
    ]);
    const batch = new BatchCall(sender, calls);
    const txReceipt = await batch.send({
        from: accountAddressesToDeploy[0]
    }).wait();
    if (waitUntilProven) {
        if (!pxeOrNode) {
            throw new Error('Need to provide a PXE or AztecNode to wait for proven.');
        } else {
            await waitForProven(pxeOrNode, txReceipt);
        }
    }
}
