text: import asyncio
  import logging
  from typing import List, Optional, Set, Tuple
  
  import aiosqlite
  from blspy import G1Element
  
  from stai.types.blockchain_format.sized_bytes import bytes32
  from stai.util.db_wrapper import DBWrapper
  from stai.util.ints import uint32
  from stai.wallet.derivation_record import DerivationRecord
  from stai.wallet.util.wallet_types import WalletType
  
  log = logging.getLogger(__name__)
  
  
  class WalletPuzzleStore:
      """
      WalletPuzzleStore keeps track of all generated puzzle_hashes and their derivation path / wallet.
      This is only used for HD wallets where each address is derived from a public key. Otherwise, use the
      WalletInterestedStore to keep track of puzzle hashes which we are interested in.
      """
  
      db_connection: aiosqlite.Connection
      lock: asyncio.Lock
      cache_size: uint32
      all_puzzle_hashes: Set[bytes32]
      db_wrapper: DBWrapper
  
      @classmethod
      async def create(cls, db_wrapper: DBWrapper, cache_size: uint32 = uint32(600000)):
          self = cls()
  
          self.cache_size = cache_size
  
          self.db_wrapper = db_wrapper
          self.db_connection = self.db_wrapper.db
          await self.db_connection.execute(
              (
                  "CREATE TABLE IF NOT EXISTS derivation_paths("
                  "derivation_index int,"
                  " pubkey text,"
                  " puzzle_hash text PRIMARY_KEY,"
                  " wallet_type int,"
                  " wallet_id int,"
                  " used tinyint)"
              )
          )
          await self.db_connection.execute(
              "CREATE INDEX IF NOT EXISTS derivation_index_index on derivation_paths(derivation_index)"
          )
  
          await self.db_connection.execute("CREATE INDEX IF NOT EXISTS ph on derivation_paths(puzzle_hash)")
  
          await self.db_connection.execute("CREATE INDEX IF NOT EXISTS pubkey on derivation_paths(pubkey)")
  
          await self.db_connection.execute("CREATE INDEX IF NOT EXISTS wallet_type on derivation_paths(wallet_type)")
  
          await self.db_connection.execute("CREATE INDEX IF NOT EXISTS wallet_id on derivation_paths(wallet_id)")
  
          await self.db_connection.execute("CREATE INDEX IF NOT EXISTS used on derivation_paths(wallet_type)")
  
          await self.db_connection.commit()
          # Lock
          self.lock = asyncio.Lock()  # external
          await self._init_cache()
          return self
  
      async def close(self):
          await self.db_connection.close()
  
      async def _init_cache(self):
          self.all_puzzle_hashes = await self.get_all_puzzle_hashes()
  
      async def _clear_database(self):
          cursor = await self.db_connection.execute("DELETE FROM derivation_paths")
          await cursor.close()
          await self.db_connection.commit()
  
      async def add_derivation_paths(self, records: List[DerivationRecord], in_transaction=False) -> None:
          """
          Insert many derivation paths into the database.
          """
  
          if not in_transaction:
              await self.db_wrapper.lock.acquire()
          try:
              sql_records = []
              for record in records:
                  self.all_puzzle_hashes.add(record.puzzle_hash)
                  sql_records.append(
                      (
                          record.index,
                          bytes(record.pubkey).hex(),
                          record.puzzle_hash.hex(),
                          record.wallet_type,
                          record.wallet_id,
                          0,
                      ),
                  )
  
              cursor = await self.db_connection.executemany(
                  "INSERT OR REPLACE INTO derivation_paths VALUES(?, ?, ?, ?, ?, ?)",
                  sql_records,
              )
  
              await cursor.close()
          finally:
              if not in_transaction:
                  await self.db_connection.commit()
                  self.db_wrapper.lock.release()
  
      async def get_derivation_record(self, index: uint32, wallet_id: uint32) -> Optional[DerivationRecord]:
          """
          Returns the derivation record by index and wallet id.
          """
          cursor = await self.db_connection.execute(
              "SELECT * FROM derivation_paths WHERE derivation_index=? and wallet_id=?;",
              (
                  index,
                  wallet_id,
              ),
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None and row[0] is not None:
              return DerivationRecord(
                  uint32(row[0]),
                  bytes32.fromhex(row[2]),
                  G1Element.from_bytes(bytes.fromhex(row[1])),
                  WalletType(row[3]),
                  uint32(row[4]),
              )
  
          return None
  
      async def get_derivation_record_for_puzzle_hash(self, puzzle_hash: bytes32) -> Optional[DerivationRecord]:
          """
          Returns the derivation record by index and wallet id.
          """
          cursor = await self.db_connection.execute(
              "SELECT * FROM derivation_paths WHERE puzzle_hash=?;",
              (puzzle_hash.hex(),),
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None and row[0] is not None:
              return DerivationRecord(
                  uint32(row[0]),
                  bytes32.fromhex(row[2]),
                  G1Element.from_bytes(bytes.fromhex(row[1])),
                  WalletType(row[3]),
                  uint32(row[4]),
              )
  
          return None
  
      async def set_used_up_to(self, index: uint32, in_transaction=False) -> None:
          """
          Sets a derivation path to used so we don't use it again.
          """
  
          if not in_transaction:
              await self.db_wrapper.lock.acquire()
          try:
              cursor = await self.db_connection.execute(
                  "UPDATE derivation_paths SET used=1 WHERE derivation_index<=?",
                  (index,),
              )
              await cursor.close()
          finally:
              if not in_transaction:
                  await self.db_connection.commit()
                  self.db_wrapper.lock.release()
  
      async def puzzle_hash_exists(self, puzzle_hash: bytes32) -> bool:
          """
          Checks if passed puzzle_hash is present in the db.
          """
  
          cursor = await self.db_connection.execute(
              "SELECT * from derivation_paths WHERE puzzle_hash=?", (puzzle_hash.hex(),)
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          return row is not None
  
      async def one_of_puzzle_hashes_exists(self, puzzle_hashes: List[bytes32]) -> bool:
          """
          Checks if one of the passed puzzle_hashes is present in the db.
          """
          if len(puzzle_hashes) < 1:
              return False
  
          for ph in puzzle_hashes:
              if ph in self.all_puzzle_hashes:
                  return True
  
          return False
  
      async def index_for_pubkey(self, pubkey: G1Element) -> Optional[uint32]:
          """
          Returns derivation paths for the given pubkey.
          Returns None if not present.
          """
  
          cursor = await self.db_connection.execute(
              "SELECT * from derivation_paths WHERE pubkey=?", (bytes(pubkey).hex(),)
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None:
              return uint32(row[0])
  
          return None
  
      async def index_for_puzzle_hash(self, puzzle_hash: bytes32) -> Optional[uint32]:
          """
          Returns the derivation path for the puzzle_hash.
          Returns None if not present.
          """
          cursor = await self.db_connection.execute(
              "SELECT * from derivation_paths WHERE puzzle_hash=?", (puzzle_hash.hex(),)
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None:
              return uint32(row[0])
  
          return None
  
      async def index_for_puzzle_hash_and_wallet(self, puzzle_hash: bytes32, wallet_id: uint32) -> Optional[uint32]:
          """
          Returns the derivation path for the puzzle_hash.
          Returns None if not present.
          """
          cursor = await self.db_connection.execute(
              "SELECT * from derivation_paths WHERE puzzle_hash=? and wallet_id=?;",
              (
                  puzzle_hash.hex(),
                  wallet_id,
              ),
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None:
              return uint32(row[0])
  
          return None
  
      async def wallet_info_for_puzzle_hash(self, puzzle_hash: bytes32) -> Optional[Tuple[uint32, WalletType]]:
          """
          Returns the derivation path for the puzzle_hash.
          Returns None if not present.
          """
  
          cursor = await self.db_connection.execute(
              "SELECT * from derivation_paths WHERE puzzle_hash=?", (puzzle_hash.hex(),)
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None:
              return row[4], WalletType(row[3])
  
          return None
  
      async def get_all_puzzle_hashes(self) -> Set[bytes32]:
          """
          Return a set containing all puzzle_hashes we generated.
          """
  
          cursor = await self.db_connection.execute("SELECT * from derivation_paths")
          rows = await cursor.fetchall()
          await cursor.close()
          result: Set[bytes32] = set()
  
          for row in rows:
              result.add(bytes32(bytes.fromhex(row[2])))
  
          return result
  
      async def get_last_derivation_path(self) -> Optional[uint32]:
          """
          Returns the last derivation path by derivation_index.
          """
  
          cursor = await self.db_connection.execute("SELECT MAX(derivation_index) FROM derivation_paths;")
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None and row[0] is not None:
              return uint32(row[0])
  
          return None
  
      async def get_last_derivation_path_for_wallet(self, wallet_id: int) -> Optional[uint32]:
          """
          Returns the last derivation path by derivation_index.
          """
  
          cursor = await self.db_connection.execute(
              f"SELECT MAX(derivation_index) FROM derivation_paths WHERE wallet_id={wallet_id};"
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None and row[0] is not None:
              return uint32(row[0])
  
          return None
  
      async def get_current_derivation_record_for_wallet(self, wallet_id: uint32) -> Optional[DerivationRecord]:
          """
          Returns the current derivation record by derivation_index.
          """
  
          cursor = await self.db_connection.execute(
              f"SELECT MAX(derivation_index) FROM derivation_paths WHERE wallet_id={wallet_id} and used=1;"
          )
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None and row[0] is not None:
              index = uint32(row[0])
              return await self.get_derivation_record(index, wallet_id)
  
          return None
  
      async def get_unused_derivation_path(self) -> Optional[uint32]:
          """
          Returns the first unused derivation path by derivation_index.
          """
          cursor = await self.db_connection.execute("SELECT MIN(derivation_index) FROM derivation_paths WHERE used=0;")
          row = await cursor.fetchone()
          await cursor.close()
  
          if row is not None and row[0] is not None:
              return uint32(row[0])
  
          return None
  