NianToB/electron/yinian/storage.ts

import type {
  YinianConfigSnapshot,
  YinianPersistedSession,
  YinianSavedCredentials,
  YinianSkillRegistry,
  YinianSkillRegistryByHotel,
} from '../../shared/yinian';

interface YinianStoreShape {
  session?: YinianPersistedSession;
  savedCredentials?: YinianSavedCredentials & {
    passwordEncrypted?: string;
    passwordEncoding?: 'electron-safe-storage' | 'plain';
  };
  configs?: Record<string, YinianConfigSnapshot>;
  skillRegistryByHotel?: YinianSkillRegistryByHotel;
}

export interface YinianStorage {
  getSession(): Promise<YinianPersistedSession | undefined>;
  setSession(session: YinianPersistedSession): Promise<void>;
  clearSession(): Promise<void>;
  getSavedCredentials(): Promise<YinianSavedCredentials | undefined>;
  setSavedCredentials(credentials: YinianSavedCredentials): Promise<void>;
  clearSavedCredentials(): Promise<void>;
  getConfig(hotelId: string): Promise<YinianConfigSnapshot | undefined>;
  setConfig(config: YinianConfigSnapshot): Promise<void>;
  clearConfig(hotelId: string): Promise<void>;
  getSkillRegistry(hotelId?: string): Promise<YinianSkillRegistryByHotel | YinianSkillRegistry | undefined>;
  setSkillRegistry(registry: YinianSkillRegistry): Promise<void>;
  clearSkillRegistry(hotelId?: string): Promise<void>;
  clearAll(): Promise<void>;
}

// Lazy-load electron-store only in the Electron main process. Tests and pure
// modules can inject createMemoryYinianStorage instead.
let storeInstance: {
  get<K extends keyof YinianStoreShape>(key: K): YinianStoreShape[K];
  set<K extends keyof YinianStoreShape>(key: K, value: YinianStoreShape[K]): void;
  delete(key: keyof YinianStoreShape): void;
  clear(): void;
} | null = null;

async function getStore() {
  if (!storeInstance) {
    const Store = (await import('electron-store')).default;
    storeInstance = new Store<YinianStoreShape>({
      name: 'yinian',
      defaults: {
        configs: {},
        skillRegistryByHotel: {},
      },
    });
  }
  return storeInstance;
}

async function encryptPassword(password: string | undefined): Promise<{
  password?: string;
  passwordEncrypted?: string;
  passwordEncoding?: 'electron-safe-storage' | 'plain';
}> {
  if (!password) return {};
  try {
    const { safeStorage } = await import('electron');
    if (safeStorage?.isEncryptionAvailable()) {
      return {
        passwordEncrypted: safeStorage.encryptString(password).toString('base64'),
        passwordEncoding: 'electron-safe-storage',
      };
    }
  } catch {
    // Fall through to the plain fallback used in dev/test environments.
  }
  return { password, passwordEncoding: 'plain' };
}

async function decryptCredentials(
  credentials: YinianStoreShape['savedCredentials'],
): Promise<YinianSavedCredentials | undefined> {
  if (!credentials) return undefined;
  if (credentials.passwordEncrypted && credentials.passwordEncoding === 'electron-safe-storage') {
    try {
      const { safeStorage } = await import('electron');
      if (safeStorage?.isEncryptionAvailable()) {
        return {
          account: credentials.account,
          password: safeStorage.decryptString(Buffer.from(credentials.passwordEncrypted, 'base64')),
          rememberPassword: credentials.rememberPassword,
          updatedAt: credentials.updatedAt,
        };
      }
    } catch {
      return {
        account: credentials.account,
        rememberPassword: false,
        updatedAt: credentials.updatedAt,
      };
    }
  }
  return {
    account: credentials.account,
    password: credentials.password,
    rememberPassword: credentials.rememberPassword,
    updatedAt: credentials.updatedAt,
  };
}

async function prepareCredentialsForStorage(credentials: YinianSavedCredentials): Promise<YinianStoreShape['savedCredentials']> {
  const passwordState = await encryptPassword(credentials.rememberPassword ? credentials.password : undefined);
  return {
    account: credentials.account,
    rememberPassword: credentials.rememberPassword,
    updatedAt: credentials.updatedAt,
    ...passwordState,
  };
}

export function createElectronYinianStorage(): YinianStorage {
  return {
    async getSession() {
      return (await getStore()).get('session');
    },
    async setSession(session) {
      (await getStore()).set('session', session);
    },
    async clearSession() {
      (await getStore()).delete('session');
    },
    async getSavedCredentials() {
      return decryptCredentials((await getStore()).get('savedCredentials'));
    },
    async setSavedCredentials(credentials) {
      (await getStore()).set('savedCredentials', await prepareCredentialsForStorage(credentials));
    },
    async clearSavedCredentials() {
      (await getStore()).delete('savedCredentials');
    },
    async getConfig(hotelId) {
      return (await getStore()).get('configs')?.[hotelId];
    },
    async setConfig(config) {
      const store = await getStore();
      store.set('configs', {
        ...(store.get('configs') ?? {}),
        [config.hotel.id]: config,
      });
    },
    async clearConfig(hotelId) {
      const store = await getStore();
      const configs = { ...(store.get('configs') ?? {}) };
      delete configs[hotelId];
      store.set('configs', configs);
    },
    async getSkillRegistry(hotelId) {
      const registries = (await getStore()).get('skillRegistryByHotel') ?? {};
      return hotelId ? registries[hotelId] : registries;
    },
    async setSkillRegistry(registry) {
      const store = await getStore();
      store.set('skillRegistryByHotel', {
        ...(store.get('skillRegistryByHotel') ?? {}),
        [registry.hotelId]: registry,
      });
    },
    async clearSkillRegistry(hotelId) {
      const store = await getStore();
      if (!hotelId) {
        store.set('skillRegistryByHotel', {});
        return;
      }
      const registries = { ...(store.get('skillRegistryByHotel') ?? {}) };
      delete registries[hotelId];
      store.set('skillRegistryByHotel', registries);
    },
    async clearAll() {
      (await getStore()).clear();
    },
  };
}

export function createMemoryYinianStorage(initial?: Partial<YinianStoreShape>): YinianStorage {
  const state: YinianStoreShape = {
    configs: {},
    skillRegistryByHotel: {},
    ...initial,
  };

  return {
    async getSession() {
      return state.session;
    },
    async setSession(session) {
      state.session = session;
    },
    async clearSession() {
      delete state.session;
    },
    async getSavedCredentials() {
      return decryptCredentials(state.savedCredentials);
    },
    async setSavedCredentials(credentials) {
      state.savedCredentials = await prepareCredentialsForStorage(credentials);
    },
    async clearSavedCredentials() {
      delete state.savedCredentials;
    },
    async getConfig(hotelId) {
      return state.configs?.[hotelId];
    },
    async setConfig(config) {
      state.configs = {
        ...(state.configs ?? {}),
        [config.hotel.id]: config,
      };
    },
    async clearConfig(hotelId) {
      if (!state.configs) return;
      delete state.configs[hotelId];
    },
    async getSkillRegistry(hotelId) {
      const registries = state.skillRegistryByHotel ?? {};
      return hotelId ? registries[hotelId] : registries;
    },
    async setSkillRegistry(registry) {
      state.skillRegistryByHotel = {
        ...(state.skillRegistryByHotel ?? {}),
        [registry.hotelId]: registry,
      };
    },
    async clearSkillRegistry(hotelId) {
      if (!hotelId) {
        state.skillRegistryByHotel = {};
        return;
      }
      if (!state.skillRegistryByHotel) return;
      delete state.skillRegistryByHotel[hotelId];
    },
    async clearAll() {
      delete state.session;
      delete state.savedCredentials;
      state.configs = {};
      state.skillRegistryByHotel = {};
    },
  };
}

let yinianStorage: YinianStorage | null = null;

export function getYinianStorage(): YinianStorage {
  yinianStorage ??= createElectronYinianStorage();
  return yinianStorage;
}

export function resetYinianStorageForTests(next?: YinianStorage): void {
  yinianStorage = next ?? null;
  storeInstance = null;
}