From 36f2a26f4655734868937bad80a6f372940fce86 Mon Sep 17 00:00:00 2001
From: Nabarun <nabarun@deepstacksoft.com>
Date: Mon, 27 Jan 2025 16:50:50 +0530
Subject: [PATCH] Add button to disconnect wallet

---
 src/app/page.tsx                | 14 +++++++++++++-
 src/components/WalletHeader.tsx |  9 ++++++++-
 src/services/walletService.ts   | 26 ++++++++++++++++++++++++--
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/src/app/page.tsx b/src/app/page.tsx
index bbcfb5f..5e21e18 100644
--- a/src/app/page.tsx
+++ b/src/app/page.tsx
@@ -5,7 +5,7 @@ import WalletHeader from '../components/WalletHeader'
 import AIServiceCard from '../components/AIServiceCard'
 import { generateWithFlux, FluxGenerationResult, FLUX_MODELS } from '../services/fluxService'
 import { processMTMPayment } from '../services/paymentService'
-import { connectWallet, WalletState } from '../services/walletService'
+import { connectWallet, disconnectWallet, WalletState } from '../services/walletService'
 import { WalletType } from '../services/types'
 
 const Page: React.FC = (): React.ReactElement => {
@@ -29,6 +29,17 @@ const Page: React.FC = (): React.ReactElement => {
     }
   }
 
+  const handleDisconnect = async (): Promise<void> => {
+    if (walletState.type) {
+      await disconnectWallet(walletState.type)
+      setWalletState({
+        connected: false,
+        publicKey: null,
+        type: null
+      })
+    }
+  }
+
   const handleFluxGeneration = (modelId: string, cost: number) => {
     return async (prompt: string): Promise<FluxGenerationResult> => {
       const type = walletState.type;
@@ -69,6 +80,7 @@ const Page: React.FC = (): React.ReactElement => {
           <WalletHeader
             walletState={walletState}
             onConnect={handleConnect}
+            onDisconnect={handleDisconnect}
 	  />
         </div>
 
diff --git a/src/components/WalletHeader.tsx b/src/components/WalletHeader.tsx
index 81954c1..25d4f76 100644
--- a/src/components/WalletHeader.tsx
+++ b/src/components/WalletHeader.tsx
@@ -6,9 +6,10 @@ import { WalletState, SUPPORTED_WALLETS } from '../services/walletService'
 interface WalletHeaderProps {
   walletState: WalletState
   onConnect: (walletType: string) => Promise<void>
+  onDisconnect: () => Promise<void>
 }
 
-const WalletHeader: React.FC<WalletHeaderProps> = ({ walletState, onConnect }) => {
+const WalletHeader: React.FC<WalletHeaderProps> = ({ walletState, onConnect, onDisconnect }) => {
   return (
     <div className="w-full bg-slate-900/50 backdrop-blur-lg rounded-xl shadow-lg border border-orange-800/50 mb-8 p-4">
       {!walletState.connected ? (
@@ -31,6 +32,12 @@ const WalletHeader: React.FC<WalletHeaderProps> = ({ walletState, onConnect }) =
           <span className="px-3 py-1 bg-amber-500/20 rounded-full text-amber-200 text-sm">
             {walletState.publicKey?.slice(0, 22)}...
           </span>
+          <button
+            onClick={onDisconnect}
+            className="bg-red-500 hover:bg-red-600 text-white font-semibold py-2 px-4 rounded-lg transition-all duration-200 shadow-lg hover:shadow-red-500/25 ml-4"
+          >
+            Disconnect
+          </button>
         </div>
       )}
     </div>
diff --git a/src/services/walletService.ts b/src/services/walletService.ts
index e15f298..96b8a67 100644
--- a/src/services/walletService.ts
+++ b/src/services/walletService.ts
@@ -10,6 +10,7 @@ export interface WalletConfig {
   type: WalletType
   name: string
   connect: () => Promise<{ publicKey: string } | null>
+  disconnect: () => Promise<void>
 }
 
 const connectSolflare = async (): Promise<{ publicKey: string } | null> => {
@@ -18,6 +19,12 @@ const connectSolflare = async (): Promise<{ publicKey: string } | null> => {
   return window.solflare.publicKey ? { publicKey: window.solflare.publicKey.toString() } : null
 }
 
+const disconnectSolflare = async (): Promise<void> => {
+  if (window.solflare) {
+    await window.solflare.disconnect()
+  }
+}
+
 const connectPhantom = async (): Promise<{ publicKey: string } | null> => {
   if (!window.phantom?.solana) return null
   try {
@@ -28,16 +35,24 @@ const connectPhantom = async (): Promise<{ publicKey: string } | null> => {
   }
 }
 
+const disconnectPhantom = async (): Promise<void> => {
+  if (window.phantom?.solana) {
+    await window.phantom.solana.disconnect()
+  }
+}
+
 export const SUPPORTED_WALLETS: WalletConfig[] = [
   {
     type: 'solflare',
     name: 'Solflare',
-    connect: connectSolflare
+    connect: connectSolflare,
+    disconnect: disconnectSolflare
   },
   {
     type: 'phantom',
     name: 'Phantom',
-    connect: connectPhantom
+    connect: connectPhantom,
+    disconnect: disconnectPhantom
   }
 ]
 
@@ -62,3 +77,10 @@ export async function connectWallet(type: WalletType): Promise<WalletState> {
     }
   }
 }
+
+export async function disconnectWallet(type: WalletType): Promise<void> {
+  const wallet = SUPPORTED_WALLETS.find(w => w.type === type)
+  if (wallet) {
+    await wallet.disconnect()
+  }
+}
-- 
2.45.2