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
  cleanup: () => 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 cleanup = () => {
    for (const key in eventMap) {
      delete eventMap[key]
    }
  }

  return { on, emit, cleanup }
}