From f93058add2f82c5c2621c730566f1dda3c992baf Mon Sep 17 00:00:00 2001
From: padreug <padreug@aiolabs.dev>
Date: Thu, 18 Sep 2025 09:24:51 +0200
Subject: [PATCH] Implement updateWalletBalance method in PaymentService and
 deprecate AuthService method

- Added updateWalletBalance method in PaymentService to handle wallet balance updates from WebSocket notifications, improving wallet management.
- Deprecated the existing updateWalletBalance method in AuthService, redirecting calls to the new PaymentService method for better consistency and maintainability.
- Updated WalletWebSocketService to utilize PaymentService for balance updates, ensuring accurate wallet state management.

These changes enhance the architecture of wallet balance handling and streamline the update process across services.
---
 src/core/services/PaymentService.ts           | 33 +++++++++++++++++++
 src/modules/base/auth/auth-service.ts         | 31 +++++------------
 .../wallet/services/WalletWebSocketService.ts |  8 ++---
 3 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/src/core/services/PaymentService.ts b/src/core/services/PaymentService.ts
index 277fabd..8f5733b 100644
--- a/src/core/services/PaymentService.ts
+++ b/src/core/services/PaymentService.ts
@@ -117,6 +117,39 @@ export class PaymentService extends BaseService {
     return wallet?.inkey || null
   }
 
+  /**
+   * Update wallet balance from WebSocket notifications
+   * @param balanceMsat - The new balance in millisatoshis
+   * @param walletId - Optional wallet ID (defaults to preferred wallet)
+   */
+  updateWalletBalance(balanceMsat: number, walletId?: string): void {
+    if (!this.authService?.user?.value?.wallets?.length) {
+      this.debug('No wallets available to update')
+      return
+    }
+
+    // Determine which wallet to update
+    let walletToUpdate
+    if (walletId) {
+      walletToUpdate = this.authService.user.value.wallets.find((w: any) => w.id === walletId)
+    } else {
+      // Use preferred wallet if no ID specified
+      walletToUpdate = this.getPreferredWallet()
+    }
+
+    if (walletToUpdate) {
+      const oldBalance = walletToUpdate.balance_msat
+      walletToUpdate.balance_msat = balanceMsat
+
+      // Trigger reactivity in AuthService
+      this.authService.user.value = { ...this.authService.user.value }
+
+      this.debug(`Wallet ${walletToUpdate.id} balance updated: ${oldBalance} -> ${balanceMsat} msat (${Math.round(balanceMsat/1000)} sats)`)
+    } else {
+      this.debug(`Wallet ${walletId} not found for balance update`)
+    }
+  }
+
   /**
    * Generate QR code for Lightning payment request
    */
diff --git a/src/modules/base/auth/auth-service.ts b/src/modules/base/auth/auth-service.ts
index c8e6eb4..ea2fef0 100644
--- a/src/modules/base/auth/auth-service.ts
+++ b/src/modules/base/auth/auth-service.ts
@@ -2,6 +2,7 @@
 import { ref, computed } from 'vue'
 import { BaseService } from '@/core/base/BaseService'
 import { eventBus } from '@/core/event-bus'
+import { tryInjectService, SERVICE_TOKENS } from '@/core/di-container'
 import type { LoginCredentials, RegisterData, User } from '@/lib/api/lnbits'
 
 export class AuthService extends BaseService {
@@ -134,32 +135,16 @@ export class AuthService extends BaseService {
   }
 
   /**
-   * Update wallet balance from WebSocket notifications
-   * @param walletId - The wallet ID to update (if not provided, updates first wallet for backwards compatibility)
-   * @param balanceMsat - The new balance in millisatoshis
+   * @deprecated Use PaymentService.updateWalletBalance() instead
+   * This method will be removed in a future version
    */
   updateWalletBalance(balanceMsat: number, walletId?: string): void {
-    if (!this.user.value?.wallets?.length) return
+    console.warn('AuthService.updateWalletBalance is deprecated. Use PaymentService.updateWalletBalance() instead')
 
-    // Find the wallet to update
-    let walletToUpdate
-    if (walletId) {
-      walletToUpdate = this.user.value.wallets.find(w => w.id === walletId)
-    } else {
-      // Fallback to first wallet for backwards compatibility
-      // TODO: This should eventually be removed once all callers provide walletId
-      walletToUpdate = this.user.value.wallets[0]
-      console.warn('updateWalletBalance called without walletId, using wallets[0] as fallback')
-    }
-
-    if (walletToUpdate) {
-      const oldBalance = walletToUpdate.balance_msat
-      walletToUpdate.balance_msat = balanceMsat
-
-      // Trigger reactivity
-      this.user.value = { ...this.user.value }
-
-      this.debug(`Wallet ${walletToUpdate.id} balance updated: ${oldBalance} -> ${balanceMsat} msat (${Math.round(balanceMsat/1000)} sats) via WebSocket`)
+    // For backwards compatibility, delegate to PaymentService if available
+    const paymentService = tryInjectService(SERVICE_TOKENS.PAYMENT_SERVICE) as any
+    if (paymentService?.updateWalletBalance) {
+      paymentService.updateWalletBalance(balanceMsat, walletId)
     }
   }
   
diff --git a/src/modules/wallet/services/WalletWebSocketService.ts b/src/modules/wallet/services/WalletWebSocketService.ts
index 911606e..9fab90e 100644
--- a/src/modules/wallet/services/WalletWebSocketService.ts
+++ b/src/modules/wallet/services/WalletWebSocketService.ts
@@ -233,7 +233,7 @@ export class WalletWebSocketService extends BaseService {
   }
 
   /**
-   * Update wallet balance in auth service
+   * Update wallet balance via PaymentService
    */
   private updateWalletBalance(balanceSats: number): void {
     console.log('WalletWebSocketService: Updating balance to', balanceSats, 'sats')
@@ -245,9 +245,9 @@ export class WalletWebSocketService extends BaseService {
     const wallet = this.paymentService?.getPreferredWallet?.()
     const walletId = wallet?.id
 
-    // Update balance in auth service (source of truth)
-    if (this.authService?.updateWalletBalance) {
-      this.authService.updateWalletBalance(balanceMsat, walletId)
+    // Update balance via PaymentService (which manages wallet state)
+    if (this.paymentService?.updateWalletBalance) {
+      this.paymentService.updateWalletBalance(balanceMsat, walletId)
     }
   }