import { MatrixClient, RoomStateEvent } from 'matrix-js-sdk';
import type { StateEventCallback } from '../hooks/stateEventCallback.types';

type ClientRegistry = {
  listener: StateEventCallback;
  subscribers: Set<StateEventCallback>;
};

const registries = new WeakMap<MatrixClient, ClientRegistry>();

const getOrCreateRegistry = (mx: MatrixClient): ClientRegistry => {
  let reg = registries.get(mx);
  if (!reg) {
    const subscribers = new Set<StateEventCallback>();
    const listener: StateEventCallback = (event, state, lastStateEvent) => {
      subscribers.forEach((cb) => cb(event, state, lastStateEvent));
    };
    mx.on(RoomStateEvent.Events, listener);
    reg = { listener, subscribers };
    registries.set(mx, reg);
  }
  return reg;
};

export const subscribeStateEvents = (
  mx: MatrixClient,
  callback: StateEventCallback
): (() => void) => {
  const reg = getOrCreateRegistry(mx);
  reg.subscribers.add(callback);
  return () => {
    reg.subscribers.delete(callback);
    if (reg.subscribers.size === 0) {
      mx.removeListener(RoomStateEvent.Events, reg.listener);
      registries.delete(mx);
    }
  };
};
