/*
 * Copyright Consensys Software Inc., 2022
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
 * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations under the License.
 */

package tech.pegasys.teku.storage.store;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static tech.pegasys.teku.infrastructure.async.SafeFutureAssert.assertThatSafeFuture;
import static tech.pegasys.teku.infrastructure.async.SafeFutureAssert.safeJoin;
import static tech.pegasys.teku.infrastructure.async.SyncAsyncRunner.SYNC_RUNNER;
import static tech.pegasys.teku.infrastructure.time.TimeUtilities.millisToSeconds;

import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;
import org.apache.tuweni.bytes.Bytes32;
import org.junit.jupiter.api.Test;
import tech.pegasys.teku.dataproviders.lookup.EarliestBlobSidecarSlotProvider;
import tech.pegasys.teku.dataproviders.lookup.StateAndBlockSummaryProvider;
import tech.pegasys.teku.infrastructure.async.SafeFuture;
import tech.pegasys.teku.infrastructure.metrics.StubMetricsSystem;
import tech.pegasys.teku.infrastructure.unsigned.UInt64;
import tech.pegasys.teku.spec.datastructures.blobs.versions.deneb.BlobSidecar;
import tech.pegasys.teku.spec.datastructures.blocks.BeaconBlock;
import tech.pegasys.teku.spec.datastructures.blocks.BlockCheckpoints;
import tech.pegasys.teku.spec.datastructures.blocks.SignedBeaconBlock;
import tech.pegasys.teku.spec.datastructures.blocks.SignedBlockAndState;
import tech.pegasys.teku.spec.datastructures.blocks.SlotAndBlockRoot;
import tech.pegasys.teku.spec.datastructures.blocks.StateAndBlockSummary;
import tech.pegasys.teku.spec.datastructures.forkchoice.InvalidCheckpointException;
import tech.pegasys.teku.spec.datastructures.forkchoice.ProtoNodeData;
import tech.pegasys.teku.spec.datastructures.forkchoice.ProtoNodeValidationStatus;
import tech.pegasys.teku.spec.datastructures.state.AnchorPoint;
import tech.pegasys.teku.spec.datastructures.state.Checkpoint;
import tech.pegasys.teku.spec.datastructures.state.CheckpointState;
import tech.pegasys.teku.spec.datastructures.state.beaconstate.BeaconState;
import tech.pegasys.teku.spec.datastructures.type.SszKZGCommitment;
import tech.pegasys.teku.spec.generator.ChainBuilder;
import tech.pegasys.teku.storage.api.StubStorageUpdateChannel;
import tech.pegasys.teku.storage.api.StubStorageUpdateChannelWithDelays;
import tech.pegasys.teku.storage.archive.nooparchive.DataArchiveNoopWriter;
import tech.pegasys.teku.storage.storageSystem.InMemoryStorageSystemBuilder;
import tech.pegasys.teku.storage.storageSystem.StorageSystem;
import tech.pegasys.teku.storage.store.UpdatableStore.StoreTransaction;

class StoreTest extends AbstractStoreTest {
  @Test
  public void create_timeLessThanGenesisTime() {
    final UInt64 genesisTime = UInt64.valueOf(100);
    final SignedBlockAndState genesis = chainBuilder.generateGenesis(genesisTime, false);
    final Checkpoint genesisCheckpoint = chainBuilder.getCurrentCheckpointForEpoch(0);

    assertThatThrownBy(
            () ->
                Store.create(
                    SYNC_RUNNER,
                    new StubMetricsSystem(),
                    spec,
                    blockProviderFromChainBuilder(),
                    StateAndBlockSummaryProvider.NOOP,
                    EarliestBlobSidecarSlotProvider.NOOP,
                    Optional.empty(),
                    genesisTime.minus(1),
                    genesisTime,
                    AnchorPoint.create(spec, genesisCheckpoint, genesis),
                    Optional.empty(),
                    genesisCheckpoint,
                    genesisCheckpoint,
                    Collections.emptyMap(),
                    Collections.emptyMap(),
                    defaultStoreConfig))
        .isInstanceOf(IllegalArgumentException.class)
        .hasMessageContaining("Time must be greater than or equal to genesisTime");
  }

  @Test
  public void retrieveSignedBlock_withLimitedCache() {
    processChainWithLimitedCache(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          final SignedBeaconBlock expectedBlock = blockAndState.getBlock();
          SafeFuture<Optional<SignedBeaconBlock>> result = store.retrieveSignedBlock(root);
          assertThat(result).isCompleted();
          assertThat(result)
              .describedAs("block %s", expectedBlock.getSlot())
              .isCompletedWithValue(Optional.of(expectedBlock));
        });
  }

  @Test
  public void isHeadWeak_withoutNodeData() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          store.computeBalanceThresholds(justifiedState(store));
          assertThat(store.isHeadWeak(root)).isFalse();
        });
  }

  @Test
  public void isHeadWeak_withSufficientWeightIsFalse() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          setProtoNodeDataForBlock(blockAndState, UInt64.valueOf("2400000001"), UInt64.MAX_VALUE);
          store.computeBalanceThresholds(justifiedState(store));
          assertThat(store.isHeadWeak(root)).isFalse();
        });
  }

  @Test
  public void isHeadWeak_Boundary() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          setProtoNodeDataForBlock(blockAndState, UInt64.valueOf("2399999999"), UInt64.MAX_VALUE);
          store.computeBalanceThresholds(justifiedState(store));
          assertThat(store.isHeadWeak(root)).isTrue();
        });
  }

  @Test
  public void isHeadWeak_withLowWeightIsTrue() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          setProtoNodeDataForBlock(blockAndState, UInt64.valueOf("1000000000"), UInt64.MAX_VALUE);
          store.computeBalanceThresholds(justifiedState(store));
          assertThat(store.isHeadWeak(root)).isTrue();
        });
  }

  @Test
  public void isParentStrong_withoutNodeData() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getBlock().getParentRoot();
          store.computeBalanceThresholds(justifiedState(store));
          assertThat(store.isParentStrong(root)).isTrue();
        });
  }

  @Test
  public void isParentStrong_withSufficientWeight() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getBlock().getParentRoot();
          setProtoNodeDataForBlock(blockAndState, UInt64.ZERO, UInt64.valueOf("19200000001"));
          store.computeBalanceThresholds(justifiedState(store));
          assertThat(store.isParentStrong(root)).isTrue();
        });
  }

  @Test
  public void isParentStrong_wityBoundaryWeight() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getBlock().getParentRoot();
          setProtoNodeDataForBlock(blockAndState, UInt64.ZERO, UInt64.valueOf("19200000000"));
          store.computeBalanceThresholds(justifiedState(store));
          assertThat(store.isParentStrong(root)).isFalse();
        });
  }

  @Test
  public void isParentStrong_wityZeroWeight() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getBlock().getParentRoot();
          setProtoNodeDataForBlock(blockAndState, UInt64.ZERO, UInt64.ZERO);
          store.computeBalanceThresholds(justifiedState(store));
          assertThat(store.isParentStrong(root)).isFalse();
        });
  }

  @Test
  public void isFfgCompetitive_checkpointMatches() {
    final BlockCheckpoints headBlockCheckpoint = mock(BlockCheckpoints.class);
    final BlockCheckpoints parentBlockCheckpoint = mock(BlockCheckpoints.class);
    final Checkpoint checkpoint = new Checkpoint(UInt64.ZERO, Bytes32.random());
    when(headBlockCheckpoint.getUnrealizedJustifiedCheckpoint()).thenReturn(checkpoint);
    when(parentBlockCheckpoint.getUnrealizedJustifiedCheckpoint()).thenReturn(checkpoint);
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          final Bytes32 parentRoot = blockAndState.getBlock().getParentRoot();
          setProtoNodeDataForBlock(blockAndState, headBlockCheckpoint, parentBlockCheckpoint);
          assertThat(store.isFfgCompetitive(root, parentRoot)).contains(true);
        });
  }

  @Test
  public void isFfgCompetitive_checkpointDifferent() {
    final BlockCheckpoints headBlockCheckpoint = mock(BlockCheckpoints.class);
    final BlockCheckpoints parentBlockCheckpoint = mock(BlockCheckpoints.class);
    final Checkpoint checkpoint = new Checkpoint(UInt64.ZERO, Bytes32.random());
    final Checkpoint checkpointParent = new Checkpoint(UInt64.ONE, Bytes32.random());
    when(headBlockCheckpoint.getUnrealizedJustifiedCheckpoint()).thenReturn(checkpoint);
    when(parentBlockCheckpoint.getUnrealizedJustifiedCheckpoint()).thenReturn(checkpointParent);
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          final Bytes32 parentRoot = blockAndState.getBlock().getParentRoot();
          setProtoNodeDataForBlock(blockAndState, headBlockCheckpoint, parentBlockCheckpoint);
          assertThat(store.isFfgCompetitive(root, parentRoot)).contains(false);
        });
  }

  @Test
  public void isFfgCompetitive_missingProtoNodeEntries() {
    processChainHeadWithMockForkChoiceStrategy(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          final Bytes32 parentRoot = blockAndState.getBlock().getParentRoot();
          assertThat(store.isFfgCompetitive(root, parentRoot)).isEmpty();
        });
  }

  @Test
  public void retrieveSignedBlock_withBlobs() {
    final UpdatableStore store = createGenesisStore();
    final UInt64 slot = UInt64.ONE;
    final SignedBlockAndState signedBlockAndState =
        chainBuilder.generateBlockAtSlot(
            slot,
            ChainBuilder.BlockOptions.create()
                .setGenerateRandomBlobs(true)
                .setStoreBlobSidecars(true));

    addBlock(store, signedBlockAndState);

    final StoreTransaction tx = store.startTransaction(new StubStorageUpdateChannel());
    chainBuilder
        .streamBlocksAndStates()
        .forEach(
            blockAndState ->
                tx.putBlockAndState(
                    blockAndState,
                    chainBuilder.getBlobSidecars(blockAndState.getBlock().getRoot()),
                    spec.calculateBlockCheckpoints(blockAndState.getState())));
    tx.commit().join();

    final SlotAndBlockRoot slotAndBlockRoot =
        new SlotAndBlockRoot(slot, signedBlockAndState.getBlock().getRoot());

    final Optional<List<BlobSidecar>> blobSidecarsFromStore =
        store.getBlobSidecarsIfAvailable(slotAndBlockRoot);

    assertThat(blobSidecarsFromStore)
        .hasValueSatisfying(blobSidecars -> assertThat(blobSidecars).isNotEmpty());

    final SafeFuture<Optional<SignedBeaconBlock>> signedBlock =
        store.retrieveSignedBlock(signedBlockAndState.getRoot());

    assertThat(signedBlock)
        .isCompletedWithValueMatching(Optional::isPresent, "Result must be present")
        .isCompletedWithValueMatching(
            signedBeaconBlock ->
                signedBeaconBlock
                    .orElseThrow()
                    .getSignedBeaconBlock()
                    .orElseThrow()
                    .equals(signedBlockAndState.getBlock()),
            " Block must match")
        .isCompletedWithValueMatching(
            signedBeaconBlock ->
                signedBeaconBlock
                    .orElseThrow()
                    .getBeaconBlock()
                    .orElseThrow()
                    .getBody()
                    .getOptionalBlobKzgCommitments()
                    .orElseThrow()
                    .asList()
                    .containsAll(
                        blobSidecarsFromStore.get().stream()
                            .map(blob -> new SszKZGCommitment(blob.getKZGCommitment()))
                            .toList()),
            "Blob sidecars must match");
  }

  @Test
  public void retrieveSignedBlock_withEmptyBlobs() {
    final UpdatableStore store = createGenesisStore();
    final UInt64 slot = UInt64.ONE;
    final SignedBlockAndState signedBlockAndState =
        chainBuilder.generateBlockAtSlot(
            slot,
            ChainBuilder.BlockOptions.create()
                .setGenerateRandomBlobs(false)
                .setStoreBlobSidecars(false));

    addBlock(store, signedBlockAndState);

    final StoreTransaction tx = store.startTransaction(new StubStorageUpdateChannel());
    chainBuilder
        .streamBlocksAndStates()
        .forEach(
            blockAndState ->
                tx.putBlockAndState(
                    blockAndState,
                    chainBuilder.getBlobSidecars(blockAndState.getBlock().getRoot()),
                    spec.calculateBlockCheckpoints(blockAndState.getState())));
    tx.commit().join();

    final SlotAndBlockRoot slotAndBlockRoot =
        new SlotAndBlockRoot(slot, signedBlockAndState.getBlock().getRoot());
    final Optional<List<BlobSidecar>> blobSidecarsFromStore =
        store.getBlobSidecarsIfAvailable(slotAndBlockRoot);

    assertThat(blobSidecarsFromStore)
        .hasValueSatisfying(blobSidecars -> assertThat(blobSidecars).isEmpty());

    final SafeFuture<Optional<SignedBeaconBlock>> signedBlock =
        store.retrieveSignedBlock(signedBlockAndState.getRoot());

    assertThat(signedBlock)
        .isCompletedWithValueMatching(Optional::isPresent, "Result must be present")
        .isCompletedWithValueMatching(
            signedBeaconBlock ->
                signedBeaconBlock
                    .orElseThrow()
                    .getSignedBeaconBlock()
                    .orElseThrow()
                    .equals(signedBlockAndState.getBlock()),
            " Block must match")
        .isCompletedWithValueMatching(
            signedBeaconBlockAndBlobsSidecar ->
                signedBeaconBlockAndBlobsSidecar
                    .orElseThrow()
                    .getBeaconBlock()
                    .orElseThrow()
                    .getBody()
                    .getOptionalBlobKzgCommitments()
                    .orElseThrow()
                    .isEmpty(),
            "Blob sidecars must be empty");
  }

  @Test
  public void retrieveSignedBlock_shouldReturnEmptyIfBlockNotPresent() {
    final UpdatableStore store = createGenesisStore();
    final SignedBlockAndState blockAndState =
        chainBuilder.generateBlockAtSlot(
            1, ChainBuilder.BlockOptions.create().setGenerateRandomBlobs(true));

    final SafeFuture<Optional<SignedBeaconBlock>> signedBeaconBlock =
        store.retrieveSignedBlock(blockAndState.getRoot());

    assertThat(signedBeaconBlock)
        .isCompletedWithValueMatching(Optional::isEmpty, "Result must be empty");
  }

  @Test
  public void retrieveEarliestBlobSidecarSlot_shouldReturnUpdatedValue() {
    final StorageSystem storageSystem = InMemoryStorageSystemBuilder.buildDefault(spec);
    storageSystem.chainUpdater().initializeGenesis();
    final UpdatableStore store =
        createGenesisStore(
            () ->
                SafeFuture.completedFuture(storageSystem.database().getEarliestBlobSidecarSlot()));

    assertThat(store.retrieveEarliestBlobSidecarSlot())
        .isCompletedWithValueMatching(
            maybeEarliestBlobSidecarSlot ->
                maybeEarliestBlobSidecarSlot.isPresent()
                    && maybeEarliestBlobSidecarSlot.get().equals(UInt64.ZERO));

    storageSystem
        .chainUpdater()
        .advanceChainUntil(
            10,
            ChainBuilder.BlockOptions.create()
                .setGenerateRandomBlobs(true)
                .setStoreBlobSidecars(true));

    assertThat(store.retrieveEarliestBlobSidecarSlot())
        .isCompletedWithValueMatching(
            maybeEarliestBlobSidecarSlot ->
                maybeEarliestBlobSidecarSlot.isPresent()
                    && maybeEarliestBlobSidecarSlot.get().equals(UInt64.ZERO));

    storageSystem
        .database()
        .pruneOldestBlobSidecars(UInt64.valueOf(5), 3, new DataArchiveNoopWriter<>());

    assertThat(store.retrieveEarliestBlobSidecarSlot())
        .isCompletedWithValueMatching(
            maybeEarliestBlobSidecarSlot ->
                maybeEarliestBlobSidecarSlot.isPresent()
                    && maybeEarliestBlobSidecarSlot.get().equals(UInt64.valueOf(4)));
  }

  private void setProtoNodeDataForBlock(
      final SignedBlockAndState blockAndState,
      final BlockCheckpoints headCheckpoint,
      final BlockCheckpoints parentCheckpoint) {
    final Bytes32 root = blockAndState.getRoot();
    final Bytes32 parentRoot = blockAndState.getParentRoot();
    final ProtoNodeData protoNodeData =
        new ProtoNodeData(
            UInt64.ONE,
            root,
            blockAndState.getParentRoot(),
            blockAndState.getStateRoot(),
            UInt64.ZERO,
            Bytes32.random(),
            ProtoNodeValidationStatus.VALID,
            headCheckpoint,
            UInt64.ZERO);
    final ProtoNodeData parentNodeData =
        new ProtoNodeData(
            UInt64.ZERO,
            parentRoot,
            Bytes32.random(),
            blockAndState.getStateRoot(),
            UInt64.ZERO,
            Bytes32.random(),
            ProtoNodeValidationStatus.VALID,
            parentCheckpoint,
            UInt64.ZERO);
    when(dummyForkChoiceStrategy.getBlockData(root)).thenReturn(Optional.of(protoNodeData));
    when(dummyForkChoiceStrategy.getBlockData(parentRoot)).thenReturn(Optional.of(parentNodeData));
  }

  private void setProtoNodeDataForBlock(
      final SignedBlockAndState blockAndState, final UInt64 headValue, final UInt64 parentValue) {
    final Bytes32 root = blockAndState.getRoot();
    final Bytes32 parentRoot = blockAndState.getParentRoot();
    final ProtoNodeData protoNodeData =
        new ProtoNodeData(
            UInt64.ONE,
            root,
            blockAndState.getParentRoot(),
            blockAndState.getStateRoot(),
            UInt64.ZERO,
            Bytes32.random(),
            ProtoNodeValidationStatus.VALID,
            null,
            headValue);
    final ProtoNodeData parentNodeData =
        new ProtoNodeData(
            UInt64.ZERO,
            parentRoot,
            Bytes32.random(),
            blockAndState.getStateRoot(),
            UInt64.ZERO,
            Bytes32.random(),
            ProtoNodeValidationStatus.VALID,
            null,
            parentValue);
    when(dummyForkChoiceStrategy.getBlockData(root)).thenReturn(Optional.of(protoNodeData));
    when(dummyForkChoiceStrategy.getBlockData(parentRoot)).thenReturn(Optional.of(parentNodeData));
  }

  @Test
  public void epochStatesCacheMostRecentlyAddedStates() {
    final int cacheSize = 2;
    final StoreConfig pruningOptions =
        StoreConfig.builder()
            .checkpointStateCacheSize(cacheSize)
            .blockCacheSize(cacheSize)
            .stateCacheSize(cacheSize)
            .epochStateCacheSize(cacheSize)
            .build();

    final UpdatableStore store = createGenesisStore(pruningOptions);

    chainBuilder.generateBlocksUpToSlot(25);
    // Add blocks
    final StoreTransaction tx = store.startTransaction(new StubStorageUpdateChannel());
    chainBuilder
        .streamBlocksAndStates(0, 25)
        .forEach(
            blockAndState -> {
              tx.putBlockAndState(
                  blockAndState, spec.calculateBlockCheckpoints(blockAndState.getState()));
            });
    tx.commit().join();

    final Store s = (Store) store;
    assertThat(store.retrieveStateAtSlot(stateAndBlockAtSlot(16, chainBuilder))).isCompleted();
    assertThat(store.retrieveStateAtSlot(stateAndBlockAtSlot(24, chainBuilder))).isCompleted();
    assertThat(store.retrieveStateAtSlot(stateAndBlockAtSlot(8, chainBuilder))).isCompleted();
    assertThat(
            s.getEpochStates().orElseThrow().values().stream()
                .map(StateAndBlockSummary::getSlot)
                .map(UInt64::intValue)
                .collect(Collectors.toList()))
        .containsExactlyInAnyOrder(24, 8);
  }

  private SlotAndBlockRoot stateAndBlockAtSlot(final int i, final ChainBuilder chainBuilder) {
    return new SlotAndBlockRoot(UInt64.valueOf(i), chainBuilder.getBlockAtSlot(i).getRoot());
  }

  @Test
  public void retrieveBlock_withLimitedCache() {
    processChainWithLimitedCache(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          final BeaconBlock expectedBlock = blockAndState.getBlock().getMessage();
          SafeFuture<Optional<BeaconBlock>> result = store.retrieveBlock(root);
          assertThat(result).isCompleted();
          assertThat(result)
              .describedAs("block %s", expectedBlock.getSlot())
              .isCompletedWithValue(Optional.of(expectedBlock));
        });
  }

  @Test
  public void retrieveBlockAndState_withLimitedCache() {
    processChainWithLimitedCache(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          SafeFuture<Optional<SignedBlockAndState>> result = store.retrieveBlockAndState(root);
          assertThat(result).isCompleted();
          assertThat(result)
              .describedAs("block and state at %s", blockAndState.getSlot())
              .isCompletedWithValue(Optional.of(blockAndState));
        });
  }

  @Test
  public void retrieveBlockState_withLimitedCache() {
    processChainWithLimitedCache(
        (store, blockAndState) -> {
          final Bytes32 root = blockAndState.getRoot();
          SafeFuture<Optional<BeaconState>> result = store.retrieveBlockState(root);
          assertThat(result).isCompleted();
          assertThat(result)
              .describedAs("State at %s", blockAndState.getSlot())
              .isCompletedWithValue(Optional.of(blockAndState.getState()));
        });
  }

  @Test
  public void retrieveCheckpointState_withLimitedCache() {
    processCheckpointsWithLimitedCache(
        (store, checkpointState) -> {
          SafeFuture<Optional<BeaconState>> result =
              store.retrieveCheckpointState(checkpointState.getCheckpoint());
          assertThat(result)
              .describedAs("Checkpoint state for checkpoint %s", checkpointState.getCheckpoint())
              .isCompletedWithValue(Optional.of(checkpointState.getState()));
        });
  }

  @Test
  public void shouldApplyChangesWhenTransactionCommits() {
    testApplyChangesWhenTransactionCommits(false);
  }

  @Test
  public void shouldApplyChangesWhenTransactionCommits_withInterleavedTx() {
    testApplyChangesWhenTransactionCommits(true);
  }

  @Test
  public void retrieveCheckpointState_shouldGenerateCheckpointStates() {
    final UpdatableStore store = createGenesisStore();
    final SignedBlockAndState genesisBlockAndState = chainBuilder.getLatestBlockAndState();
    final Checkpoint checkpoint = new Checkpoint(UInt64.ONE, genesisBlockAndState.getRoot());

    final SafeFuture<Optional<BeaconState>> result = store.retrieveCheckpointState(checkpoint);
    assertThatSafeFuture(result).isCompletedWithNonEmptyOptional();
    final BeaconState checkpointState = safeJoin(result).orElseThrow();
    assertThat(checkpointState.getSlot()).isEqualTo(checkpoint.getEpochStartSlot(spec));
    assertThat(checkpointState.getLatestBlockHeader().hashTreeRoot())
        .isEqualTo(checkpoint.getRoot());
  }

  @Test
  public void retrieveCheckpointState_forGenesis() {
    final UpdatableStore store = createGenesisStore();
    final SignedBlockAndState genesisBlockAndState = chainBuilder.getLatestBlockAndState();
    final Checkpoint checkpoint = new Checkpoint(UInt64.ZERO, genesisBlockAndState.getRoot());

    final BeaconState baseState = genesisBlockAndState.getState();
    final SafeFuture<Optional<BeaconState>> result =
        store.retrieveCheckpointState(checkpoint, baseState);
    assertThatSafeFuture(result).isCompletedWithOptionalContaining(baseState);
  }

  @Test
  public void retrieveCheckpointState_forEpochPastGenesis() {
    final UpdatableStore store = createGenesisStore();
    final SignedBlockAndState genesisBlockAndState = chainBuilder.getLatestBlockAndState();
    final Checkpoint checkpoint = new Checkpoint(UInt64.ONE, genesisBlockAndState.getRoot());

    final BeaconState baseState = genesisBlockAndState.getState();
    final SafeFuture<Optional<BeaconState>> resultFuture =
        store.retrieveCheckpointState(checkpoint, baseState);
    assertThatSafeFuture(resultFuture).isCompletedWithNonEmptyOptional();
    final BeaconState result = safeJoin(resultFuture).orElseThrow();
    assertThat(result.getSlot()).isGreaterThan(baseState.getSlot());
    assertThat(result.getSlot()).isEqualTo(checkpoint.getEpochStartSlot(spec));
    assertThat(result.getLatestBlockHeader().hashTreeRoot()).isEqualTo(checkpoint.getRoot());
  }

  @Test
  public void retrieveCheckpointState_invalidState() {
    final UpdatableStore store = createGenesisStore();

    final UInt64 epoch = UInt64.valueOf(2);
    final UInt64 startSlot = spec.computeStartSlotAtEpoch(epoch);
    final SignedBlockAndState futureBlockAndState = chainBuilder.generateBlockAtSlot(startSlot);

    final Checkpoint checkpoint = new Checkpoint(UInt64.ONE, futureBlockAndState.getRoot());

    assertThatSafeFuture(store.retrieveCheckpointState(checkpoint, futureBlockAndState.getState()))
        .isCompletedExceptionallyWith(InvalidCheckpointException.class);
  }

  @Test
  public void retrieveFinalizedCheckpointAndState() {
    final UpdatableStore store = createGenesisStore();
    final SignedBlockAndState finalizedBlockAndState =
        chainBuilder.generateBlockAtSlot(spec.slotsPerEpoch(UInt64.ZERO) - 1);
    final Checkpoint finalizedCheckpoint =
        new Checkpoint(UInt64.ONE, finalizedBlockAndState.getRoot());

    final StoreTransaction tx = store.startTransaction(new StubStorageUpdateChannel());
    tx.putBlockAndState(
        finalizedBlockAndState, spec.calculateBlockCheckpoints(finalizedBlockAndState.getState()));
    tx.setFinalizedCheckpoint(finalizedCheckpoint, false);
    assertThat(tx.commit()).isCompleted();

    final SafeFuture<CheckpointState> result = store.retrieveFinalizedCheckpointAndState();
    assertThat(result).isCompleted();
    assertThat(safeJoin(result).getCheckpoint()).isEqualTo(finalizedCheckpoint);
    assertThat(safeJoin(result).getRoot()).isEqualTo(finalizedBlockAndState.getRoot());
    assertThat(safeJoin(result).getState()).isNotEqualTo(finalizedBlockAndState.getState());
    assertThat(safeJoin(result).getState().getSlot())
        .isEqualTo(finalizedBlockAndState.getSlot().plus(1));
  }

  @Test
  public void
      retrieveCheckpointState_shouldThrowInvalidCheckpointExceptionWhenEpochBeforeBlockRoot() {
    final UpdatableStore store = createGenesisStore();
    final UInt64 epoch = UInt64.valueOf(2);
    final UInt64 startSlot = spec.computeStartSlotAtEpoch(epoch);
    final Bytes32 futureRoot = chainBuilder.generateBlockAtSlot(startSlot).getRoot();

    // Add blocks
    final StoreTransaction tx = store.startTransaction(new StubStorageUpdateChannel());
    chainBuilder
        .streamBlocksAndStates()
        .forEach(
            blockAndState ->
                tx.putBlockAndState(
                    blockAndState, spec.calculateBlockCheckpoints(blockAndState.getState())));
    tx.commit().join();

    final Checkpoint checkpoint = new Checkpoint(UInt64.ONE, futureRoot);
    final SafeFuture<Optional<BeaconState>> result = store.retrieveCheckpointState(checkpoint);
    assertThat(result).isCompletedExceptionally();
    assertThatThrownBy(result::get).hasCauseInstanceOf(InvalidCheckpointException.class);
  }

  @Test
  public void shouldKeepOnlyMostRecentBlocksInBlockCache() {
    final UpdatableStore store = createGenesisStore();
    final UInt64 epoch = UInt64.valueOf(5);
    final UInt64 startSlot = spec.computeStartSlotAtEpoch(epoch);
    chainBuilder.generateBlocksUpToSlot(startSlot);

    // Add blocks
    final StoreTransaction tx = store.startTransaction(new StubStorageUpdateChannel());
    chainBuilder
        .streamBlocksAndStates()
        .forEach(
            blockAndState ->
                tx.putBlockAndState(
                    blockAndState, spec.calculateBlockCheckpoints(blockAndState.getState())));
    safeJoin(tx.commit());
    final List<SignedBlockAndState> last32 =
        chainBuilder
            .streamBlocksAndStates()
            .dropWhile(
                signedBlockAndState ->
                    signedBlockAndState
                        .getSlot()
                        .isLessThanOrEqualTo(
                            chainBuilder
                                .getLatestBlockAndState()
                                .getSlot()
                                .minus(defaultStoreConfig.getBlockCacheSize())))
            .toList();
    for (final SignedBlockAndState signedBlockAndState : last32) {
      assertThat(store.getBlockIfAvailable(signedBlockAndState.getRoot())).isPresent();
    }
  }

  private BeaconState justifiedState(final UpdatableStore store) {
    return safeJoin(store.retrieveCheckpointState(store.getJustifiedCheckpoint())).orElseThrow();
  }

  private void testApplyChangesWhenTransactionCommits(final boolean withInterleavedTransaction) {
    final UpdatableStore store = createGenesisStore();
    final UInt64 epoch3 = UInt64.valueOf(4);
    final UInt64 epoch3Slot = spec.computeStartSlotAtEpoch(epoch3);
    chainBuilder.generateBlocksUpToSlot(epoch3Slot);

    final Checkpoint genesisCheckpoint = store.getFinalizedCheckpoint();
    final UInt64 initialTimeMillis = store.getTimeInMillis();
    final UInt64 genesisTime = store.getGenesisTime();

    final Checkpoint checkpoint1 = chainBuilder.getCurrentCheckpointForEpoch(UInt64.valueOf(1));
    final Checkpoint checkpoint2 = chainBuilder.getCurrentCheckpointForEpoch(UInt64.valueOf(2));
    final Checkpoint checkpoint3 = chainBuilder.getCurrentCheckpointForEpoch(UInt64.valueOf(3));

    // Start transaction
    final StubStorageUpdateChannelWithDelays updateChannel =
        new StubStorageUpdateChannelWithDelays();
    final StoreTransaction tx = store.startTransaction(updateChannel);
    // Add blocks
    chainBuilder
        .streamBlocksAndStates()
        .forEach(
            blockAndState ->
                tx.putBlockAndState(
                    blockAndState, spec.calculateBlockCheckpoints(blockAndState.getState())));
    // Update checkpoints
    tx.setFinalizedCheckpoint(checkpoint1, false);
    tx.setJustifiedCheckpoint(checkpoint2);
    tx.setBestJustifiedCheckpoint(checkpoint3);
    // Update time
    UInt64 firstUpdateTimeMillis = initialTimeMillis.plus(1300);
    tx.setTimeMillis(firstUpdateTimeMillis);
    UInt64 updatedGenesisTime = genesisTime.plus(UInt64.ONE);
    tx.setGenesisTime(updatedGenesisTime);

    // Check that store is not yet updated
    // Check blocks
    chainBuilder
        .streamBlocksAndStates(1, chainBuilder.getLatestSlot().longValue())
        .forEach(b -> assertThat(store.containsBlock(b.getRoot())).isFalse());
    // Check checkpoints
    assertThat(store.getJustifiedCheckpoint()).isEqualTo(genesisCheckpoint);
    assertThat(store.getBestJustifiedCheckpoint()).isEqualTo(genesisCheckpoint);
    assertThat(store.getFinalizedCheckpoint()).isEqualTo(genesisCheckpoint);
    // Check time
    assertThat(store.getTimeSeconds()).isEqualTo(millisToSeconds(initialTimeMillis));
    assertThat(store.getTimeInMillis()).isEqualTo(initialTimeMillis);
    assertThat(store.getGenesisTime()).isEqualTo(genesisTime);

    // Check that transaction is updated
    chainBuilder
        .streamBlocksAndStates(1, chainBuilder.getLatestSlot().longValue())
        .forEach(
            b ->
                assertThat(tx.retrieveBlockAndState(b.getRoot()))
                    .isCompletedWithValue(Optional.of(b)));
    // Check checkpoints
    assertThat(tx.getFinalizedCheckpoint()).isEqualTo(checkpoint1);
    assertThat(tx.getJustifiedCheckpoint()).isEqualTo(checkpoint2);
    assertThat(tx.getBestJustifiedCheckpoint()).isEqualTo(checkpoint3);
    // Check time
    assertThat(tx.getTimeSeconds()).isEqualTo(millisToSeconds(firstUpdateTimeMillis));
    assertThat(tx.getTimeInMillis()).isEqualTo(firstUpdateTimeMillis);
    assertThat(tx.getGenesisTime()).isEqualTo(updatedGenesisTime);

    // Commit transaction
    final SafeFuture<Void> txResult = tx.commit();

    final UInt64 expectedTimeMillis;
    final SafeFuture<Void> txResult2;
    if (withInterleavedTransaction) {
      expectedTimeMillis = firstUpdateTimeMillis.plus(1500);
      StoreTransaction tx2 = store.startTransaction(updateChannel);
      tx2.setTimeMillis(expectedTimeMillis);
      txResult2 = tx2.commit();
    } else {
      expectedTimeMillis = firstUpdateTimeMillis;
      txResult2 = SafeFuture.COMPLETE;
    }

    // Complete transactions
    assertThat(updateChannel.getAsyncRunner().countDelayedActions()).isLessThanOrEqualTo(2);
    updateChannel.getAsyncRunner().executeUntilDone();
    assertThat(txResult).isCompleted();
    assertThat(txResult2).isCompleted();

    // Check store is updated
    chainBuilder
        .streamBlocksAndStates(checkpoint3.getEpochStartSlot(spec), chainBuilder.getLatestSlot())
        .forEach(
            b ->
                assertThat(store.retrieveBlockAndState(b.getRoot()))
                    .isCompletedWithValue(Optional.of(b)));
    // Check checkpoints
    assertThat(store.getFinalizedCheckpoint()).isEqualTo(checkpoint1);
    assertThat(store.getJustifiedCheckpoint()).isEqualTo(checkpoint2);
    assertThat(store.getBestJustifiedCheckpoint()).isEqualTo(checkpoint3);
    // Extra checks for finalized checkpoint
    final SafeFuture<CheckpointState> finalizedCheckpointState =
        store.retrieveFinalizedCheckpointAndState();
    assertThat(finalizedCheckpointState).isCompleted();
    assertThat(safeJoin(finalizedCheckpointState).getCheckpoint()).isEqualTo(checkpoint1);
    // Check time
    assertThat(store.getTimeInMillis()).isEqualTo(expectedTimeMillis);
    assertThat(store.getTimeSeconds()).isEqualTo(millisToSeconds(expectedTimeMillis));
    assertThat(store.getGenesisTime()).isEqualTo(updatedGenesisTime);

    // Check store was pruned as expected
    final List<Bytes32> expectedBlockRoots =
        chainBuilder
            .streamBlocksAndStates(checkpoint1.getEpochStartSlot(spec))
            .map(SignedBlockAndState::getRoot)
            .collect(Collectors.toList());
    assertThat(store.getOrderedBlockRoots()).containsExactlyElementsOf(expectedBlockRoots);
  }
}
