cinny/src/app/state/sessions.ts

import { atom } from 'jotai';
import {
  atomWithLocalStorage,
  getLocalStorageItem,
  setLocalStorageItem,
} from './utils/atomWithLocalStorage';

export type Session = {
  baseUrl: string;
  userId: string;
  deviceId: string;
  accessToken: string;
  expiresInMs?: number;
  refreshToken?: string;
  fallbackSdkStores?: boolean;
};

export type Sessions = Session[];
export type SessionStoreName = {
  sync: string;
  crypto: string;
};

/**
 * Migration code for old session
 */
const FALLBACK_STORE_NAME: SessionStoreName = {
  sync: 'web-sync-store',
  crypto: 'crypto-store',
} as const;

export function setFallbackSession(
  accessToken: string,
  deviceId: string,
  userId: string,
  baseUrl: string
) {
  localStorage.setItem('cinny_access_token', accessToken);
  localStorage.setItem('cinny_device_id', deviceId);
  localStorage.setItem('cinny_user_id', userId);
  localStorage.setItem('cinny_hs_base_url', baseUrl);
}
export const removeFallbackSession = () => {
  localStorage.removeItem('cinny_hs_base_url');
  localStorage.removeItem('cinny_user_id');
  localStorage.removeItem('cinny_device_id');
  localStorage.removeItem('cinny_access_token');
};
export const getFallbackSession = (): Session | undefined => {
  const baseUrl = localStorage.getItem('cinny_hs_base_url');
  const userId = localStorage.getItem('cinny_user_id');
  const deviceId = localStorage.getItem('cinny_device_id');
  const accessToken = localStorage.getItem('cinny_access_token');

  if (baseUrl && userId && deviceId && accessToken) {
    const session: Session = {
      baseUrl,
      userId,
      deviceId,
      accessToken,
      fallbackSdkStores: true,
    };

    return session;
  }

  return undefined;
};
/**
 * End of migration code for old session
 */

export const getSessionStoreName = (session: Session): SessionStoreName => {
  if (session.fallbackSdkStores) {
    return FALLBACK_STORE_NAME;
  }

  return {
    sync: `sync@${session.userId}@${session.deviceId}`,
    crypto: `crypto@${session.userId}@${session.deviceId}`,
  };
};

export const MATRIX_SESSIONS_KEY = 'matrixSessions';
const baseSessionsAtom = atomWithLocalStorage<Sessions>(
  MATRIX_SESSIONS_KEY,
  (key) => {
    const defaultSessions: Sessions = [];
    const sessions = getLocalStorageItem(key, defaultSessions);

    // Before multi account support session was stored
    // as multiple item in local storage.
    // So we need these migration code.
    const fallbackSession = getFallbackSession();
    if (fallbackSession && sessions.length === 0) {
      // Only use fallback stores if this is the ONLY session
      // This prevents multiple accounts from trying to use the same crypto-store
      removeFallbackSession();
      sessions.push(fallbackSession);
      setLocalStorageItem(key, sessions);
    } else if (fallbackSession && sessions.length > 0) {
      // If there are already sessions, add the fallback but don't use fallback stores
      removeFallbackSession();
      const fallbackWithoutStoreFlag = { ...fallbackSession, fallbackSdkStores: false };
      // Check if this session already exists (by userId + deviceId)
      const exists = sessions.some(
        s => s.userId === fallbackSession.userId && s.deviceId === fallbackSession.deviceId
      );
      if (!exists) {
        sessions.push(fallbackWithoutStoreFlag);
        setLocalStorageItem(key, sessions);
      }
    }
    return sessions;
  },
  (key, value) => {
    setLocalStorageItem(key, value);
  }
);

export type SessionsAction =
  | {
      type: 'PUT';
      session: Session;
    }
  | {
      type: 'DELETE';
      session: Session;
    };

export const sessionsAtom = atom<Sessions, [SessionsAction], undefined>(
  (get) => get(baseSessionsAtom),
  (get, set, action) => {
    if (action.type === 'PUT') {
      const sessions = [...get(baseSessionsAtom)];
      const sessionIndex = sessions.findIndex(
        (session) => session.userId === action.session.userId && session.deviceId === action.session.deviceId
      );
      if (sessionIndex === -1) {
        sessions.push(action.session);
      } else {
        sessions.splice(sessionIndex, 1, action.session);
      }
      set(baseSessionsAtom, sessions);
      return;
    }
    if (action.type === 'DELETE') {
      const sessions = get(baseSessionsAtom).filter(
        (session) => !(session.userId === action.session.userId && session.deviceId === action.session.deviceId)
      );
      set(baseSessionsAtom, sessions);
    }
  }
);

// Current selected session
export const CURRENT_SESSION_KEY = 'currentSessionUserId';
export const currentSessionAtom = atomWithLocalStorage<string | undefined>(
  CURRENT_SESSION_KEY,
  (key) => getLocalStorageItem(key, undefined),
  (key, value) => setLocalStorageItem(key, value)
);

// Helper to get current session from sessions list
export const getCurrentSession = (): Session | undefined => {
  const currentUserId = getLocalStorageItem<string | undefined>(CURRENT_SESSION_KEY, undefined);
  if (!currentUserId) {
    // Fallback to old session for migration
    return getFallbackSession();
  }
  
  const sessions = getLocalStorageItem<Sessions>(MATRIX_SESSIONS_KEY, []);
  return sessions.find(s => s.userId === currentUserId);
};

// Helper to get a session by userId
export const getSessionByUserId = (userId: string): Session | undefined => {
  const sessions = getLocalStorageItem<Sessions>(MATRIX_SESSIONS_KEY, []);
  return sessions.find(s => s.userId === userId);
};