import { SLOTS_PER_EPOCH } from "@lodestar/params";
import { DataAvailabilityStatus, ExecutionPayloadStatus, StateHashTreeRootSource, computeEpochAtSlot, computeStartSlotAtEpoch, processSlots, stateTransition, } from "@lodestar/state-transition";
import { isGloasBeaconBlock } from "@lodestar/types";
import { fromHex, toRootHex } from "@lodestar/utils";
import { nextEventLoop } from "../../util/eventLoop.js";
import { getCheckpointFromState } from "../blocks/utils/checkpoint.js";
import { ChainEvent } from "../emitter.js";
import { RegenError, RegenErrorCode } from "./errors.js";
/**
 * Regenerates states that have already been processed by the fork choice
 * Since Feb 2024, we support reloading checkpoint state from disk via allowDiskReload flag. Due to its performance impact
 * this flag is only set to true in this case:
 *    - getPreState: this is for block processing, it's important to reload state in unfinality time
 *    - updateHeadState: rarely happen, but it's important to make sure we always can regen head state
 */
export class StateRegenerator {
    modules;
    constructor(modules) {
        this.modules = modules;
    }
    /**
     * Get the state to run with `block`. May be:
     * - If parent is in same epoch -> Exact state at `block.parentRoot`
     * - If parent is in prev epoch -> State after `block.parentRoot` dialed forward through epoch transition
     * - reload state if needed in this flow
     */
    async getPreState(block, opts, regenCaller) {
        const parentRoot = toRootHex(block.parentRoot);
        const parentBlock = isGloasBeaconBlock(block)
            ? this.modules.forkChoice.getBlockHexAndBlockHash(parentRoot, toRootHex(block.body.signedExecutionPayloadBid.message.parentBlockHash))
            : this.modules.forkChoice.getBlockHexDefaultStatus(parentRoot);
        if (!parentBlock) {
            throw new RegenError({
                code: RegenErrorCode.BLOCK_NOT_IN_FORKCHOICE,
                blockRoot: block.parentRoot,
            });
        }
        const parentEpoch = computeEpochAtSlot(parentBlock.slot);
        const blockEpoch = computeEpochAtSlot(block.slot);
        const allowDiskReload = true;
        // This may save us at least one epoch transition.
        // If the requested state crosses an epoch boundary
        // then we may use the checkpoint state before the block
        // We may have the checkpoint state with parent root inside the checkpoint state cache
        // through gossip validation.
        if (parentEpoch < blockEpoch) {
            return this.getBlockSlotState(parentBlock, block.slot, opts, regenCaller, allowDiskReload);
        }
        // Otherwise, get the state normally.
        return this.getState(parentBlock.stateRoot, regenCaller, allowDiskReload);
    }
    /**
     * Get state after block `blockRoot` dialed forward to `slot`
     *   - allowDiskReload should be used with care, as it will cause the state to be reloaded from disk
     */
    async getBlockSlotState(block, slot, opts, regenCaller, allowDiskReload = false) {
        if (slot < block.slot) {
            throw new RegenError({
                code: RegenErrorCode.SLOT_BEFORE_BLOCK_SLOT,
                slot,
                blockSlot: block.slot,
            });
        }
        const { blockRoot } = block;
        const { checkpointStateCache } = this.modules;
        const epoch = computeEpochAtSlot(slot);
        const latestCheckpointStateCtx = allowDiskReload
            ? await checkpointStateCache.getOrReloadLatest(blockRoot, epoch)
            : checkpointStateCache.getLatest(blockRoot, epoch);
        // If a checkpoint state exists with the given checkpoint root, it either is in requested epoch
        // or needs to have empty slots processed until the requested epoch
        if (latestCheckpointStateCtx) {
            return processSlotsByCheckpoint(this.modules, latestCheckpointStateCtx, slot, regenCaller, opts);
        }
        // Otherwise, use the fork choice to get the stateRoot from block at the checkpoint root
        // regenerate that state,
        // then process empty slots until the requested epoch
        const blockStateCtx = await this.getState(block.stateRoot, regenCaller, allowDiskReload);
        return processSlotsByCheckpoint(this.modules, blockStateCtx, slot, regenCaller, opts);
    }
    /**
     * Get state by exact root. If not in cache directly, requires finding the block that references the state from the
     * forkchoice and replaying blocks to get to it.
     *   - allowDiskReload should be used with care, as it will cause the state to be reloaded from disk
     */
    async getState(stateRoot, caller, 
    // internal option, don't want to expose to external caller
    allowDiskReload = false) {
        // Trivial case, state at stateRoot is already cached
        const cachedStateCtx = this.modules.blockStateCache.get(stateRoot);
        if (cachedStateCtx) {
            return cachedStateCtx;
        }
        // Otherwise we have to use the fork choice to traverse backwards, block by block,
        // searching the state caches
        // then replay blocks forward to the desired stateRoot
        const block = this.findFirstStateBlock(stateRoot);
        // blocks to replay, ordered highest to lowest
        // gets reversed when replayed
        const blocksToReplay = [block];
        let state = null;
        const { checkpointStateCache } = this.modules;
        const getSeedStateTimer = this.modules.metrics?.regenGetState.getSeedState.startTimer({ caller });
        // iterateAncestorBlocks only returns ancestor blocks, not the block itself
        for (const b of this.modules.forkChoice.iterateAncestorBlocks(block.blockRoot, block.payloadStatus)) {
            state = this.modules.blockStateCache.get(b.stateRoot);
            if (state) {
                break;
            }
            const lastBlockToReplay = blocksToReplay.at(-1);
            if (!lastBlockToReplay)
                continue;
            const epoch = computeEpochAtSlot(lastBlockToReplay.slot - 1);
            state = allowDiskReload
                ? await checkpointStateCache.getOrReloadLatest(b.blockRoot, epoch)
                : checkpointStateCache.getLatest(b.blockRoot, epoch);
            if (state) {
                break;
            }
            blocksToReplay.push(b);
        }
        getSeedStateTimer?.();
        if (state === null) {
            throw new RegenError({
                code: RegenErrorCode.NO_SEED_STATE,
            });
        }
        const blockCount = blocksToReplay.length;
        const MAX_EPOCH_TO_PROCESS = 5;
        if (blockCount > MAX_EPOCH_TO_PROCESS * SLOTS_PER_EPOCH) {
            throw new RegenError({
                code: RegenErrorCode.TOO_MANY_BLOCK_PROCESSED,
                stateRoot,
            });
        }
        this.modules.metrics?.regenGetState.blockCount.observe({ caller }, blockCount);
        const replaySlots = new Array(blockCount);
        const blockPromises = new Array(blockCount);
        const protoBlocksAsc = blocksToReplay.reverse();
        for (const [i, protoBlock] of protoBlocksAsc.entries()) {
            replaySlots[i] = protoBlock.slot;
            const blockInput = this.modules.seenBlockInputCache.get(protoBlock.blockRoot);
            blockPromises[i] = blockInput?.hasBlock()
                ? Promise.resolve(blockInput.getBlock())
                : this.modules.db.block.get(fromHex(protoBlock.blockRoot));
        }
        const logCtx = { stateRoot, caller, replaySlots: replaySlots.join(",") };
        this.modules.logger.debug("Replaying blocks to get state", logCtx);
        const loadBlocksTimer = this.modules.metrics?.regenGetState.loadBlocks.startTimer({ caller });
        const blockOrNulls = await Promise.all(blockPromises);
        loadBlocksTimer?.();
        const blocksByRoot = new Map();
        for (const [i, blockOrNull] of blockOrNulls.entries()) {
            // checking early here helps prevent unneccessary state transition below
            if (blockOrNull === null) {
                throw new RegenError({
                    code: RegenErrorCode.BLOCK_NOT_IN_DB,
                    blockRoot: protoBlocksAsc[i].blockRoot,
                });
            }
            blocksByRoot.set(protoBlocksAsc[i].blockRoot, blockOrNull);
        }
        const stateTransitionTimer = this.modules.metrics?.regenGetState.stateTransition.startTimer({ caller });
        for (const b of protoBlocksAsc) {
            const block = blocksByRoot.get(b.blockRoot);
            // just to make compiler happy, we checked in the above for loop already
            if (block === undefined) {
                throw new RegenError({
                    code: RegenErrorCode.BLOCK_NOT_IN_DB,
                    blockRoot: b.blockRoot,
                });
            }
            try {
                // Only advances state trusting block's signture and hashes.
                // We are only running the state transition to get a specific state's data.
                // stateTransition() does the clone() inside, transfer cache to make the regen faster
                state = stateTransition(state, block, {
                    // Replay previously imported blocks, assume valid and available
                    executionPayloadStatus: ExecutionPayloadStatus.valid,
                    dataAvailabilityStatus: DataAvailabilityStatus.Available,
                    verifyStateRoot: false,
                    verifyProposer: false,
                    verifySignatures: false,
                    dontTransferCache: false,
                }, this.modules);
                const hashTreeRootTimer = this.modules.metrics?.stateHashTreeRootTime.startTimer({
                    source: StateHashTreeRootSource.regenState,
                });
                const stateRoot = toRootHex(state.hashTreeRoot());
                hashTreeRootTimer?.();
                if (b.stateRoot !== stateRoot) {
                    throw new RegenError({
                        slot: b.slot,
                        code: RegenErrorCode.INVALID_STATE_ROOT,
                        actual: stateRoot,
                        expected: b.stateRoot,
                    });
                }
                if (allowDiskReload) {
                    // also with allowDiskReload flag, we "reload" it to the state cache too
                    this.modules.blockStateCache.add(state);
                }
            }
            catch (e) {
                throw new RegenError({
                    code: RegenErrorCode.STATE_TRANSITION_ERROR,
                    error: e,
                });
            }
        }
        stateTransitionTimer?.();
        this.modules.logger.debug("Replayed blocks to get state", { ...logCtx, stateSlot: state.slot });
        return state;
    }
    findFirstStateBlock(stateRoot) {
        for (const block of this.modules.forkChoice.forwarditerateAncestorBlocks()) {
            if (block.stateRoot === stateRoot) {
                return block;
            }
        }
        throw new RegenError({
            code: RegenErrorCode.STATE_NOT_IN_FORKCHOICE,
            stateRoot,
        });
    }
}
/**
 * Starting at `state.slot`,
 * process slots forward towards `slot`,
 * emitting "checkpoint" events after every epoch processed.
 */
async function processSlotsByCheckpoint(modules, preState, slot, regenCaller, opts) {
    let postState = await processSlotsToNearestCheckpoint(modules, preState, slot, regenCaller, opts);
    if (postState.slot < slot) {
        postState = processSlots(postState, slot, opts, modules);
    }
    return postState;
}
/**
 * Starting at `state.slot`,
 * process slots forward towards `slot`,
 * emitting "checkpoint" events after every epoch processed.
 *
 * Stops processing after no more full epochs can be processed.
 */
export async function processSlotsToNearestCheckpoint(modules, preState, slot, regenCaller, opts) {
    const preSlot = preState.slot;
    const postSlot = slot;
    const preEpoch = computeEpochAtSlot(preSlot);
    let postState = preState;
    const { checkpointStateCache, emitter, metrics, logger } = modules;
    let count = 0;
    for (let nextEpochSlot = computeStartSlotAtEpoch(preEpoch + 1); nextEpochSlot <= postSlot; nextEpochSlot += SLOTS_PER_EPOCH) {
        logger?.verbose("Processing slots over epochs", {
            slot: postState.slot,
            nextEpochSlot,
            postSlot,
            caller: regenCaller,
        });
        // processSlots calls .clone() before mutating
        postState = processSlots(postState, nextEpochSlot, opts, modules);
        metrics?.epochTransitionByCaller.inc({ caller: regenCaller });
        // this is usually added when we prepare for next slot or validate gossip block
        // then when we process the 1st block of epoch, we don't have to do state transition again
        // This adds Previous Root Checkpoint State to the checkpoint state cache
        // This may becomes the "official" checkpoint state if the 1st block of epoch is skipped
        const checkpointState = postState;
        const cp = getCheckpointFromState(checkpointState);
        checkpointStateCache.add(cp, checkpointState);
        // consumers should not mutate state ever
        emitter?.emit(ChainEvent.checkpoint, cp, checkpointState);
        if (count >= 1) {
            // in normal condition, we only process 1 epoch so never reach this
            // in that case, we want to prune state at the last 1/3 slot of slot 0 of the next epoch after importing the 1st block of epoch
            // in non-finality time, we may process a lot of epochs so need to prune the cache to keep the node healthy
            // this happened to holesky on Feb 2025, see https://github.com/ChainSafe/lodestar/issues/7495#issuecomment-2680800898
            // cannot use getBlockRootAtSlot() because nextEpochSlot = postState
            const latestBlockHex = toRootHex(cp.root);
            try {
                const persistCount = await checkpointStateCache.processState(latestBlockHex, checkpointState);
                logger?.verbose("pruning checkpointStateCache during processSlotsToNearestCheckpoint", {
                    root: latestBlockHex,
                    epoch: cp.epoch,
                    persistCount,
                });
            }
            catch (e) {
                logger?.debug("CheckpointStateCache failed to process checkpoint state", { root: latestBlockHex, epoch: cp.epoch }, e);
            }
        }
        count++;
        // this avoids keeping our node busy processing blocks
        await nextEventLoop();
    }
    return postState;
}
//# sourceMappingURL=regen.js.map