Amba-Agent-App/lib/stores/userWalletStore.ts

/**
 * User Wallet Store - Platform-aware wallet management
 * Uses Firebase abstraction layer for cross-platform support
 */
import { create } from 'zustand';
import { doc } from '../firebase';
import type { UserWallet } from '../services/walletService';

export interface WalletCacheEntry {
    wallet: UserWallet | null;
    loading: boolean;
    error: string | null;
    lastFetched?: number;
}

interface UserWalletStoreState {
    wallets: Record<string, WalletCacheEntry>;
    subscriptions: Record<string, () => void>;
    ensureSubscription: (uid: string) => void;
    refreshWallet: (uid: string) => Promise<UserWallet | null>;
    invalidateWallet: (uid: string) => Promise<UserWallet | null>;
    removeWallet: (uid: string) => void;
    clearAll: () => void;
    setWalletState: (uid: string, updates: Partial<WalletCacheEntry>) => void;
}

const buildEntry = (
    existing: WalletCacheEntry | undefined,
    overrides: Partial<WalletCacheEntry>
): WalletCacheEntry => ({
    wallet: existing?.wallet ?? null,
    loading: existing?.loading ?? false,
    error: existing?.error ?? null,
    lastFetched: existing?.lastFetched,
    ...overrides,
});

const snapshotExists = (docSnap: any): boolean => {
    const existsValue = docSnap?.exists;
    if (typeof existsValue === 'function') {
        try {
            return !!existsValue.call(docSnap);
        } catch {
            return false;
        }
    }
    return !!existsValue;
};

export const useUserWalletStore = create<UserWalletStoreState>((set, get) => ({
    wallets: {},
    subscriptions: {},

    ensureSubscription: (uid: string) => {
        if (!uid) {
            return;
        }

        // If already subscribed, don't create another subscription
        if (get().subscriptions[uid]) {
            return;
        }

        const existing = get().wallets[uid];
        if (existing?.loading) {
            return;
        }

        // If we have cached data, set up a real-time listener
        if (existing?.lastFetched) {
            const walletDocRef = doc('wallets', uid);
            const unsubscribe = walletDocRef.onSnapshot(
                (docSnap: any) => {
                    const wallet = snapshotExists(docSnap) ? (docSnap.data() as UserWallet) : null;
                    set((state) => ({
                        wallets: {
                            ...state.wallets,
                            [uid]: buildEntry(state.wallets[uid], {
                                wallet,
                                loading: false,
                                error: wallet ? null : 'Wallet not found',
                                lastFetched: Date.now(),
                            }),
                        },
                    }));
                },
                (error: any) => {
                    set((state) => ({
                        wallets: {
                            ...state.wallets,
                            [uid]: buildEntry(state.wallets[uid], {
                                loading: false,
                                error: error.message || 'Failed to fetch wallet',
                            }),
                        },
                    }));
                }
            );

            set((state) => ({
                subscriptions: {
                    ...state.subscriptions,
                    [uid]: unsubscribe,
                },
            }));
            return;
        }

        // Initial fetch
        void get().refreshWallet(uid).catch((error) => {
            set((state) => ({
                wallets: {
                    ...state.wallets,
                    [uid]: buildEntry(state.wallets[uid], {
                        loading: false,
                        error: error instanceof Error ? error.message : 'Failed to fetch wallet',
                    }),
                },
            }));
        });
    },

    refreshWallet: async (uid: string) => {
        if (!uid) {
            return null;
        }

        set((state) => ({
            wallets: {
                ...state.wallets,
                [uid]: buildEntry(state.wallets[uid], {
                    loading: true,
                    error: null,
                }),
            },
        }));

        try {
            const walletDocRef = doc('wallets', uid);
            const docSnap = await walletDocRef.get();
            const wallet = snapshotExists(docSnap) ? (docSnap.data() as UserWallet) : null;

            set((state) => ({
                wallets: {
                    ...state.wallets,
                    [uid]: buildEntry(state.wallets[uid], {
                        wallet,
                        loading: false,
                        error: wallet ? null : 'Wallet not found',
                        lastFetched: Date.now(),
                    }),
                },
            }));

            // Set up real-time listener after initial fetch if not already subscribed
            if (!get().subscriptions[uid] && wallet) {
                const unsubscribe = walletDocRef.onSnapshot(
                    (docSnap: any) => {
                        const updatedWallet = snapshotExists(docSnap) ? (docSnap.data() as UserWallet) : null;
                        set((state) => ({
                            wallets: {
                                ...state.wallets,
                                [uid]: buildEntry(state.wallets[uid], {
                                    wallet: updatedWallet,
                                    loading: false,
                                    error: updatedWallet ? null : 'Wallet not found',
                                    lastFetched: Date.now(),
                                }),
                            },
                        }));
                    },
                    (error: any) => {
                        set((state) => ({
                            wallets: {
                                ...state.wallets,
                                [uid]: buildEntry(state.wallets[uid], {
                                    loading: false,
                                    error: error.message || 'Failed to fetch wallet',
                                }),
                            },
                        }));
                    }
                );

                set((state) => ({
                    subscriptions: {
                        ...state.subscriptions,
                        [uid]: unsubscribe,
                    },
                }));
            }

            return wallet;
        } catch (error) {
            set((state) => ({
                wallets: {
                    ...state.wallets,
                    [uid]: buildEntry(state.wallets[uid], {
                        loading: false,
                        error: error instanceof Error ? error.message : 'Failed to fetch wallet',
                    }),
                },
            }));
            throw error;
        }
    },

    invalidateWallet: async (uid: string) => {
        if (!uid) {
            return null;
        }

        set((state) => ({
            wallets: {
                ...state.wallets,
                [uid]: buildEntry(state.wallets[uid], {
                    loading: true,
                    error: null,
                    lastFetched: undefined,
                }),
            },
        }));

        return get().refreshWallet(uid);
    },

    removeWallet: (uid: string) => {
        // Unsubscribe from listener if exists
        const unsubscribe = get().subscriptions[uid];
        if (unsubscribe) {
            unsubscribe();
        }

        set((state) => {
            const { [uid]: _removedWallet, ...restWallets } = state.wallets;
            const { [uid]: _removedSub, ...restSubs } = state.subscriptions;
            return {
                wallets: restWallets,
                subscriptions: restSubs,
            };
        });
    },

    clearAll: () => {
        // Unsubscribe from all listeners
        const subscriptions = get().subscriptions;
        Object.values(subscriptions).forEach((unsubscribe) => {
            unsubscribe();
        });
        set({ wallets: {}, subscriptions: {} });
    },

    setWalletState: (uid: string, updates: Partial<WalletCacheEntry>) => {
        if (!uid) {
            return;
        }
        set((state) => ({
            wallets: {
                ...state.wallets,
                [uid]: buildEntry(state.wallets[uid], updates),
            },
        }));
    },
}));

export const getWalletFromCache = (uid: string): UserWallet | null => {
    if (!uid) {
        return null;
    }
    return useUserWalletStore.getState().wallets[uid]?.wallet ?? null;
};