// Copyright (c) Mysten Labs, Inc.
// Modifications Copyright (c) 2024 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

import type { Wallet, WalletAccount, WalletWithRequiredFeatures } from '@iota/wallet-standard';
import { createStore } from 'zustand';
import type { StateStorage } from 'zustand/middleware';
import { createJSONStorage, persist } from 'zustand/middleware';

import { getWalletUniqueIdentifier } from './utils/walletUtils.js';
import type { ChainType } from '@iota/iota-sdk/client';

type WalletConnectionStatus = 'disconnected' | 'connecting' | 'connected';

export type WalletActions = {
    setAccountSwitched: (selectedAccount: WalletAccount) => void;
    setConnectionStatus: (connectionStatus: WalletConnectionStatus) => void;
    setWalletConnected: (
        wallet: WalletWithRequiredFeatures,
        connectedAccounts: readonly WalletAccount[],
        selectedAccount: WalletAccount | null,
        supportedIntents?: string[],
    ) => void;
    updateWalletAccounts: (accounts: readonly WalletAccount[]) => void;
    setWalletDisconnected: () => void;
    setWalletRegistered: (updatedWallets: WalletWithRequiredFeatures[]) => void;
    setWalletUnregistered: (
        updatedWallets: WalletWithRequiredFeatures[],
        unregisteredWallet: Wallet,
    ) => void;
    setChain: (chain: ChainType | undefined) => void;
};

export type WalletStore = ReturnType<typeof createWalletStore>;

export type StoreState = {
    autoConnectEnabled: boolean;
    wallets: WalletWithRequiredFeatures[];
    accounts: readonly WalletAccount[];
    currentWallet: WalletWithRequiredFeatures | null;
    currentAccount: WalletAccount | null;
    lastConnectedAccountAddress: string | null;
    lastConnectedWalletName: string | null;
    connectionStatus: WalletConnectionStatus;
    supportedIntents: string[];
    chain: ChainType | undefined;
} & WalletActions;

type WalletConfiguration = {
    autoConnectEnabled: boolean;
    wallets: WalletWithRequiredFeatures[];
    storage: StateStorage;
    storageKey: string;
    chain?: ChainType;
};

export function createWalletStore({
    wallets,
    storage,
    storageKey,
    autoConnectEnabled,
    chain,
}: WalletConfiguration) {
    return createStore<StoreState>()(
        persist(
            (set, get) => ({
                autoConnectEnabled,
                wallets,
                accounts: [] as WalletAccount[],
                currentWallet: null,
                currentAccount: null,
                lastConnectedAccountAddress: null,
                lastConnectedWalletName: null,
                connectionStatus: 'disconnected',
                supportedIntents: [],
                chain,
                setConnectionStatus(connectionStatus) {
                    set(() => ({
                        connectionStatus,
                    }));
                },
                setWalletConnected(
                    wallet,
                    connectedAccounts,
                    selectedAccount,
                    supportedIntents = [],
                ) {
                    set(() => ({
                        accounts: connectedAccounts,
                        currentWallet: wallet,
                        currentAccount: selectedAccount,
                        lastConnectedWalletName: getWalletUniqueIdentifier(wallet),
                        lastConnectedAccountAddress: selectedAccount?.address,
                        connectionStatus: 'connected',
                        supportedIntents,
                    }));
                },
                setWalletDisconnected() {
                    set(() => ({
                        accounts: [],
                        currentWallet: null,
                        currentAccount: null,
                        lastConnectedWalletName: null,
                        lastConnectedAccountAddress: null,
                        connectionStatus: 'disconnected',
                        supportedIntents: [],
                    }));
                },
                setAccountSwitched(selectedAccount) {
                    set(() => ({
                        currentAccount: selectedAccount,
                        lastConnectedAccountAddress: selectedAccount.address,
                    }));
                },
                setWalletRegistered(updatedWallets) {
                    set(() => ({ wallets: updatedWallets }));
                },
                setWalletUnregistered(updatedWallets, unregisteredWallet) {
                    if (unregisteredWallet === get().currentWallet) {
                        set(() => ({
                            wallets: updatedWallets,
                            accounts: [],
                            currentWallet: null,
                            currentAccount: null,
                            lastConnectedWalletName: null,
                            lastConnectedAccountAddress: null,
                            connectionStatus: 'disconnected',
                            supportedIntents: [],
                        }));
                    } else {
                        set(() => ({ wallets: updatedWallets }));
                    }
                },
                updateWalletAccounts(accounts) {
                    const currentAccount = get().currentAccount;

                    set(() => ({
                        accounts,
                        currentAccount:
                            (currentAccount &&
                                accounts.find(
                                    ({ address }) => address === currentAccount.address,
                                )) ||
                            accounts[0],
                    }));
                },
                setChain(chain) {
                    set(() => ({ chain }));
                },
            }),
            {
                name: storageKey,
                storage: createJSONStorage(() => storage),
                partialize: ({ lastConnectedWalletName, lastConnectedAccountAddress }) => ({
                    lastConnectedWalletName,
                    lastConnectedAccountAddress,
                }),
            },
        ),
    );
}