From 59365bad4e5b39f6ceebc5eb4a566fed9ca0bc3f Mon Sep 17 00:00:00 2001
From: lukema95 <867273263@qq.com>
Date: Thu, 22 Aug 2024 21:13:54 +0800
Subject: [PATCH] use wagmi context for Balances and UniversalKitProvider

---
 src/components/Balances/index.tsx      |  7 +++-
 src/providers/UniversalKitProvider.tsx | 55 +++++++++++++-------------
 2 files changed, 33 insertions(+), 29 deletions(-)

diff --git a/src/components/Balances/index.tsx b/src/components/Balances/index.tsx
index c6e2329..5e56df6 100644
--- a/src/components/Balances/index.tsx
+++ b/src/components/Balances/index.tsx
@@ -10,7 +10,7 @@ import { roundNumber } from "@/lib/utils";
 import { useBitcoinWallet } from "@/index";
 import { useAccount, useWalletClient } from "wagmi";
 import { useEthersSigner } from "@/hooks/useEthersSigner";
-import { useZetaChainClient } from "@/providers/UniversalKitProvider";
+import { useZetaChainClient, wagmiContextValue } from "@/providers/UniversalKitProvider";
 
 interface Token {
   id: string;
@@ -22,15 +22,18 @@ interface Token {
 interface BalancesProps {
   config?: any;
   balances?: any;
+  wagmiContextValue?: wagmiContextValue;
   onBalanceClick?: (balance: Token) => void;
 }
 
 export const Balances = ({
   config,
   balances: balancesProp,
+  wagmiContextValue,
   onBalanceClick = () => {},
 }: BalancesProps) => {
-  const { address, status } = useAccount();
+  const useWagmiAccount = wagmiContextValue?.useAccount || useAccount;
+  const { address, status } = useWagmiAccount();
 
   const { address: bitcoin } = useBitcoinWallet();
   const client = useZetaChainClient();
diff --git a/src/providers/UniversalKitProvider.tsx b/src/providers/UniversalKitProvider.tsx
index 240ef2f..e390577 100644
--- a/src/providers/UniversalKitProvider.tsx
+++ b/src/providers/UniversalKitProvider.tsx
@@ -16,17 +16,32 @@ const ZetaChainClientContext = createContext<any>(null);
 
 export const useZetaChainClient = () => useContext(ZetaChainClientContext);
 
+export interface wagmiContextValue {
+  useAccount: typeof useAccount;
+  useChainId: typeof useChainId;
+  useWalletClient: typeof useWalletClient;
+}
+
+interface ZetaChainClientProviderProps {
+  children: ReactNode;
+  zetaChainConfig?: any;
+  wagmiContextValue?: wagmiContextValue;
+}
+
 const ZetaChainClientProvider = ({
   children,
   zetaChainConfig,
-}: {
-  children: ReactNode;
-  zetaChainConfig?: any;
-}) => {
-  const { status } = useAccount();
-  const chainId = useChainId();
-  const { data: walletClient } = useWalletClient({ chainId });
+  wagmiContextValue
+}: ZetaChainClientProviderProps ) => {  
+  const useWagmiAccount = wagmiContextValue?.useAccount || useAccount;
+  const useWagmiChainId = wagmiContextValue?.useChainId || useChainId;
+  const useWagmiWalletClient = wagmiContextValue?.useWalletClient || useWalletClient;
+
+  const { status } = useWagmiAccount();
+  const chainId = useWagmiChainId();
+  const { data: walletClient } = useWagmiWalletClient({ chainId });
   const signer = useEthersSigner({ walletClient });
+
   const [zetaChainClient, setZetaChainClient] = useState<any>(null);
 
   useEffect(() => {
@@ -57,26 +72,12 @@ const ZetaChainClientProvider = ({
   );
 };
 
-export const UniversalKitProvider = ({
-  children,
-  config,
-  client,
-  zetaChainConfig,
-}: {
-  children: ReactNode;
-  config?: any;
-  client: any;
-  zetaChainConfig?: any;
-}) => {
+export const UniversalKitProvider = ({ children, zetaChainConfig, wagmiContextValue }: ZetaChainClientProviderProps) => {
   return (
-    <WagmiProvider config={config || { autoConnect: true }}>
-      <QueryClientProvider client={client}>
-        <BitcoinWalletProvider>
-          <ZetaChainClientProvider zetaChainConfig={zetaChainConfig}>
-            {children}
-          </ZetaChainClientProvider>
-        </BitcoinWalletProvider>
-      </QueryClientProvider>
-    </WagmiProvider>
+    <BitcoinWalletProvider>
+      <ZetaChainClientProvider zetaChainConfig={zetaChainConfig} wagmiContextValue={wagmiContextValue}>
+        {children}
+      </ZetaChainClientProvider>
+    </BitcoinWalletProvider>
   );
 };