cinny/src/app/state/sessions.ts

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

// OIDC/next-gen-auth (MSC3861) metadata needed to refresh + revoke tokens. Kept
// as a generic claims object so this core module stays decoupled from
// oidc-client-ts (the refresher casts idTokenClaims to IdTokenClaims).
export type OidcSessionMeta = {
  issuer: string;
  clientId: string;
  redirectUri: string;
  idTokenClaims?: Record<string, unknown>;
};

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

// Legacy per-field localStorage keys. Kept for dual-write (see below) so a
// rollback to an older build that only understands these keys still works.
const LEGACY_KEYS = {
  accessToken: 'cinny_access_token',
  deviceId: 'cinny_device_id',
  userId: 'cinny_user_id',
  baseUrl: 'cinny_hs_base_url',
} as const;

// OIDC-only localStorage keys (absent for password/legacy-SSO sessions).
const OIDC_KEYS = {
  refreshToken: 'cinny_refresh_token',
  expiresAt: 'cinny_expires_at',
  issuer: 'cinny_oidc_issuer',
  clientId: 'cinny_oidc_client_id',
  redirectUri: 'cinny_oidc_redirect_uri',
  idTokenClaims: 'cinny_oidc_id_token_claims',
} as const;

// Single-key atomic session blob. The whole session is serialised and written
// in ONE `setItem`, so a reader can never observe a torn/partial session the
// way the multi-key legacy layout could. Bumping the schema means bumping the
// `_v1` suffix.
const SESSION_BLOB_KEY = 'cinny_session_v1';

// The exact shape stored inside SESSION_BLOB_KEY. Note it stores an ABSOLUTE
// `expiresAt` (ms since epoch) rather than a relative lifetime — identical to
// the legacy `cinny_expires_at` semantics — so reads stay drift-free.
type PersistedSession = {
  accessToken: string;
  deviceId: string;
  userId: string;
  baseUrl: string;
  refreshToken?: string;
  expiresAt?: number;
  oidc?: OidcSessionMeta;
};

// Build the persisted shape from the public setFallbackSession arguments. This
// is the single source of truth written to BOTH the blob and the legacy keys.
const buildPersisted = (
  accessToken: string,
  deviceId: string,
  userId: string,
  baseUrl: string,
  extra?: FallbackSessionExtra,
): PersistedSession => {
  const persisted: PersistedSession = { accessToken, deviceId, userId, baseUrl };
  if (extra?.refreshToken) persisted.refreshToken = extra.refreshToken;
  // Store ABSOLUTE expiry to avoid drift across reloads.
  if (typeof extra?.expiresInMs === 'number') persisted.expiresAt = Date.now() + extra.expiresInMs;
  if (extra?.oidc) persisted.oidc = extra.oidc;
  return persisted;
};

// Convert a persisted shape into the public Session returned to callers. Keeps
// behaviour identical to the original getFallbackSession assembly: derives the
// REMAINING lifetime from the absolute expiry, and only surfaces `oidc` when the
// three required OIDC fields are present.
const sessionFromPersisted = (p: PersistedSession): Session => {
  const session: Session = {
    baseUrl: p.baseUrl,
    userId: p.userId,
    deviceId: p.deviceId,
    accessToken: p.accessToken,
    fallbackSdkStores: true,
  };

  if (p.refreshToken) session.refreshToken = p.refreshToken;

  if (typeof p.expiresAt === 'number' && Number.isFinite(p.expiresAt)) {
    // Expose the REMAINING lifetime (clamped at 0); the SDK refreshes on 401.
    session.expiresInMs = Math.max(0, p.expiresAt - Date.now());
  }

  if (p.oidc && p.oidc.issuer && p.oidc.clientId && p.oidc.redirectUri) {
    session.oidc = {
      issuer: p.oidc.issuer,
      clientId: p.oidc.clientId,
      redirectUri: p.oidc.redirectUri,
      idTokenClaims: p.oidc.idTokenClaims,
    };
  }

  return session;
};

// Read the atomic blob. Returns undefined when absent, unparseable, or missing
// any of the four required credential fields — callers then fall back to the
// legacy keys.
const readSessionBlob = (): PersistedSession | undefined => {
  const raw = localStorage.getItem(SESSION_BLOB_KEY);
  if (!raw) return undefined;
  let parsed: unknown;
  try {
    parsed = JSON.parse(raw);
  } catch {
    // Corrupt JSON — treat as absent and let the legacy path take over.
    return undefined;
  }
  if (!parsed || typeof parsed !== 'object') return undefined;
  const p = parsed as Partial<PersistedSession>;
  if (
    typeof p.accessToken !== 'string' ||
    typeof p.deviceId !== 'string' ||
    typeof p.userId !== 'string' ||
    typeof p.baseUrl !== 'string'
  ) {
    // Partial/corrupt blob — fall back to legacy assembly.
    return undefined;
  }
  return p as PersistedSession;
};

// Assemble a session from the legacy per-field keys, or undefined when the four
// required keys are not all present. Used for transparent migration from builds
// that predate the atomic blob.
const readLegacyKeys = (): PersistedSession | undefined => {
  const baseUrl = localStorage.getItem(LEGACY_KEYS.baseUrl);
  const userId = localStorage.getItem(LEGACY_KEYS.userId);
  const deviceId = localStorage.getItem(LEGACY_KEYS.deviceId);
  const accessToken = localStorage.getItem(LEGACY_KEYS.accessToken);

  if (!(baseUrl && userId && deviceId && accessToken)) return undefined;

  const persisted: PersistedSession = { accessToken, deviceId, userId, baseUrl };

  const refreshToken = localStorage.getItem(OIDC_KEYS.refreshToken);
  if (refreshToken) persisted.refreshToken = refreshToken;

  const expiresAtRaw = localStorage.getItem(OIDC_KEYS.expiresAt);
  if (expiresAtRaw) {
    const expiresAt = Number(expiresAtRaw);
    if (Number.isFinite(expiresAt)) persisted.expiresAt = expiresAt;
  }

  const issuer = localStorage.getItem(OIDC_KEYS.issuer);
  const clientId = localStorage.getItem(OIDC_KEYS.clientId);
  const redirectUri = localStorage.getItem(OIDC_KEYS.redirectUri);
  if (issuer && clientId && redirectUri) {
    let idTokenClaims: Record<string, unknown> | undefined;
    const claimsRaw = localStorage.getItem(OIDC_KEYS.idTokenClaims);
    if (claimsRaw) {
      try {
        idTokenClaims = JSON.parse(claimsRaw);
      } catch {
        /* corrupt claims — ignore, the refresher will re-validate on use */
      }
    }
    persisted.oidc = { issuer, clientId, redirectUri, idTokenClaims };
  }

  return persisted;
};

// Write the legacy per-field keys (dual-write half). Mirrors the original
// setFallbackSession body so a rollback to an older build keeps working.
const writeLegacyKeys = (p: PersistedSession): void => {
  localStorage.setItem(LEGACY_KEYS.accessToken, p.accessToken);
  localStorage.setItem(LEGACY_KEYS.deviceId, p.deviceId);
  localStorage.setItem(LEGACY_KEYS.userId, p.userId);
  localStorage.setItem(LEGACY_KEYS.baseUrl, p.baseUrl);

  // OIDC fields — written only when present; otherwise cleared so a password
  // session never carries stale OIDC state.
  if (p.refreshToken) localStorage.setItem(OIDC_KEYS.refreshToken, p.refreshToken);
  else localStorage.removeItem(OIDC_KEYS.refreshToken);

  if (typeof p.expiresAt === 'number')
    localStorage.setItem(OIDC_KEYS.expiresAt, String(p.expiresAt));
  else localStorage.removeItem(OIDC_KEYS.expiresAt);

  if (p.oidc) {
    localStorage.setItem(OIDC_KEYS.issuer, p.oidc.issuer);
    localStorage.setItem(OIDC_KEYS.clientId, p.oidc.clientId);
    localStorage.setItem(OIDC_KEYS.redirectUri, p.oidc.redirectUri);
    if (p.oidc.idTokenClaims) {
      localStorage.setItem(OIDC_KEYS.idTokenClaims, JSON.stringify(p.oidc.idTokenClaims));
    } else localStorage.removeItem(OIDC_KEYS.idTokenClaims);
  } else {
    localStorage.removeItem(OIDC_KEYS.issuer);
    localStorage.removeItem(OIDC_KEYS.clientId);
    localStorage.removeItem(OIDC_KEYS.redirectUri);
    localStorage.removeItem(OIDC_KEYS.idTokenClaims);
  }
};

export type FallbackSessionExtra = {
  refreshToken?: string;
  expiresInMs?: number;
  oidc?: OidcSessionMeta;
};

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;

// Persist the session. Writes the atomic blob FIRST (so the consistent,
// never-torn copy is established before the multi-key legacy write), then
// dual-writes the legacy keys for rollback safety. Signature is unchanged —
// callers (login/register/OIDC callback/token refresher) are untouched.
export function setFallbackSession(
  accessToken: string,
  deviceId: string,
  userId: string,
  baseUrl: string,
  extra?: FallbackSessionExtra,
) {
  const persisted = buildPersisted(accessToken, deviceId, userId, baseUrl, extra);
  // ONE setItem — the blob can never be observed half-written.
  localStorage.setItem(SESSION_BLOB_KEY, JSON.stringify(persisted));
  // Dual-write the legacy keys (removal of this half is a future release).
  writeLegacyKeys(persisted);
}

// Clear BOTH the atomic blob and every legacy key so no reader (blob-preferring
// or legacy-fallback) can resurrect a logged-out session.
export const removeFallbackSession = () => {
  localStorage.removeItem(SESSION_BLOB_KEY);
  Object.values(LEGACY_KEYS).forEach((key) => localStorage.removeItem(key));
  Object.values(OIDC_KEYS).forEach((key) => localStorage.removeItem(key));
};

// Read the session, preferring the atomic blob. If the blob is absent or
// corrupt/partial we transparently assemble from the legacy keys (migration);
// the next setFallbackSession then persists the blob. When both exist the blob
// wins by construction.
export const getFallbackSession = (): Session | undefined => {
  const persisted = readSessionBlob() ?? readLegacyKeys();
  if (!persisted) return undefined;
  return sessionFromPersisted(persisted);
};
/**
 * End of migration code for old session
 */

// Session keys whose cross-tab change indicates a login/logout/token-rotation
// in another tab. localStorage.clear() dispatches a storage event with a null
// key, which we also treat as a session change.
const SESSION_STORAGE_KEYS = new Set<string>([
  SESSION_BLOB_KEY,
  ...Object.values(LEGACY_KEYS),
  ...Object.values(OIDC_KEYS),
]);

/**
 * Subscribe to session changes made in OTHER tabs/windows. The browser only
 * dispatches `storage` events to tabs that did NOT perform the write, so this
 * is inherently guarded against reacting to our own same-tab writes — no
 * echo-suppression needed. The callback receives the freshly-read session, or
 * `null` when the session was removed (logout in another tab, or a full
 * localStorage.clear()). Returns an unsubscribe function.
 */
export const subscribeSessionChanges = (
  callback: (session: Session | null) => void,
): (() => void) => {
  const handleStorage = (evt: StorageEvent) => {
    // A null key means localStorage.clear(); otherwise only react to our keys.
    if (evt.key !== null && !SESSION_STORAGE_KEYS.has(evt.key)) return;
    callback(getFallbackSession() ?? null);
  };
  window.addEventListener('storage', handleStorage);
  return () => {
    window.removeEventListener('storage', handleStorage);
  };
};

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

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

// 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) {
//       removeFallbackSession();
//       sessions.push(fallbackSession);
//       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
//       );
//       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
//       );
//       set(baseSessionsAtom, sessions);
//     }
//   }
// );