From df37902a77cc64c440b06198dc472b1e030b3437 Mon Sep 17 00:00:00 2001
From: 0xKheops <26880866+0xKheops@users.noreply.github.com>
Date: Thu, 12 Dec 2024 16:33:43 +0900
Subject: [PATCH] feat: change hook to read stake by hotkey from balance
 instead of chain

---
 .../bittensor/useGetBittensorStakeByHotKey.ts | 82 ++-----------------
 .../Staking/shared/useGetUnbondInfo.ts        |  2 +-
 2 files changed, 10 insertions(+), 74 deletions(-)

diff --git a/apps/extension/src/ui/domains/Staking/hooks/bittensor/useGetBittensorStakeByHotKey.ts b/apps/extension/src/ui/domains/Staking/hooks/bittensor/useGetBittensorStakeByHotKey.ts
index 0fc64dd06e..d1e7a1ca2f 100644
--- a/apps/extension/src/ui/domains/Staking/hooks/bittensor/useGetBittensorStakeByHotKey.ts
+++ b/apps/extension/src/ui/domains/Staking/hooks/bittensor/useGetBittensorStakeByHotKey.ts
@@ -1,6 +1,6 @@
-import { useQueries, useQuery } from "@tanstack/react-query"
+import { useMemo } from "react"
 
-import { useScaleApi } from "@ui/hooks/sapi/useScaleApi"
+import { useBalance } from "@ui/state"
 
 type GetBittensorStakeByHotKey = {
   address: string | null | undefined
@@ -8,76 +8,12 @@ type GetBittensorStakeByHotKey = {
   isEnabled?: boolean
 }
 
-type GetBittensorStakeByHotKeys = Omit<GetBittensorStakeByHotKey, "hotkey"> & {
-  hotkeys: (string | undefined)[]
-}
-type GetBittensorStakesByHotKeys = Omit<GetBittensorStakeByHotKey, "address" | "hotkey"> & {
-  addresses: (string | undefined)[]
-  hotkeys: (string[] | undefined)[]
-}
-
-export const useGetBittensorStakeByHotKey = ({
-  address,
-  hotkey,
-  isEnabled = true,
-}: GetBittensorStakeByHotKey) => {
-  const { data: sapi } = useScaleApi("bittensor")
-  return useQuery({
-    queryKey: ["useGetBittensorStakeByHotKey", sapi?.id, address, hotkey],
-    queryFn: async () => {
-      return sapi?.getStorage<bigint>("SubtensorModule", "Stake", [hotkey, address])
-    },
-    enabled: isEnabled && !!sapi && !!address && !!hotkey,
-  })
-}
-
-export const useGetBittensorStakeByHotKeys = ({
-  address,
-  hotkeys,
-  isEnabled = true,
-}: GetBittensorStakeByHotKeys) => {
-  const { data: sapi } = useScaleApi("bittensor")
-  return useQueries({
-    queries: hotkeys.map((hotkey) => ({
-      queryKey: ["useGetBittensorStakeByHotKey", sapi?.id, address, hotkey],
-      queryFn: () => sapi?.getStorage<bigint>("SubtensorModule", "Stake", [hotkey, address]),
-      enabled: isEnabled && !!sapi && !!address && !!hotkey,
-    })),
-    combine: (results) => {
-      return {
-        data: results.map((result) => result.data),
-        isPending: results.some((result) => result.isPending),
-        isLoading: results.some((result) => result.isLoading),
-        error: results.find((result) => result.isError),
-      }
-    },
-  })
-}
+export const useGetBittensorStakeByHotKey = ({ address, hotkey }: GetBittensorStakeByHotKey) => {
+  const balance = useBalance(hotkey ? address : null, "bittensor-substrate-native")
 
-export const useGetBittensorStakesByHotKeys = ({
-  addresses,
-  hotkeys,
-  isEnabled = true,
-}: GetBittensorStakesByHotKeys) => {
-  const { data: sapi } = useScaleApi("bittensor")
-  return useQueries({
-    queries: addresses
-      .map((address, index) => {
-        const hotkeysForAddress = hotkeys[index] ?? []
-        return hotkeysForAddress.map((hotkey) => ({
-          queryKey: ["useGetBittensorStakeByHotKey", sapi?.id, address, hotkey],
-          queryFn: () => sapi?.getStorage<bigint>("SubtensorModule", "Stake", [hotkey, address]),
-          enabled: isEnabled && !!sapi && !!address && !!hotkey,
-        }))
-      })
-      .flat(),
-    combine: (results) => {
-      return {
-        data: results.map((result) => result.data),
-        isPending: results.some((result) => result.isPending),
-        isLoading: results.some((result) => result.isLoading),
-        error: results.find((result) => result.isError),
-      }
-    },
-  })
+  return useMemo(() => {
+    if (!balance || !hotkey) return undefined
+    const value = balance?.subtensor.find((b) => (b.meta as { hotkey?: string })?.hotkey === hotkey)
+    return value?.amount.planck
+  }, [balance, hotkey])
 }
diff --git a/apps/extension/src/ui/domains/Staking/shared/useGetUnbondInfo.ts b/apps/extension/src/ui/domains/Staking/shared/useGetUnbondInfo.ts
index 79f6a3fe01..c3d026a0f3 100644
--- a/apps/extension/src/ui/domains/Staking/shared/useGetUnbondInfo.ts
+++ b/apps/extension/src/ui/domains/Staking/shared/useGetUnbondInfo.ts
@@ -34,7 +34,7 @@ export const useGetUnbondInfo = ({ sapi, chainId, address, unstakePoolId }: GetU
     isEnabled: chainId !== "bittensor",
   })
 
-  const { data: bittensorPlanks } = useGetBittensorStakeByHotKey({
+  const bittensorPlanks = useGetBittensorStakeByHotKey({
     address,
     hotkey: unstakePoolId,
   })