text: import logging
  from typing import Dict, List, Optional, Tuple
  from replaceme.util.ints import uint32
  from replaceme.types.blockchain_format.sized_bytes import bytes32
  from replaceme.types.blockchain_format.sub_epoch_summary import SubEpochSummary
  from pathlib import Path
  import aiofiles
  from dataclasses import dataclass
  from replaceme.util.streamable import Streamable, streamable
  from replaceme.util.files import write_file_async
  from replaceme.util.db_wrapper import DBWrapper
  
  log = logging.getLogger(__name__)
  
  
  @dataclass(frozen=True)
  @streamable
  class SesCache(Streamable):
      content: List[Tuple[uint32, bytes]]
  
  
  class BlockHeightMap:
      db: DBWrapper
  
      # the below dictionaries are loaded from the database, from the peak
      # and back in time on startup.
  
      # Defines the path from genesis to the peak, no orphan blocks
      # this buffer contains all block hashes that are part of the current peak
      # ordered by height. i.e. __height_to_hash[0..32] is the genesis hash
      # __height_to_hash[32..64] is the hash for height 1 and so on
      __height_to_hash: bytearray
  
      # All sub-epoch summaries that have been included in the blockchain from the beginning until and including the peak
      # (height_included, SubEpochSummary). Note: ONLY for the blocks in the path to the peak
      # The value is a serialized SubEpochSummary object
      __sub_epoch_summaries: Dict[uint32, bytes]
  
      # count how many blocks have been added since the cache was last written to
      # disk
      __dirty: int
  
      # the file we're saving the height-to-hash cache to
      __height_to_hash_filename: Path
  
      # the file we're saving the sub epoch summary cache to
      __ses_filename: Path
  
      @classmethod
      async def create(cls, blockchain_dir: Path, db: DBWrapper) -> "BlockHeightMap":
          self = BlockHeightMap()
          self.db = db
  
          self.__dirty = 0
          self.__height_to_hash = bytearray()
          self.__sub_epoch_summaries = {}
          self.__height_to_hash_filename = blockchain_dir / "height-to-hash"
          self.__ses_filename = blockchain_dir / "sub-epoch-summaries"
  
          res = await self.db.db.execute(
              "SELECT header_hash,prev_hash,height,sub_epoch_summary from block_records WHERE is_peak=1"
          )
          row = await res.fetchone()
          await res.close()
  
          if row is None:
              return self
  
          try:
              async with aiofiles.open(self.__height_to_hash_filename, "rb") as f:
                  self.__height_to_hash = bytearray(await f.read())
          except Exception:
              # it's OK if this file doesn't exist, we can rebuild it
              pass
  
          try:
              async with aiofiles.open(self.__ses_filename, "rb") as f:
                  self.__sub_epoch_summaries = {k: v for (k, v) in SesCache.from_bytes(await f.read()).content}
          except Exception:
              # it's OK if this file doesn't exist, we can rebuild it
              pass
  
          peak: bytes32
          prev_hash: bytes32
          if db.db_version == 2:
              peak = row[0]
              prev_hash = row[1]
          else:
              # TODO: address hint errors and remove ignores
              #       error: Incompatible types in assignment (expression has type "bytes", variable has type "bytes32")
              #       [assignment]
              peak = bytes.fromhex(row[0])  # type: ignore[assignment]
              # TODO: address hint errors and remove ignores
              #       error: Incompatible types in assignment (expression has type "bytes", variable has type "bytes32")
              #       [assignment]
              prev_hash = bytes.fromhex(row[1])  # type: ignore[assignment]
          height = row[2]
  
          # allocate memory for height to hash map
          # this may also truncate it, if thie file on disk had an invalid size
          new_size = (height + 1) * 32
          size = len(self.__height_to_hash)
          if size > new_size:
              del self.__height_to_hash[new_size:]
          else:
              self.__height_to_hash += bytearray([0] * (new_size - size))
  
          # if the peak hash is already in the height-to-hash map, we don't need
          # to load anything more from the DB
          if self.get_hash(height) != peak:
              self.__set_hash(height, peak)
  
              if row[3] is not None:
                  self.__sub_epoch_summaries[height] = row[3]
  
              # prepopulate the height -> hash mapping
              await self._load_blocks_from(height, prev_hash)
  
              await self.maybe_flush()
  
          return self
  
      def update_height(self, height: uint32, header_hash: bytes32, ses: Optional[SubEpochSummary]):
          # we're only updating the last hash. If we've reorged, we already rolled
          # back, making this the new peak
          idx = height * 32
          assert idx <= len(self.__height_to_hash)
          self.__height_to_hash[idx : idx + 32] = header_hash
          if ses is not None:
              self.__sub_epoch_summaries[height] = bytes(ses)
  
      async def maybe_flush(self):
          if self.__dirty < 1000:
              return
  
          assert (len(self.__height_to_hash) % 32) == 0
          map_buf = self.__height_to_hash.copy()
  
          ses_buf = bytes(SesCache([(k, v) for (k, v) in self.__sub_epoch_summaries.items()]))
  
          self.__dirty = 0
  
          await write_file_async(self.__height_to_hash_filename, map_buf)
          await write_file_async(self.__ses_filename, ses_buf)
  
      # load height-to-hash map entries from the DB starting at height back in
      # time until we hit a match in the existing map, at which point we can
      # assume all previous blocks have already been populated
      async def _load_blocks_from(self, height: uint32, prev_hash: bytes32):
  
          while height > 0:
              # load 5000 blocks at a time
              window_end = max(0, height - 5000)
              cursor = await self.db.db.execute(
                  "SELECT header_hash,prev_hash,height,sub_epoch_summary from block_records "
                  "INDEXED BY height WHERE height>=? AND height <?",
                  (window_end, height),
              )
  
              rows = await cursor.fetchall()
              await cursor.close()
  
              # maps block-hash -> (height, prev-hash, sub-epoch-summary)
              ordered: Dict[bytes32, Tuple[uint32, bytes32, Optional[bytes]]] = {}
  
              if self.db.db_version == 2:
                  for r in rows:
                      ordered[r[0]] = (r[2], r[1], r[3])
              else:
                  for r in rows:
                      # TODO: address hint errors and remove ignores
                      #       error: Invalid index type "bytes" for "Dict[bytes32, Tuple[uint32, bytes32,
                      #          Optional[bytes]]]";
                      #       expected type "bytes32"  [index]
                      #       error: Incompatible types in assignment (expression has type "Tuple[Any,
                      #          bytes, Any]", target
                      #       has type "Tuple[uint32, bytes32, Optional[bytes]]")  [assignment]
                      ordered[bytes.fromhex(r[0])] = (r[2], bytes.fromhex(r[1]), r[3])  # type: ignore[index,assignment]
  
              while height > window_end:
                  entry = ordered[prev_hash]
                  assert height == entry[0] + 1
                  height = entry[0]
                  if entry[2] is not None:
  
                      if (
                          self.get_hash(height) == prev_hash
                          and height in self.__sub_epoch_summaries
                          and self.__sub_epoch_summaries[height] == entry[2]
                      ):
                          return
                      self.__sub_epoch_summaries[height] = entry[2]
                  self.__set_hash(height, prev_hash)
                  prev_hash = entry[1]
  
      def __set_hash(self, height: int, block_hash: bytes32):
          idx = height * 32
          self.__height_to_hash[idx : idx + 32] = block_hash
          self.__dirty += 1
  
      def get_hash(self, height: uint32) -> bytes32:
          idx = height * 32
          assert idx + 32 <= len(self.__height_to_hash)
          # TODO: address hint errors and remove ignores
          #       error: Incompatible return value type (got "bytes", expected "bytes32")  [return-value]
          return bytes(self.__height_to_hash[idx : idx + 32])  # type: ignore[return-value]
  
      def contains_height(self, height: uint32) -> bool:
          return height * 32 < len(self.__height_to_hash)
  
      def rollback(self, fork_height: int):
          # fork height may be -1, in which case all blocks are different and we
          # should clear all sub epoch summaries
          heights_to_delete = []
          for ses_included_height in self.__sub_epoch_summaries.keys():
              if ses_included_height > fork_height:
                  heights_to_delete.append(ses_included_height)
          for height in heights_to_delete:
              del self.__sub_epoch_summaries[height]
  
      def get_ses(self, height: uint32) -> SubEpochSummary:
          return SubEpochSummary.from_bytes(self.__sub_epoch_summaries[height])
  
      def get_ses_heights(self) -> List[uint32]:
          return sorted(self.__sub_epoch_summaries.keys())
  