SRCmail/stores/account-store.ts

import { create } from 'zustand';
import { persist, createJSONStorage } from 'zustand/middleware';
import { encryptedStorage } from '@/stores/encrypted-storage';
import { generateAccountId, generateAvatarColor, getMaxAccounts } from '@/lib/account-utils';

export interface AccountEntry {
  /** Unique key: `${username}@${serverHostname}` */
  id: string;
  /** Display label (defaults to email, user-editable) */
  label: string;
  /** Full server URL */
  serverUrl: string;
  /** Username / email used to authenticate */
  username: string;
  /** Authentication mode */
  authMode: 'basic' | 'oauth';
  /** Cookie slot index for session/token cookies (0 ≤ slot < MAX_ACCOUNT_SLOTS) */
  cookieSlot: number;
  /** Whether "Remember Me" was checked (basic auth only) */
  rememberMe: boolean;
  /**
   * Server-confirmed account identifiers captured at login: the account-id form
   * ({@link generateAccountId}) of the JMAP Session.username and the primary
   * sending-identity email. The account-switch guard matches a reconnected
   * session against these so a short login username (e.g. `linus`) that the
   * server canonicalizes to a full address (`linus@example.com`) is still
   * recognized as the same account. `undefined` = never captured (legacy entry).
   */
  serverIdentifiers?: string[];
  /** Cached display info */
  displayName: string;
  email: string;
  avatarColor: string;
  /** Timestamp of last successful login */
  lastLoginAt: number;
  /** Whether this account is currently connected */
  isConnected: boolean;
  /** Whether this account had a connection error */
  hasError: boolean;
  errorMessage?: string;
  /** Whether this is the default account (loaded on app start) */
  isDefault: boolean;
}

interface AccountState {
  accounts: AccountEntry[];
  activeAccountId: string | null;
  defaultAccountId: string | null;

  addAccount: (entry: Omit<AccountEntry, 'id' | 'cookieSlot' | 'avatarColor'>) => string;
  removeAccount: (accountId: string) => void;
  setActiveAccount: (accountId: string) => void;
  setDefaultAccount: (accountId: string) => void;
  getDefaultAccount: () => AccountEntry | null;
  updateAccount: (accountId: string, updates: Partial<AccountEntry>) => void;
  reorderAccounts: (orderedIds: string[]) => void;
  getActiveAccount: () => AccountEntry | null;
  getAccountById: (accountId: string) => AccountEntry | undefined;
  getNextCookieSlot: () => number;
  hasAccount: (username: string, serverUrl: string) => boolean;
}

export const useAccountStore = create<AccountState>()(
  persist(
    (set, get) => ({
      accounts: [],
      activeAccountId: null,
      defaultAccountId: null,

      addAccount: (entry) => {
        const state = get();

        const id = generateAccountId(entry.username, entry.serverUrl);
        if (state.accounts.some((a) => a.id === id)) {
          // Already exists - update mutable fields and return existing id
          set((s) => ({
            accounts: s.accounts.map((a) =>
              a.id === id
                ? {
                    ...a,
                    rememberMe: entry.rememberMe,
                    isConnected: entry.isConnected,
                    hasError: entry.hasError,
                    errorMessage: undefined,
                    lastLoginAt: entry.lastLoginAt,
                    authMode: entry.authMode,
                  }
                : a
            ),
          }));
          return id;
        }

        const max = getMaxAccounts();
        if (state.accounts.length >= max) {
          throw new Error(`Maximum of ${max} accounts reached`);
        }

        const cookieSlot = state.getNextCookieSlot();
        const avatarColor = generateAvatarColor(entry.email || entry.username);
        const isDefault = state.accounts.length === 0; // first account is default

        const account: AccountEntry = {
          ...entry,
          id,
          cookieSlot,
          avatarColor,
          isDefault,
        };

        set((s) => ({
          accounts: [...s.accounts, account],
          // If there is no active account, activate this one
          activeAccountId: s.activeAccountId ?? id,
          defaultAccountId: isDefault ? id : s.defaultAccountId,
        }));

        return id;
      },

      removeAccount: (accountId) => {
        set((s) => {
          const remaining = s.accounts.filter((a) => a.id !== accountId);
          const wasDefault = s.defaultAccountId === accountId;
          const wasActive = s.activeAccountId === accountId;

          let newDefault = s.defaultAccountId;
          if (wasDefault) {
            newDefault = remaining[0]?.id ?? null;
            // Mark new default
            if (newDefault) {
              const idx = remaining.findIndex((a) => a.id === newDefault);
              if (idx >= 0) {
                remaining[idx] = { ...remaining[idx], isDefault: true };
              }
            }
          }

          return {
            accounts: remaining,
            activeAccountId: wasActive ? (remaining[0]?.id ?? null) : s.activeAccountId,
            defaultAccountId: newDefault,
          };
        });
      },

      setActiveAccount: (accountId) => {
        const account = get().accounts.find((a) => a.id === accountId);
        if (!account) return;
        set({ activeAccountId: accountId });
      },

      setDefaultAccount: (accountId) => {
        const account = get().accounts.find((a) => a.id === accountId);
        if (!account) return;
        set((s) => ({
          defaultAccountId: accountId,
          accounts: s.accounts.map((a) => ({
            ...a,
            isDefault: a.id === accountId,
          })),
        }));
      },

      getDefaultAccount: () => {
        const state = get();
        if (state.defaultAccountId) {
          const account = state.accounts.find((a) => a.id === state.defaultAccountId);
          if (account) return account;
        }
        return state.accounts[0] ?? null;
      },

      updateAccount: (accountId, updates) => {
        set((s) => ({
          accounts: s.accounts.map((a) =>
            a.id === accountId ? { ...a, ...updates } : a
          ),
        }));
      },

      reorderAccounts: (orderedIds) => {
        set((s) => {
          const byId = new Map(s.accounts.map((a) => [a.id, a]));
          const reordered: AccountEntry[] = [];
          for (const id of orderedIds) {
            const a = byId.get(id);
            if (a) {
              reordered.push(a);
              byId.delete(id);
            }
          }
          // Append any accounts that weren't in the ordered list (defensive)
          for (const a of byId.values()) reordered.push(a);
          return { accounts: reordered };
        });
      },

      getActiveAccount: () => {
        const state = get();
        return state.accounts.find((a) => a.id === state.activeAccountId) ?? null;
      },

      getAccountById: (accountId) => {
        return get().accounts.find((a) => a.id === accountId);
      },

      getNextCookieSlot: () => {
        const used = new Set(get().accounts.map((a) => a.cookieSlot));
        let i = 0;
        while (used.has(i)) i++;
        return i;
      },

      hasAccount: (username, serverUrl) => {
        const id = generateAccountId(username, serverUrl);
        return get().accounts.some((a) => a.id === id);
      },
    }),
    {
      name: 'account-registry',
      storage: createJSONStorage(() => encryptedStorage),
      partialize: (state) => ({
        accounts: state.accounts,
        activeAccountId: state.activeAccountId,
        defaultAccountId: state.defaultAccountId,
      }),
    }
  )
);