From b4cd40daa88742118225c9f3771788bad7e5a54f Mon Sep 17 00:00:00 2001
From: Hien Le <leduyhien152@gmail.com>
Date: Mon, 15 Apr 2024 21:09:38 +0700
Subject: [PATCH] fix: auto ask to change network

---
 .../stake/flexible/flexible-stake-content.tsx | 16 ++++++++++----
 hooks/useWalletNetwork.ts                     | 22 +++++++++++++++----
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/components/stake/flexible/flexible-stake-content.tsx b/components/stake/flexible/flexible-stake-content.tsx
index 50d1109..036943d 100644
--- a/components/stake/flexible/flexible-stake-content.tsx
+++ b/components/stake/flexible/flexible-stake-content.tsx
@@ -32,9 +32,10 @@ export const FlexibleStakeContent = (props: Props) => {
     value: 0,
     display: "",
   });
-  const { isConnected, isCorrectNetwork, changeNetwork } = useWalletNetwork({
-    chain,
-  });
+  const { isConnected, isCorrectNetwork, changeNetwork, checkNetwork } =
+    useWalletNetwork({
+      chain,
+    });
 
   const convertedBalance = Number(
     formatUnits(balance, stakingToken?.token_decimal)
@@ -104,7 +105,14 @@ export const FlexibleStakeContent = (props: Props) => {
             <DrawerOverlay />
             <DrawerContent className="!w-full rounded-xl h-[calc(100%-60px)]">
               <div className="flex justify-center items-center h-full">
-                <LoginWidget raw onchain chain="evm" />
+                <LoginWidget
+                  raw
+                  onchain
+                  chain="evm"
+                  onWalletConnectSuccess={async () => {
+                    checkNetwork();
+                  }}
+                />
               </div>
             </DrawerContent>
           </DrawerPortal>
diff --git a/hooks/useWalletNetwork.ts b/hooks/useWalletNetwork.ts
index a3f93e6..83583fa 100644
--- a/hooks/useWalletNetwork.ts
+++ b/hooks/useWalletNetwork.ts
@@ -1,7 +1,7 @@
 import { tryCatch } from "@/utils/tryCatch";
 import { toast } from "@mochi-ui/core";
 import { ChainProvider, useLoginWidget } from "@mochi-web3/login-widget";
-import { useEffect, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
 
 interface Props {
   chain?: { id?: number; name?: string; rpc?: string; currency?: string };
@@ -21,9 +21,8 @@ export const useWalletNetwork = (props: Props) => {
     () => ({})
   );
 
-  const changeNetwork = async () => {
+  const changeNetwork = useCallback(async () => {
     if (!provider) return;
-    console.log("sdgdfg");
     try {
       await provider.request({
         method: "wallet_switchEthereumChain",
@@ -37,7 +36,7 @@ export const useWalletNetwork = (props: Props) => {
       // This error code indicates that the chain has not been added to MetaMask.
       if (switchError.code === 4902) {
         try {
-          await provider?.provider.request({
+          await provider.request({
             method: "wallet_addEthereumChain",
             params: [
               {
@@ -73,8 +72,22 @@ export const useWalletNetwork = (props: Props) => {
         });
       }
     }
+  }, [chain?.currency, chain?.id, chain?.name, chain?.rpc, provider]);
+
+  const [shouldChangeNetwork, setShouldChangeNetwork] = useState(false);
+  const checkNetwork = () => {
+    if (!isCorrectNetwork) {
+      setShouldChangeNetwork(true);
+    }
   };
 
+  useEffect(() => {
+    if (shouldChangeNetwork && !!provider) {
+      changeNetwork();
+      setShouldChangeNetwork(false);
+    }
+  }, [changeNetwork, provider, shouldChangeNetwork]);
+
   useEffect(() => {
     setIsCorrectNetwork(chainId === `0x${chain?.id?.toString(16)}`);
   }, [chain?.id, chainId]);
@@ -83,5 +96,6 @@ export const useWalletNetwork = (props: Props) => {
     isConnected: !!connectedAddress,
     isCorrectNetwork,
     changeNetwork,
+    checkNetwork,
   };
 };