trackgo-fe/packages/core-storage/src/indexed-db/indexed-db.service.ts

import { EncryptionUtils } from '@repo/utils';
import type { IStorageService } from '../types/storage.interface';
import type { StorageOptions } from '../local-storage/local-storage.service';

// ─── Types ──────────────────────────────────────────────────────

export interface IndexedDBConfig<TKey extends string> extends StorageOptions<TKey> {
  /** Database name. @default 'app_db' */
  dbName?: string;
  /** Object store name. @default 'kv_store' */
  storeName?: string;
  /** Database version. @default 1 */
  version?: number;
}

// ─── Helpers ────────────────────────────────────────────────────

/**
 * Open (or create) an IndexedDB database with a simple key-value store.
 * Returns a Promise that resolves with the IDBDatabase instance.
 */
function openDatabase(dbName: string, storeName: string, version: number): Promise<IDBDatabase> {
  return new Promise((resolve, reject) => {
    const request = indexedDB.open(dbName, version);

    request.onupgradeneeded = () => {
      const db = request.result;
      if (!db.objectStoreNames.contains(storeName)) {
        db.createObjectStore(storeName);
      }
    };

    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
  });
}

/**
 * Execute a single IndexedDB transaction and return the result.
 * Resolves only when the entire transaction is completed to prevent silent failures.
 */
function withTransaction<R>(
  db: IDBDatabase,
  storeName: string,
  mode: IDBTransactionMode,
  operation: (store: IDBObjectStore) => IDBRequest<R>,
): Promise<R> {
  return new Promise((resolve, reject) => {
    const tx = db.transaction(storeName, mode);
    const store = tx.objectStore(storeName);
    let requestResult: R;

    const request = operation(store);

    request.onsuccess = () => {
      requestResult = request.result;
    };

    request.onerror = () => reject(request.error);

    // Resolve promise ONLY after transaction is successfully committed
    tx.oncomplete = () => resolve(requestResult);
    tx.onerror = () => reject(tx.error);
    tx.onabort = () => reject(tx.error);
  });
}

// ─── Service ────────────────────────────────────────────────────

/**
 * Enterprise-grade IndexedDB wrapper with optional AES encryption.
 */
export class IndexedDBService<TKey extends string> implements IStorageService<TKey> {
  private readonly encryption: EncryptionUtils;
  private readonly dbName: string;
  private readonly storeName: string;
  private readonly version: number;
  private readonly encryptedKeys: Set<TKey>;
  private readonly plainTextKeys: Set<TKey>;
  private readonly personalizedKeys?: Set<TKey>;
  private readonly getUserId?: () => string | null | undefined;
  private dbPromise: Promise<IDBDatabase> | null = null;

  constructor(config?: IndexedDBConfig<TKey>, encryptionUtils?: EncryptionUtils) {
    this.encryption = encryptionUtils ?? EncryptionUtils.getInstance();
    this.dbName = config?.dbName ?? 'app_db';
    this.storeName = config?.storeName ?? 'kv_store';
    this.version = config?.version ?? 1;
    this.encryptedKeys = config?.encryptedKeys ?? new Set();
    this.plainTextKeys = config?.plainTextKeys ?? new Set();
    this.personalizedKeys = config?.personalizedKeys;
    this.getUserId = config?.getUserId;
  }

  /** Lazy-open the database connection (cached). */
  private getDB(): Promise<IDBDatabase> {
    if (!this.dbPromise) {
      this.dbPromise = openDatabase(this.dbName, this.storeName, this.version);
    }
    return this.dbPromise;
  }

  private validateKey(key: TKey): void {
    if (!this.encryptedKeys.has(key) && !this.plainTextKeys.has(key)) {
      throw new Error(`[Storage Engine] Security Exception: Key '${key}' is not registered and cannot be accessed.`);
    }
  }

  private shouldEncrypt(key: TKey): boolean {
    return this.encryptedKeys.has(key);
  }

  private resolveKey(key: TKey): string {
    if (this.personalizedKeys?.has(key)) {
      const userId = this.getUserId?.() || 'guest';
      return `${String(key)}_${userId}`;
    }
    return String(key);
  }

  async setItem<T>(key: TKey, value: T): Promise<void> {
    this.validateKey(key);
    const db = await this.getDB();
    const resolvedKey = this.resolveKey(key);

    const serialized = JSON.stringify(value);
    const payload = this.shouldEncrypt(key) ? this.encryption.encrypt(serialized) : serialized;

    await withTransaction(db, this.storeName, 'readwrite', (store) => store.put(payload, resolvedKey));
  }

  async getItem<T>(key: TKey): Promise<T | null> {
    this.validateKey(key);
    const db = await this.getDB();
    const resolvedKey = this.resolveKey(key);

    const raw = await withTransaction<string | undefined>(
      db,
      this.storeName,
      'readonly',
      (store) => store.get(resolvedKey) as IDBRequest<string | undefined>,
    );

    if (raw === undefined || raw === null) return null;

    try {
      const decrypted = this.shouldEncrypt(key) ? this.encryption.decrypt(raw) : raw;
      if (!decrypted) return null;

      return JSON.parse(decrypted) as T;
    } catch {
      console.warn(`[core-storage] Failed to parse IndexedDB key "${String(key)}". Removing corrupt entry.`);
      await this.removeItem(key);
      return null;
    }
  }

  async removeItem(key: TKey): Promise<void> {
    this.validateKey(key);
    const db = await this.getDB();
    const resolvedKey = this.resolveKey(key);

    await withTransaction(db, this.storeName, 'readwrite', (store) => store.delete(resolvedKey));
  }

  async clear(): Promise<void> {
    const allValidKeys = await this.keys();
    for (const key of allValidKeys) {
      await this.removeItem(key);
    }
  }

  async hasItem(key: TKey): Promise<boolean> {
    const value = await this.getItem(key);
    return value !== null;
  }

  async keys(): Promise<TKey[]> {
    const db = await this.getDB();
    const rawKeys = await withTransaction<string[]>(
      db,
      this.storeName,
      'readonly',
      (store) => store.getAllKeys() as IDBRequest<string[]>,
    );

    const resultSet = new Set<TKey>();
    const currentUserId = this.getUserId?.() || 'guest';

    for (const rawKey of rawKeys) {
      const isGlobalKey = this.encryptedKeys.has(rawKey as TKey) || this.plainTextKeys.has(rawKey as TKey);

      if (isGlobalKey && !this.personalizedKeys?.has(rawKey as TKey)) {
        resultSet.add(rawKey as TKey);
      }

      if (this.personalizedKeys) {
        for (const pKey of this.personalizedKeys) {
          if (rawKey === `${String(pKey)}_${currentUserId}`) {
            resultSet.add(pKey);
            break;
          }
        }
      }
    }

    return Array.from(resultSet);
  }
}

export function createIndexedDB<TKey extends string>(config: IndexedDBConfig<TKey>): IStorageService<TKey> {
  return new IndexedDBService<TKey>(config);
}