type CallbackFn<T> = (args: T) => void

export type EventEmitter<T extends Record<string, unknown>> = {
  on: <U extends keyof T>(event: U, callback: CallbackFn<T[U]>) => () => void
  emit: <U extends keyof T>(event: U, arg?: T[U]) => void
  removeAllListeners: <U extends keyof T>(event?: U) => void
}

export const createEventEmitter = <T extends Record<string, unknown>>(): EventEmitter<T> => {
  const eventMap: Partial<Record<keyof T, Set<CallbackFn<any>>>> = {}

  const on = <U extends keyof T>(event: U, callback: CallbackFn<T[U]>) => {
    if (!eventMap[event]) {
      eventMap[event] = new Set<CallbackFn<T[U]>>()
    }

    eventMap[event]?.add(callback)

    return () => eventMap[event]?.delete(callback)
  }

  const emit = <U extends keyof T>(event: U, arg?: T[U]) => {
    eventMap[event]?.forEach((callback) => callback(arg))
    eventMap['*']?.forEach((callback) => callback({ type: event, payload: arg }))
  }

  const removeAllListeners = <U extends keyof T>(event?: U): void => {
    if (event) {
      eventMap[event]?.clear()
    } else {
      Object.keys(eventMap).forEach((key) => {
        eventMap[key as keyof T]?.clear()
      })
    }
  }

  return { on, emit, removeAllListeners }
}