/*
 * 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.statetransition.util;

import static org.assertj.core.api.Assertions.assertThat;

import java.util.ArrayList;
import java.util.List;
import org.apache.tuweni.bytes.Bytes32;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import tech.pegasys.teku.infrastructure.metrics.StubMetricsSystem;
import tech.pegasys.teku.infrastructure.unsigned.UInt64;
import tech.pegasys.teku.spec.Spec;
import tech.pegasys.teku.spec.TestSpecFactory;
import tech.pegasys.teku.spec.datastructures.blocks.SignedBeaconBlock;
import tech.pegasys.teku.spec.datastructures.state.Checkpoint;
import tech.pegasys.teku.spec.util.DataStructureUtil;

public class PendingPoolTest {
  private final Spec spec = TestSpecFactory.createDefault();
  private final DataStructureUtil dataStructureUtil = new DataStructureUtil(spec);
  private final UInt64 historicalTolerance = UInt64.valueOf(5);
  private final UInt64 futureTolerance = UInt64.valueOf(2);
  private final StubMetricsSystem metricsSystem = new StubMetricsSystem();
  private final int maxItems = 15;
  private final PendingPool<SignedBeaconBlock> pendingPool =
      new PoolFactory(metricsSystem)
          .createPendingPoolForBlocks(spec, historicalTolerance, futureTolerance, maxItems);
  private UInt64 currentSlot = historicalTolerance.times(2);
  private final List<Bytes32> requiredRootEvents = new ArrayList<>();
  private final List<Bytes32> requiredRootDroppedEvents = new ArrayList<>();

  @BeforeEach
  public void setup() {
    // Set up slot
    pendingPool.subscribeRequiredBlockRoot(requiredRootEvents::add);
    pendingPool.subscribeRequiredBlockRootDropped(requiredRootDroppedEvents::add);
    setSlot(currentSlot);
  }

  private void setSlot(final long slot) {
    setSlot(UInt64.valueOf(slot));
  }

  private void setSlot(final UInt64 slot) {
    currentSlot = slot;
    pendingPool.onSlot(slot);
  }

  @Test
  public void add_shouldDeferSubscribersCallToAvoidConcurrentModificationException() {
    final SignedBeaconBlock block =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());

    final SignedBeaconBlock block2 =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());

    pendingPool.subscribeRequiredBlockRoot(
        blockRoot -> {
          if (blockRoot.equals(block.getParentRoot())) {
            pendingPool.add(block2);
          }
        });
    pendingPool.add(block);

    assertThat(pendingPool.contains(block2)).isTrue();
  }

  @Test
  public void add_blockForCurrentSlot() {
    final SignedBeaconBlock block =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
    pendingPool.add(block);

    assertThat(pendingPool.contains(block)).isTrue();
    assertThat(pendingPool.size()).isEqualTo(1);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false))
        .containsExactlyInAnyOrder(block);
    assertThat(requiredRootEvents).containsExactly(block.getParentRoot());
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void add_historicalBlockWithinWindow() {
    final UInt64 slot = currentSlot.minus(historicalTolerance);
    final SignedBeaconBlock block = dataStructureUtil.randomSignedBeaconBlock(slot.longValue());
    pendingPool.add(block);

    assertThat(pendingPool.contains(block)).isTrue();
    assertThat(pendingPool.size()).isEqualTo(1);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false))
        .containsExactlyInAnyOrder(block);
    assertThat(requiredRootEvents).containsExactly(block.getParentRoot());
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void add_historicalBlockOutsideWindow() {
    final UInt64 slot = currentSlot.minus(historicalTolerance).minus(UInt64.ONE);
    final SignedBeaconBlock block = dataStructureUtil.randomSignedBeaconBlock(slot.longValue());
    pendingPool.add(block);

    assertThat(pendingPool.contains(block)).isFalse();
    assertThat(pendingPool.size()).isEqualTo(0);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false)).isEmpty();
    assertThat(requiredRootEvents).isEmpty();
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void add_futureBlockWithinWindow() {
    final UInt64 slot = currentSlot.plus(futureTolerance);
    final SignedBeaconBlock block = dataStructureUtil.randomSignedBeaconBlock(slot.longValue());
    pendingPool.add(block);

    assertThat(pendingPool.contains(block)).isTrue();
    assertThat(pendingPool.size()).isEqualTo(1);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false))
        .containsExactlyInAnyOrder(block);
    assertThat(requiredRootEvents).containsExactly(block.getParentRoot());
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void add_futureBlockOutsideWindow() {
    final UInt64 slot = currentSlot.plus(futureTolerance).plus(UInt64.ONE);
    final SignedBeaconBlock block = dataStructureUtil.randomSignedBeaconBlock(slot.longValue());
    pendingPool.add(block);

    assertThat(pendingPool.contains(block)).isFalse();
    assertThat(pendingPool.size()).isEqualTo(0);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false)).isEmpty();
    assertThat(requiredRootEvents).isEmpty();
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void add_nonFinalizedBlock() {
    final SignedBeaconBlock finalizedBlock = dataStructureUtil.randomSignedBeaconBlock(10);
    final Checkpoint checkpoint = finalizedCheckpoint(finalizedBlock);
    pendingPool.onNewFinalizedCheckpoint(checkpoint, false);

    final UInt64 slot = checkpoint.getEpochStartSlot(spec).plus(UInt64.ONE);
    setSlot(slot);
    final SignedBeaconBlock block = dataStructureUtil.randomSignedBeaconBlock(slot.longValue());

    pendingPool.add(block);
    assertThat(pendingPool.contains(block)).isTrue();
    assertThat(pendingPool.size()).isEqualTo(1);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false))
        .containsExactlyInAnyOrder(block);
    assertThat(requiredRootEvents).containsExactly(block.getParentRoot());
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void add_finalizedBlock() {
    final SignedBeaconBlock finalizedBlock = dataStructureUtil.randomSignedBeaconBlock(10);
    final Checkpoint checkpoint = finalizedCheckpoint(finalizedBlock);
    pendingPool.onNewFinalizedCheckpoint(checkpoint, false);
    final long slot = checkpoint.getEpochStartSlot(spec).longValue() + 10;
    setSlot(slot);

    final long blockSlot = checkpoint.getEpochStartSlot(spec).longValue();
    final SignedBeaconBlock block = dataStructureUtil.randomSignedBeaconBlock(blockSlot);

    pendingPool.add(block);
    assertThat(pendingPool.contains(block)).isFalse();
    assertThat(pendingPool.size()).isEqualTo(0);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false)).isEmpty();
    assertThat(requiredRootEvents).isEmpty();
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void add_moreThanMaxItems() {
    for (int i = 0; i < maxItems * 2; i++) {
      final SignedBeaconBlock block =
          dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
      pendingPool.add(block);

      final int expectedSize = Math.min(maxItems, i + 1);
      assertThat(pendingPool.contains(block)).isTrue();
      assertThat(pendingPool.size()).isEqualTo(expectedSize);
    }

    // Final sanity check
    assertThat(pendingPool.size()).isEqualTo(maxItems);
  }

  @Test
  public void add_dontDropItemsWhenNewItemsIgnored() {
    final List<SignedBeaconBlock> blocks = new ArrayList<>();
    for (int i = 0; i < maxItems; i++) {
      final SignedBeaconBlock block =
          dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
      pendingPool.add(block);
      blocks.add(block);

      final int expectedSize = Math.min(maxItems, i + 1);
      assertThat(pendingPool.contains(block)).isTrue();
      assertThat(pendingPool.size()).isEqualTo(expectedSize);
    }
    assertThat(pendingPool.size()).isEqualTo(maxItems);

    // Add some blocks that should be ignored
    SignedBeaconBlock ignoredBlock =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.plus(futureTolerance).plus(1));
    pendingPool.add(ignoredBlock);
    ignoredBlock =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.minus(historicalTolerance).minus(1));
    pendingPool.add(ignoredBlock);

    // All initial blocks should still be present
    assertThat(pendingPool.size()).isEqualTo(maxItems);
    blocks.forEach(b -> assertThat(pendingPool.contains(b)).isTrue());
  }

  @Test
  public void add_shouldDropOldestItemsWhenEnforcingLimits() {
    final UInt64 olderSlot = currentSlot;
    final UInt64 newerSlot = currentSlot.plus(1);
    SignedBeaconBlock block;

    // Add older block
    SignedBeaconBlock olderBlock = dataStructureUtil.randomSignedBeaconBlock(olderSlot);
    pendingPool.add(olderBlock);
    assertThat(pendingPool.contains(olderBlock)).isTrue();

    // Add max newer blocks
    for (int i = 0; i < maxItems; i++) {
      block = dataStructureUtil.randomSignedBeaconBlock(newerSlot);
      pendingPool.add(block);

      assertThat(pendingPool.contains(block)).isTrue();
      assertThat(pendingPool.size()).isLessThanOrEqualTo(maxItems);
    }

    // Older block should've been dropped
    assertThat(pendingPool.contains(olderBlock)).isFalse();

    // Add older block
    olderBlock = dataStructureUtil.randomSignedBeaconBlock(olderSlot);
    pendingPool.add(olderBlock);
    assertThat(pendingPool.contains(olderBlock)).isTrue();
    // Add newer block
    block = dataStructureUtil.randomSignedBeaconBlock(newerSlot);
    pendingPool.add(block);

    // Older block should be dropped when newer block is added
    assertThat(pendingPool.contains(olderBlock)).isFalse();
    assertThat(pendingPool.size()).isEqualTo(maxItems);
  }

  private Checkpoint finalizedCheckpoint(final SignedBeaconBlock block) {
    final UInt64 epoch = spec.computeEpochAtSlot(block.getSlot()).plus(UInt64.ONE);
    final Bytes32 root = block.getMessage().hashTreeRoot();

    return new Checkpoint(epoch, root);
  }

  @Test
  public void add_duplicateBlock() {
    final SignedBeaconBlock block =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
    final Bytes32 parentRoot = block.getParentRoot();
    pendingPool.add(block);
    pendingPool.add(block);

    assertThat(pendingPool.contains(block)).isTrue();
    assertThat(pendingPool.size()).isEqualTo(1);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false))
        .containsExactlyInAnyOrder(block);
    assertThat(requiredRootEvents).containsExactly(parentRoot);
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void add_siblingBlocks() {
    final SignedBeaconBlock blockA =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
    final Bytes32 parentRoot = blockA.getParentRoot();
    final SignedBeaconBlock blockB =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue(), parentRoot);
    pendingPool.add(blockA);
    pendingPool.add(blockB);

    assertThat(pendingPool.contains(blockA)).isTrue();
    assertThat(pendingPool.contains(blockB)).isTrue();
    assertThat(pendingPool.size()).isEqualTo(2);
    assertThat(pendingPool.getItemsDependingOn(blockA.getParentRoot(), false))
        .containsExactlyInAnyOrder(blockA, blockB);
    assertThat(requiredRootEvents).containsExactly(parentRoot);
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void remove_existingBlock() {
    final SignedBeaconBlock block =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
    pendingPool.add(block);
    pendingPool.remove(block);

    assertThat(pendingPool.contains(block)).isFalse();
    assertThat(pendingPool.size()).isEqualTo(0);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false)).isEmpty();
    assertThat(requiredRootEvents).containsExactly(block.getParentRoot());
    assertThat(requiredRootDroppedEvents).containsExactly(block.getParentRoot());
  }

  @Test
  public void remove_unknownBlock() {
    final SignedBeaconBlock block =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
    pendingPool.remove(block);

    assertThat(pendingPool.contains(block)).isFalse();
    assertThat(pendingPool.size()).isEqualTo(0);
    assertThat(pendingPool.getItemsDependingOn(block.getParentRoot(), false)).isEmpty();
    assertThat(requiredRootEvents).isEmpty();
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  public void remove_siblingBlock() {
    final SignedBeaconBlock blockA =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
    final Bytes32 parentRoot = blockA.getParentRoot();
    final SignedBeaconBlock blockB =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue(), parentRoot);
    pendingPool.add(blockA);
    pendingPool.add(blockB);
    pendingPool.remove(blockA);

    assertThat(pendingPool.contains(blockB)).isTrue();
    assertThat(pendingPool.size()).isEqualTo(1);
    assertThat(pendingPool.getItemsDependingOn(blockA.getParentRoot(), false))
        .containsExactlyInAnyOrder(blockB);
    assertThat(requiredRootEvents).containsExactly(parentRoot);
    assertThat(requiredRootDroppedEvents).isEmpty();
  }

  @Test
  void getAllRequiredBlockRoots_shouldExcludeRootsForAlreadyPendingItems() {
    final SignedBeaconBlock blockA =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue());
    final Bytes32 parentRoot = blockA.getParentRoot();
    final SignedBeaconBlock blockB =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue(), parentRoot);
    // Should not be included as blockB is already present
    final SignedBeaconBlock blockC =
        dataStructureUtil.randomSignedBeaconBlock(
            currentSlot.longValue() + 1, blockB.getParentRoot());
    final SignedBeaconBlock blockD =
        dataStructureUtil.randomSignedBeaconBlock(currentSlot.longValue() + 1);

    pendingPool.add(blockA);
    pendingPool.add(blockB);
    pendingPool.add(blockC);
    pendingPool.add(blockD);
    assertThat(pendingPool.size()).isEqualTo(4);

    assertThat(pendingPool.getAllRequiredBlockRoots())
        .containsExactlyInAnyOrder(parentRoot, blockD.getParentRoot());
  }

  @Test
  public void getItemsDependingOn_includeIndirect() {
    final int chainDepth = 2;
    final int descendentChainCount = 2;
    final List<SignedBeaconBlock> directDescendents = new ArrayList<>();
    final List<SignedBeaconBlock> indirectDescendents = new ArrayList<>();

    final Bytes32 commonAncestorRoot = dataStructureUtil.randomBytes32();
    for (int chainIndex = 0; chainIndex < descendentChainCount; chainIndex++) {
      Bytes32 parentRoot = commonAncestorRoot;
      for (int depth = 0; depth < chainDepth; depth++) {
        final long slot = currentSlot.longValue() + 1 + depth;
        final SignedBeaconBlock block = dataStructureUtil.randomSignedBeaconBlock(slot, parentRoot);
        final List<SignedBeaconBlock> blockSet =
            depth == 0 ? directDescendents : indirectDescendents;
        blockSet.add(block);

        parentRoot = block.getMessage().hashTreeRoot();
      }
    }

    indirectDescendents.forEach(pendingPool::add);
    directDescendents.forEach(pendingPool::add);

    List<SignedBeaconBlock> dependentBlocks =
        pendingPool.getItemsDependingOn(commonAncestorRoot, true);
    assertThat(dependentBlocks.size())
        .isEqualTo(directDescendents.size() + indirectDescendents.size());
    assertThat(dependentBlocks).containsAll(directDescendents);
    assertThat(dependentBlocks).containsAll(indirectDescendents);
  }

  @Test
  public void getItemsDependingOn_directOnly() {
    final int chainDepth = 2;
    final int descendentChainCount = 2;
    final List<SignedBeaconBlock> directDescendents = new ArrayList<>();
    final List<SignedBeaconBlock> indirectDescendents = new ArrayList<>();

    final Bytes32 commonAncestorRoot = dataStructureUtil.randomBytes32();
    for (int chainIndex = 0; chainIndex < descendentChainCount; chainIndex++) {
      Bytes32 parentRoot = commonAncestorRoot;
      for (int depth = 0; depth < chainDepth; depth++) {
        final long slot = currentSlot.longValue() + 1 + depth;
        final SignedBeaconBlock block = dataStructureUtil.randomSignedBeaconBlock(slot, parentRoot);
        final List<SignedBeaconBlock> blockSet =
            depth == 0 ? directDescendents : indirectDescendents;
        blockSet.add(block);

        parentRoot = block.getMessage().hashTreeRoot();
      }
    }

    indirectDescendents.forEach(pendingPool::add);
    directDescendents.forEach(pendingPool::add);

    List<SignedBeaconBlock> dependentBlocks =
        pendingPool.getItemsDependingOn(commonAncestorRoot, false);
    assertThat(dependentBlocks).containsExactlyInAnyOrderElementsOf(directDescendents);
  }

  @Test
  public void prune_finalizedBlocks() {
    final SignedBeaconBlock finalizedBlock = dataStructureUtil.randomSignedBeaconBlock(10);
    final Checkpoint checkpoint = finalizedCheckpoint(finalizedBlock);
    final long finalizedSlot = checkpoint.getEpochStartSlot(spec).longValue();
    setSlot(finalizedSlot);

    // Add a bunch of blocks
    List<SignedBeaconBlock> nonFinalBlocks =
        List.of(dataStructureUtil.randomSignedBeaconBlock(finalizedSlot + 1));
    List<SignedBeaconBlock> finalizedBlocks =
        List.of(
            dataStructureUtil.randomSignedBeaconBlock(finalizedSlot),
            dataStructureUtil.randomSignedBeaconBlock(finalizedSlot - 1));
    List<SignedBeaconBlock> allBlocks = new ArrayList<>();
    allBlocks.addAll(nonFinalBlocks);
    allBlocks.addAll(finalizedBlocks);
    nonFinalBlocks.forEach(pendingPool::add);
    finalizedBlocks.forEach(pendingPool::add);

    // Check that all blocks are in the collection
    assertThat(pendingPool.size()).isEqualTo(finalizedBlocks.size() + nonFinalBlocks.size());
    for (SignedBeaconBlock block : allBlocks) {
      assertThat(pendingPool.contains(block)).isTrue();
    }

    // Update finalized checkpoint and prune
    pendingPool.onNewFinalizedCheckpoint(checkpoint, false);
    pendingPool.prune();

    // Check that all final blocks have been pruned
    assertThat(pendingPool.size()).isEqualTo(nonFinalBlocks.size());
    for (SignedBeaconBlock block : nonFinalBlocks) {
      assertThat(pendingPool.contains(block)).isTrue();
    }
  }

  @Test
  public void onSlot_prunesOldAndFinalizedBlocks() {
    // We prune on PRUNE_SLOT every epoch everything
    // older than (current - historicalTolerance) or finalized
    // Let's put some blocks around (pruneSlot - historicalTolerance)
    final UInt64 startSlot =
        AbstractIgnoringFutureHistoricalSlot.PRUNE_SLOT
            .plus(spec.getGenesisSpec().getSlotsPerEpoch())
            .minus(historicalTolerance);
    setSlot(startSlot);

    // Interleave blocks to keep and blocks to prune
    final List<SignedBeaconBlock> blocksToPrune = new ArrayList<>();
    final List<SignedBeaconBlock> blocksToKeep = new ArrayList<>();
    for (int i = 0; i < 3; i++) {
      final SignedBeaconBlock toPrune =
          dataStructureUtil.randomSignedBeaconBlock(startSlot.longValue() - 1L);
      final SignedBeaconBlock toKeep =
          dataStructureUtil.randomSignedBeaconBlock(startSlot.longValue());

      blocksToPrune.add(toPrune);
      blocksToKeep.add(toKeep);
      pendingPool.add(toPrune);
      pendingPool.add(toKeep);
    }

    blocksToKeep.forEach(b -> assertThat(pendingPool.contains(b)).isTrue());
    blocksToPrune.forEach(b -> assertThat(pendingPool.contains(b)).isTrue());

    UInt64 newSlot = startSlot;
    for (int i = 0; i < historicalTolerance.intValue() - 1; i++) {
      newSlot = newSlot.plus(UInt64.ONE);
      pendingPool.onSlot(newSlot);
      blocksToKeep.forEach(b -> assertThat(pendingPool.contains(b)).isTrue());
      blocksToPrune.forEach(b -> assertThat(pendingPool.contains(b)).isTrue());
    }

    // Next slot should prune blocksToPrune
    pendingPool.onSlot(newSlot.plus(UInt64.ONE));

    blocksToKeep.forEach(b -> assertThat(pendingPool.contains(b)).isTrue());
    blocksToPrune.forEach(b -> assertThat(pendingPool.contains(b)).isFalse());

    // Make all finalized
    pendingPool.onNewFinalizedCheckpoint(dataStructureUtil.randomCheckpoint(1), false);
    pendingPool.onSlot(newSlot.plus(UInt64.ONE));
    blocksToKeep.forEach(b -> assertThat(pendingPool.contains(b)).isFalse());
  }
}
