text: from typing import List, Optional, Tuple
  
  from fork.consensus.block_record import BlockRecord
  from fork.consensus.blockchain_interface import BlockchainInterface
  from fork.consensus.constants import ConsensusConstants
  from fork.types.blockchain_format.sized_bytes import bytes32
  from fork.util.ints import uint8, uint32, uint64, uint128
  from fork.util.significant_bits import count_significant_bits, truncate_to_significant_bits
  
  
  def _get_blocks_at_height(
      blocks: BlockchainInterface,
      prev_b: BlockRecord,
      target_height: uint32,
      max_num_blocks: uint32 = uint32(1),
  ) -> List[BlockRecord]:
      """
      Return a consecutive list of BlockRecords starting at target_height, returning a maximum of
      max_num_blocks. Assumes all block records are present. Does a slot linear search, if the blocks are not
      in the path of the peak. Can only fetch ancestors of prev_b.
  
      Args:
          blocks: dict from header hash to BlockRecord.
          prev_b: prev_b (to start backwards search).
          target_height: target block to start
          max_num_blocks: max number of blocks to fetch (although less might be fetched)
  
      """
      if blocks.contains_height(prev_b.height):
          header_hash = blocks.height_to_hash(prev_b.height)
          if header_hash == prev_b.header_hash:
              # Efficient fetching, since we are fetching ancestor blocks within the heaviest chain. We can directly
              # use the height_to_block_record method
              block_list: List[BlockRecord] = []
              for h in range(target_height, target_height + max_num_blocks):
                  assert blocks.contains_height(uint32(h))
                  block_list.append(blocks.height_to_block_record(uint32(h)))
              return block_list
  
      # Slow fetching, goes back one by one, since we are in a fork
      curr_b: BlockRecord = prev_b
      target_blocks = []
      while curr_b.height >= target_height:
          if curr_b.height < target_height + max_num_blocks:
              target_blocks.append(curr_b)
          if curr_b.height == 0:
              break
          curr_b = blocks.block_record(curr_b.prev_hash)
      return list(reversed(target_blocks))
  
  
  def _get_second_to_last_transaction_block_in_previous_epoch(
      constants: ConsensusConstants,
      blocks: BlockchainInterface,
      last_b: BlockRecord,
  ) -> BlockRecord:
      """
      Retrieves the second to last transaction block in the previous epoch.
  
      Args:
          constants: consensus constants being used for this chain
          blocks: dict from header hash to block of all relevant blocks
          last_b: last-block in the current epoch, or last block we have seen, if potentially finishing epoch soon
  
             prev epoch surpassed  prev epoch started                  epoch sur.  epoch started
              v                       v                                v         v
        |.B...B....B. B....B...|......B....B.....B...B.|.B.B.B..|..B...B.B.B...|.B.B.B. B.|........
              PREV EPOCH                 CURR EPOCH                               NEW EPOCH
  
       The blocks selected for the timestamps are the second to last transaction blocks in each epoch.
       Block at height 0 is an exception. Note that H mod EPOCH_BLOCKS where H is the height of the first block in the
       epoch, must be >= 0, and < 128.
      """
  
      # This height is guaranteed to be in the next epoch (even when last_b is not actually the last block)
      height_in_next_epoch = (
          last_b.height + 2 * constants.MAX_SUB_SLOT_BLOCKS + constants.MIN_BLOCKS_PER_CHALLENGE_BLOCK + 5
      )
      height_epoch_surpass: uint32 = uint32(height_in_next_epoch - (height_in_next_epoch % constants.EPOCH_BLOCKS))
      height_prev_epoch_surpass: uint32 = uint32(height_epoch_surpass - constants.EPOCH_BLOCKS)
  
      assert height_prev_epoch_surpass % constants.EPOCH_BLOCKS == height_prev_epoch_surpass % constants.EPOCH_BLOCKS == 0
  
      # Sanity check, don't go too far past epoch barrier
      assert (height_in_next_epoch - height_epoch_surpass) < (5 * constants.MAX_SUB_SLOT_BLOCKS)
  
      if height_prev_epoch_surpass == 0:
          # The genesis block is an edge case, where we measure from the first block in epoch (height 0), as opposed to
          # a block in the previous epoch, which would be height < 0
          return _get_blocks_at_height(blocks, last_b, uint32(0))[0]
  
      # If the prev slot is the first slot, the iterations start at 0
      # We will compute the timestamps of the 2nd to last block in epoch, as well as the total iterations at infusion
      prev_slot_start_iters: uint128
      prev_slot_time_start: uint64
  
      # The target block must be in this range. Either the surpass block must be a transaction block, or something
      # in it's sub slot must be a transaction block. If that is the only transaction block in the sub-slot, the last
      # block in the previous sub-slot from that must also be a transaction block (therefore -1 is used).
      # The max height for the new epoch to start is surpass + 2*MAX_SUB_SLOT_BLOCKS + MIN_BLOCKS_PER_CHALLENGE_BLOCK - 3,
      # since we might have a deficit > 0 when surpass is hit. The +3 is added just in case
      fetched_blocks = _get_blocks_at_height(
          blocks,
          last_b,
          uint32(height_prev_epoch_surpass - constants.MAX_SUB_SLOT_BLOCKS - 1),
          uint32(3 * constants.MAX_SUB_SLOT_BLOCKS + constants.MIN_BLOCKS_PER_CHALLENGE_BLOCK + 3),
      )
  
      # We want to find the last block in the slot at which we surpass the height.
      # The last block in epoch will be before this.
      fetched_index: int = constants.MAX_SUB_SLOT_BLOCKS
      curr_b: BlockRecord = fetched_blocks[fetched_index]
      fetched_index += 1
      assert curr_b.height == height_prev_epoch_surpass - 1
      next_b: BlockRecord = fetched_blocks[fetched_index]
      assert next_b.height == height_prev_epoch_surpass
  
      # Wait until the slot finishes with a challenge chain infusion at start of slot
      # Note that there are no overflow blocks at the start of new epochs
      while next_b.sub_epoch_summary_included is None:
          curr_b = next_b
          next_b = fetched_blocks[fetched_index]
          fetched_index += 1
  
      # Backtrack to find the second to last tx block
      found_tx_block = 1 if curr_b.is_transaction_block else 0
      while found_tx_block < 2:
          curr_b = blocks.block_record(curr_b.prev_hash)
          if curr_b.is_transaction_block:
              found_tx_block += 1
  
      return curr_b
  
  
  def height_can_be_first_in_epoch(constants: ConsensusConstants, height: uint32) -> bool:
      return (height - (height % constants.SUB_EPOCH_BLOCKS)) % constants.EPOCH_BLOCKS == 0
  
  
  def can_finish_sub_and_full_epoch(
      constants: ConsensusConstants,
      blocks: BlockchainInterface,
      height: uint32,
      prev_header_hash: Optional[bytes32],
      deficit: uint8,
      block_at_height_included_ses: bool,
  ) -> Tuple[bool, bool]:
      """
      Returns a bool tuple
      first bool is true if the next sub-slot after height will form part of a new sub-epoch. Therefore
      block height is the last block, and height + 1 is in a new sub-epoch.
      second bool is true if the next sub-slot after height will form part of a new sub-epoch and epoch.
      Therefore, block height is the last block, and height + 1 is in a new epoch.
  
      Args:
          constants: consensus constants being used for this chain
          blocks: dictionary from header hash to SBR of all included SBR
          height: block height of the (potentially) last block in the sub-epoch
          prev_header_hash: prev_header hash of the block at height, assuming not genesis
          deficit: deficit of block at height height
          block_at_height_included_ses: whether or not the block at height height already included a SES
      """
  
      if height < constants.SUB_EPOCH_BLOCKS - 1:
          return False, False
  
      assert prev_header_hash is not None
  
      if deficit > 0:
          return False, False
  
      if block_at_height_included_ses:
          # If we just included a sub_epoch_summary, we cannot include one again
          return False, False
  
      # This does not check the two edge cases where (height + 1) % constants.SUB_EPOCH_BLOCKS is 0 or 1
      # If it's 0, height+1 is the first place that a sub-epoch can be included
      # If it's 1, we just checked whether 0 included it in the previous check
      if (height + 1) % constants.SUB_EPOCH_BLOCKS > 1:
          curr: BlockRecord = blocks.block_record(prev_header_hash)
          while curr.height % constants.SUB_EPOCH_BLOCKS > 0:
              if curr.sub_epoch_summary_included is not None:
                  return False, False
              curr = blocks.block_record(curr.prev_hash)
  
          if curr.sub_epoch_summary_included is not None:
              return False, False
  
      # For checking new epoch, make sure the epoch blocks are aligned
      return True, height_can_be_first_in_epoch(constants, uint32(height + 1))
  
  
  def _get_next_sub_slot_iters(
      constants: ConsensusConstants,
      blocks: BlockchainInterface,
      prev_header_hash: bytes32,
      height: uint32,
      curr_sub_slot_iters: uint64,
      deficit: uint8,
      block_at_height_included_ses: bool,
      new_slot: bool,
      signage_point_total_iters: uint128,
      skip_epoch_check=False,
  ) -> uint64:
      """
      Returns the slot iterations required for the next block after the one at height, where new_slot is true
      iff the next block will be in the next slot. WARNING: assumes that the block at height is not the first block
      in a sub-epoch.
  
      Args:
          constants: consensus constants being used for this chain
          blocks: dictionary from header hash to SBR of all included SBR
          prev_header_hash: header hash of the previous block
          height: the block height of the block to look at
          curr_sub_slot_iters: sub-slot iters at the infusion point of the block at height
          deficit: deficit of block at height height
          new_slot: whether or not there is a new slot after height
          signage_point_total_iters: signage point iters of the block at height
          skip_epoch_check: don't check correct epoch
      """
      next_height: uint32 = uint32(height + 1)
  
      if next_height < constants.EPOCH_BLOCKS:
          return uint64(constants.SUB_SLOT_ITERS_STARTING)
  
      if not blocks.contains_block(prev_header_hash):
          raise ValueError(f"Header hash {prev_header_hash} not in blocks")
  
      prev_b: BlockRecord = blocks.block_record(prev_header_hash)
  
      # If we are in the same epoch, return same ssi
      if not skip_epoch_check:
          _, can_finish_epoch = can_finish_sub_and_full_epoch(
              constants, blocks, height, prev_header_hash, deficit, block_at_height_included_ses
          )
          if not new_slot or not can_finish_epoch:
              return curr_sub_slot_iters
  
      last_block_prev: BlockRecord = _get_second_to_last_transaction_block_in_previous_epoch(constants, blocks, prev_b)
  
      # This gets the last transaction block before this block's signage point. Assuming the block at height height
      # is the last block infused in the epoch: If this block ends up being a
      # transaction block, then last_block_curr will be the second to last tx block in the epoch. If this block
      # is not a transaction block, that means there was exactly one other tx block included in between our signage
      # point and infusion point, and therefore last_block_curr is the second to last as well.
      last_block_curr = prev_b
      while last_block_curr.total_iters > signage_point_total_iters or not last_block_curr.is_transaction_block:
          last_block_curr = blocks.block_record(last_block_curr.prev_hash)
      assert last_block_curr.timestamp is not None and last_block_prev.timestamp is not None
  
      # This is computed as the iterations per second in last epoch, times the target number of seconds per slot
      new_ssi_precise: uint64 = uint64(
          constants.SUB_SLOT_TIME_TARGET
          * (last_block_curr.total_iters - last_block_prev.total_iters)
          // (last_block_curr.timestamp - last_block_prev.timestamp)
      )
  
      # Only change by a max factor as a sanity check
      max_ssi = uint64(constants.DIFFICULTY_CHANGE_MAX_FACTOR * last_block_curr.sub_slot_iters)
      min_ssi = uint64(last_block_curr.sub_slot_iters // constants.DIFFICULTY_CHANGE_MAX_FACTOR)
      if new_ssi_precise >= last_block_curr.sub_slot_iters:
          new_ssi_precise = uint64(min(new_ssi_precise, max_ssi))
      else:
          new_ssi_precise = uint64(max([constants.NUM_SPS_SUB_SLOT, new_ssi_precise, min_ssi]))
  
      new_ssi = truncate_to_significant_bits(new_ssi_precise, constants.SIGNIFICANT_BITS)
      new_ssi = uint64(new_ssi - new_ssi % constants.NUM_SPS_SUB_SLOT)  # Must divide the sub slot
      assert count_significant_bits(new_ssi) <= constants.SIGNIFICANT_BITS
      return new_ssi
  
  
  def _get_next_difficulty(
      constants: ConsensusConstants,
      blocks: BlockchainInterface,
      prev_header_hash: bytes32,
      height: uint32,
      current_difficulty: uint64,
      deficit: uint8,
      block_at_height_included_ses: bool,
      new_slot: bool,
      signage_point_total_iters: uint128,
      skip_epoch_check=False,
  ) -> uint64:
      """
      Returns the difficulty of the next block that extends onto block.
      Used to calculate the number of iterations. WARNING: assumes that the block at height is not the first block
      in a sub-epoch.
  
      Args:
          constants: consensus constants being used for this chain
          blocks: dictionary from header hash to SBR of all included SBR
          prev_header_hash: header hash of the previous block
          height: the block height of the block to look at
          deficit: deficit of block at height height
          current_difficulty: difficulty at the infusion point of the block at height
          new_slot: whether or not there is a new slot after height
          signage_point_total_iters: signage point iters of the block at height
          skip_epoch_check: don't check correct epoch
      """
      next_height: uint32 = uint32(height + 1)
  
      if next_height < (constants.EPOCH_BLOCKS - 3 * constants.MAX_SUB_SLOT_BLOCKS):
          # We are in the first epoch
          return uint64(constants.DIFFICULTY_STARTING)
  
      if not blocks.contains_block(prev_header_hash):
          raise ValueError(f"Header hash {prev_header_hash} not in blocks")
  
      prev_b: BlockRecord = blocks.block_record(prev_header_hash)
  
      # If we are in the same slot as previous block, return same difficulty
      if not skip_epoch_check:
          _, can_finish_epoch = can_finish_sub_and_full_epoch(
              constants, blocks, height, prev_header_hash, deficit, block_at_height_included_ses
          )
          if not new_slot or not can_finish_epoch:
              return current_difficulty
  
      last_block_prev: BlockRecord = _get_second_to_last_transaction_block_in_previous_epoch(constants, blocks, prev_b)
  
      # This gets the last transaction block before this block's signage point. Assuming the block at height height
      # is the last block infused in the epoch: If this block ends up being a
      # transaction block, then last_block_curr will be the second to last tx block in the epoch. If this block
      # is not a transaction block, that means there was exactly one other tx block included in between our signage
      # point and infusion point, and therefore last_block_curr is the second to last as well.
      last_block_curr = prev_b
      while last_block_curr.total_iters > signage_point_total_iters or not last_block_curr.is_transaction_block:
          last_block_curr = blocks.block_record(last_block_curr.prev_hash)
  
      assert last_block_curr.timestamp is not None
      assert last_block_prev.timestamp is not None
      actual_epoch_time: uint64 = uint64(last_block_curr.timestamp - last_block_prev.timestamp)
  
      old_difficulty = uint64(prev_b.weight - blocks.block_record(prev_b.prev_hash).weight)
  
      # Terms are rearranged so there is only one division.
      new_difficulty_precise = uint64(
          (last_block_curr.weight - last_block_prev.weight)
          * constants.SUB_SLOT_TIME_TARGET
          // (constants.SLOT_BLOCKS_TARGET * actual_epoch_time)
      )
  
      # Only change by a max factor, to prevent attacks, as in greenpaper, and must be at least 1
      max_diff = uint64(constants.DIFFICULTY_CHANGE_MAX_FACTOR * old_difficulty)
      min_diff = uint64(old_difficulty // constants.DIFFICULTY_CHANGE_MAX_FACTOR)
  
      if new_difficulty_precise >= old_difficulty:
          new_difficulty_precise = uint64(min(new_difficulty_precise, max_diff))
      else:
          new_difficulty_precise = uint64(max([uint64(1), new_difficulty_precise, min_diff]))
      new_difficulty = truncate_to_significant_bits(new_difficulty_precise, constants.SIGNIFICANT_BITS)
      assert count_significant_bits(new_difficulty) <= constants.SIGNIFICANT_BITS
      return uint64(new_difficulty)
  
  
  def get_next_sub_slot_iters_and_difficulty(
      constants: ConsensusConstants,
      is_first_in_sub_slot: bool,
      prev_b: Optional[BlockRecord],
      blocks: BlockchainInterface,
  ) -> Tuple[uint64, uint64]:
      """
      Retrieves the current sub_slot iters and difficulty of the next block after prev_b.
  
      Args:
          constants: consensus constants being used for this chain
          is_first_in_sub_slot: Whether the next block is the first in the sub slot
          prev_b: the previous block (last block in the epoch)
          blocks: dictionary from header hash to SBR of all included SBR
  
      """
  
      # genesis
      if prev_b is None:
          return constants.SUB_SLOT_ITERS_STARTING, constants.DIFFICULTY_STARTING
  
      if prev_b.height != 0:
          prev_difficulty: uint64 = uint64(prev_b.weight - blocks.block_record(prev_b.prev_hash).weight)
      else:
          # prev block is genesis
          prev_difficulty = uint64(prev_b.weight)
  
      if prev_b.sub_epoch_summary_included is not None:
          return prev_b.sub_slot_iters, prev_difficulty
  
      sp_total_iters = prev_b.sp_total_iters(constants)
      difficulty: uint64 = _get_next_difficulty(
          constants,
          blocks,
          prev_b.prev_hash,
          prev_b.height,
          prev_difficulty,
          prev_b.deficit,
          False,  # Already checked above
          is_first_in_sub_slot,
          sp_total_iters,
      )
  
      sub_slot_iters: uint64 = _get_next_sub_slot_iters(
          constants,
          blocks,
          prev_b.prev_hash,
          prev_b.height,
          prev_b.sub_slot_iters,
          prev_b.deficit,
          False,  # Already checked above
          is_first_in_sub_slot,
          sp_total_iters,
      )
  
      return sub_slot_iters, difficulty
  