"use client"

import * as React from "react"
import type { NodeKey } from "lexical"

export interface LexicalImageUploadState {
  progress?: number
}

const uploadStates = new Map<NodeKey, LexicalImageUploadState>()
const listeners = new Map<NodeKey, Set<VoidFunction>>()

function emit(nodeKey: NodeKey) {
  listeners.get(nodeKey)?.forEach((listener) => listener())
}

export function getImageUploadState(
  nodeKey: NodeKey
): LexicalImageUploadState | undefined {
  return uploadStates.get(nodeKey)
}

export function setImageUploadState(
  nodeKey: NodeKey,
  state: LexicalImageUploadState
) {
  uploadStates.set(nodeKey, state)
  emit(nodeKey)
}

export function deleteImageUploadState(nodeKey: NodeKey) {
  if (!uploadStates.delete(nodeKey)) return
  emit(nodeKey)
}

function subscribe(nodeKey: NodeKey, listener: VoidFunction) {
  const nodeListeners = listeners.get(nodeKey) ?? new Set<VoidFunction>()
  nodeListeners.add(listener)
  listeners.set(nodeKey, nodeListeners)

  return () => {
    nodeListeners.delete(listener)
    if (nodeListeners.size === 0) listeners.delete(nodeKey)
  }
}

export function useImageUploadState(nodeKey: NodeKey) {
  return React.useSyncExternalStore(
    (listener) => subscribe(nodeKey, listener),
    () => getImageUploadState(nodeKey),
    () => undefined
  )
}