import { EventEmitter, EventSubscription } from 'fbemitter';

export default class MockWebSocket {
  static readonly CONNECTING = 0;
  static readonly OPEN = 1;
  static readonly CLOSING = 2;
  static readonly CLOSED = 3;

  public readyState = MockWebSocket.CONNECTING;

  private readonly emitter = new EventEmitter();
  private subscriptions: EventSubscription[] = [];
  private broadcastSubscriptions: EventSubscription[] = [];

  // We use the broadcastEmitter to simulate event passing on a WebSocket server.
  // The map is {address -> EventEmitter}
  private static broadcastEmitter: Record<string, EventEmitter> = {};

  constructor(private readonly address: string) {
    MockWebSocket.broadcastEmitter[address] ??= new EventEmitter();
    setTimeout(() => {
      this.readyState = MockWebSocket.OPEN;
      this.emitter.emit('open');

      const subscription = MockWebSocket.broadcastEmitter[this.address].addListener(
        'broadcast',
        this.handleBroadcast
      );
      this.broadcastSubscriptions.push(subscription);
    }, 0);
  }

  addEventListener = jest.fn().mockImplementation((event, listener) => {
    const subscription = this.emitter.addListener(event, listener);
    this.subscriptions.push(subscription);

    const broadcastSubscription = MockWebSocket.broadcastEmitter[this.address].addListener(
      event,
      listener
    );
    this.broadcastSubscriptions.push(broadcastSubscription);
  });
  removeEventListener = jest.fn().mockImplementation((event, listener) => {
    const index = this.subscriptions.findIndex(
      (subscription) => subscription.listener === listener
    );
    if (index >= 0) {
      this.subscriptions[index].remove();
      this.subscriptions.splice(index, 1);
    }

    const broadcastIndex = this.broadcastSubscriptions.findIndex(
      (subscription) => subscription.listener === listener
    );
    if (broadcastIndex >= 0) {
      this.broadcastSubscriptions[broadcastIndex].remove();
      this.broadcastSubscriptions.splice(broadcastIndex, 1);
    }
  });

  close = jest.fn().mockImplementation(() => {
    for (const subscription of this.broadcastSubscriptions) {
      subscription.remove();
    }
    this.broadcastSubscriptions = [];
    this.readyState = MockWebSocket.CLOSED;
  });

  send = jest.fn().mockImplementation((data) => {
    MockWebSocket.broadcastEmitter[this.address].emit('broadcast', { sender: this, data });
  });

  private handleBroadcast = ({ sender, data }: any) => {
    if (sender !== this) {
      this.emitter.emit('message', { data });
    }
  };
}