sse-session/src/utils/event-emitter.ts

export type EventMap = Record<string, unknown>;

type Listener<T> = (detail: T) => void;

/**
 * A listener entry.
 * @template T - The event type.
 */
interface ListenerEntry<T> {
  listener: Listener<T>;
  wrappedListener: Listener<T>;
  debounceTime?: number;
  once?: boolean;
}

export type OffCallback = () => void;

/**
 * A simple event emitter implementation.
 * @template T - The event map type.
 */
export class EventEmitter<T extends EventMap> {
  /**
   * The listeners map.
   * @private
   */
  private listeners: Map<keyof T, Set<ListenerEntry<T[keyof T]>>> = new Map();

  /**
   * Add a listener for an event.
   * @param type - The event type.
   * @param listener - The listener function.
   * @param debounceMilliseconds - The debounce time in milliseconds.
   * @returns An off callback that can be called to stop listening for events.
   */
  on<K extends keyof T>(
    type: K,
    listener: Listener<T[K]>,
    debounceMilliseconds?: number,
  ): OffCallback {
    // Create a wrapped listener so that the debounce can be applied.
    const wrappedListener =
      debounceMilliseconds && debounceMilliseconds > 0
        ? this.debounce(listener, debounceMilliseconds)
        : listener;

    // If the listeners map does not have the event type, create a new set.
    if (!this.listeners.has(type)) {
      this.listeners.set(type, new Set());
    }

    // Create a listener entry.
    const listenerEntry: ListenerEntry<T[K]> = {
      listener,
      wrappedListener,
      ...(debounceMilliseconds !== undefined
        ? { debounceTime: debounceMilliseconds }
        : {}),
    };

    // Add the listener entry to the listeners map.
    this.listeners.get(type)?.add(listenerEntry as ListenerEntry<T[keyof T]>);

    // Return an "off" callback that can be called to stop listening for events.
    return () => this.off(type, listener);
  }

  /**
   * Add a one-time listener for an event.
   * @param type - The event type.
   * @param listener - The listener function.
   * @param debounceMilliseconds - The debounce time in milliseconds.
   * @returns An off callback that can be called to stop listening for events.
   */
  once<K extends keyof T>(
    type: K,
    listener: Listener<T[K]>,
    debounceMilliseconds?: number,
  ): OffCallback {
    const wrappedListener: Listener<T[K]> = (detail: T[K]) => {
      this.off(type, listener);
      listener(detail);
    };

    // Create a debounced listener.
    const debouncedListener =
      debounceMilliseconds && debounceMilliseconds > 0
        ? this.debounce(wrappedListener, debounceMilliseconds)
        : wrappedListener;

    // If the listeners map does not have the event type, create a new set.
    if (!this.listeners.has(type)) {
      this.listeners.set(type, new Set());
    }

    // Create a listener entry.
    const listenerEntry: ListenerEntry<T[K]> = {
      listener,
      wrappedListener: debouncedListener,
      once: true,
      ...(debounceMilliseconds !== undefined
        ? { debounceTime: debounceMilliseconds }
        : {}),
    };

    // Add the listener entry to the listeners map.
    this.listeners.get(type)?.add(listenerEntry as ListenerEntry<T[keyof T]>);

    // Return an "off" callback that can be called to stop listening for events.
    return () => this.off(type, listener);
  }

  /**
   * Remove a listener for an event.
   * @param type - The event type.
   * @param listener - The listener function.
   */
  off<K extends keyof T>(type: K, listener: Listener<T[K]>): void {
    // Get the listeners for the event type.
    const listeners = this.listeners.get(type);
    if (!listeners) return;

    // Find the listener entry.
    const listenerEntry = Array.from(listeners).find(
      (entry) =>
        entry.listener === listener || entry.wrappedListener === listener,
    );

    // If the listener entry is found, remove it from the listeners map.
    if (listenerEntry) {
      listeners.delete(listenerEntry);
    }
  }

  /**
   * Emit an event.
   * @param type - The event type.
   * @param payload - The event payload.
   * @returns True if there are listeners for the event, false otherwise.
   */
  emit<K extends keyof T>(type: K, payload: T[K]): boolean {
    // Get the listeners for the event type.
    const listeners = this.listeners.get(type);
    if (!listeners) return false;

    // Emit the event to all listeners.
    listeners.forEach((entry) => {
      entry.wrappedListener(payload);
    });

    // Return true if there are listeners for the event, false otherwise.
    return listeners.size > 0;
  }

  /**
   * Remove all listeners.
   */
  removeAllListeners(): void {
    this.listeners.clear();
  }

  /**
   * Wait for an event to be emitted.
   * @param type - The event type.
   * @param predicate - The predicate function.
   * @param timeoutMs - The timeout in milliseconds.
   * @returns The event payload.
   */
  async waitFor<K extends keyof T>(
    type: K,
    predicate: (payload: T[K]) => boolean,
    timeoutMs?: number,
  ): Promise<T[K]> {
    // Create a promise to wait for the event to be emitted.
    return new Promise((resolve, reject) => {
      let timeoutId: ReturnType<typeof setTimeout> | undefined;

      // Create a listener function.
      const listener = (payload: T[K]): void => {
        if (predicate(payload)) {
          // Clean up
          this.off(type, listener);
          if (timeoutId !== undefined) {
            clearTimeout(timeoutId);
          }

          resolve(payload);
        }
      };

      // Set up timeout if specified
      if (timeoutMs !== undefined) {
        timeoutId = setTimeout(() => {
          this.off(type, listener);
          reject(new Error(`Timeout waiting for event "${String(type)}"`));
        }, timeoutMs);
      }

      // Add the listener to the listeners map.
      this.on(type, listener);
    });
  }

  /**
   * Debounce a function.
   * @param func - The function to debounce.
   * @param wait - The wait time in milliseconds.
   * @returns The debounced function.
   */
  private debounce<K extends keyof T>(
    func: Listener<T[K]>,
    wait: number,
  ): Listener<T[K]> {
    // Create a timeout variable.
    let timeout: ReturnType<typeof setTimeout>;

    return (detail: T[K]) => {
      // If a debounce timer is already pending, clear it before scheduling the next one.
      if (timeout !== undefined) {
        clearTimeout(timeout);
      }

      timeout = setTimeout(() => {
        func(detail);
      }, wait);
    };
  }
}