import React, { useCallback, useEffect, useMemo, useRef } from 'react'

import { createUseContext } from '@lib/react'

export type LockKeyId = number & { readonly __key: never }

const ScrollLockContext = React.createContext<{
  /**
   * Locks the scroll and returns a unique ID that should be used to release the lock.
   */
  lock: () => LockKeyId

  /**
   * Releases the lock and removes the scroll lock if there are no more locks.
   */
  release: (id: LockKeyId) => void
} | null>(null)

/**
 * Global context that takes care of locking the scroll when overflowing elements are visible.
 */
export const ScrollLockContextProvider = ({ children }: React.PropsWithChildren) => {
  const visibleModals = useRef(new Set<LockKeyId>())

  const nextKeyId = useRef<LockKeyId>(0 as LockKeyId)

  const release = useCallback((id: LockKeyId) => {
    visibleModals.current.delete(id)

    if (visibleModals.current.size > 0) {
      document.body.classList.add('scroll-locked')
    } else {
      document.body.classList.remove('scroll-locked')
    }
  }, [])

  const lock = useCallback(() => {
    const id = nextKeyId.current++ as LockKeyId

    visibleModals.current.add(id)

    document.body.classList.add('scroll-locked')

    return id
  }, [])

  // Remove lock on unmount.
  useEffect(() => {
    return () => {
      document.body.classList.remove('scroll-locked')
    }
  }, [])

  const context = useMemo(() => {
    return { lock, release }
  }, [lock, release])

  return <ScrollLockContext.Provider value={context}>{children}</ScrollLockContext.Provider>
}

/**
 * Lets you lock the scroll when overflowing elements are visible.
 */
export const useScrollLockContext = createUseContext(ScrollLockContext, 'ScrollLockContext')

/**
 * Acquires and manages a scroll lock based on the value of `show`.
 */
export function useScrollLock(show: boolean) {
  const { lock, release } = useScrollLockContext()

  const lockKeyRef = useRef<LockKeyId | null>(null)

  useEffect(() => {
    // NOTE: If there's an existing lock, we always release it first.
    if (lockKeyRef.current != null) {
      release(lockKeyRef.current)
      lockKeyRef.current = null
    }

    if (show) {
      lockKeyRef.current = lock()
    }

    return () => {
      // NOTE: This is to make sure that the modal is not marked visible when it is unmounted.
      if (lockKeyRef.current != null) {
        release(lockKeyRef.current)
        lockKeyRef.current = null
      }
    }
  }, [lock, release, show])
}
