import { useRef, useState } from "react"
import { useIsomorphicLayoutEffect } from "usehooks-ts"

// Modified hook useScrollLock originally from usehooks-ts

/** Hook options. */
type UseScrollLockOptions = {
  /**
   * Whether to lock the scroll initially.
   * @default true
   */
  autoLock?: boolean
  /**
   * The target element to lock the scroll (default is the body element).
   * @default document.body
   */
  lockTarget?: HTMLElement | string
  /**
   * Whether to prevent width reflow when locking the scroll.
   * @default true
   */
  widthReflow?: boolean
}

/** Hook return type. */
type UseScrollLockReturn = {
  /** Whether the scroll is locked. */
  isLocked: boolean
  /** Lock the scroll. */
  lockScroll: () => void
  /** Unlock the scroll. */
  unlockScroll: () => void
}

type OriginalStyle = {
  overflow: CSSStyleDeclaration["overflow"]
  paddingRight: CSSStyleDeclaration["paddingRight"]
  position: CSSStyleDeclaration["position"]
  top: CSSStyleDeclaration["top"]
  scrollPosition: number
}

const IS_SERVER = typeof window === "undefined"

/**
 * A custom hook that locks and unlocks scroll.
 * @param {UseScrollLockOptions} [options] - Options to configure the hook, by default it will lock the scroll automatically.
 * @returns {UseScrollLockReturn} - An object containing the lockScroll and unlockScroll functions.
 * @public
 * @see [Documentation](https://usehooks-ts.com/react-hook/use-scroll-lock)
 * @example
 * ```tsx
 * // Lock the scroll when the modal is mounted, and unlock it when it's unmounted
 * useScrollLock()
 * ```
 * @example
 * ```tsx
 * // Manually lock and unlock the scroll
 * const { lockScroll, unlockScroll } = useScrollLock({ autoLock: false })
 *
 * return (
 *  <div>
 *   <button onClick={lockScroll}>Lock</button>
 *   <button onClick={unlockScroll}>Unlock</button>
 *  </div>
 * )
 * ```
 */
export function useScrollLock(
  options: UseScrollLockOptions = {}
): UseScrollLockReturn {
  const { autoLock = true, lockTarget, widthReflow = true } = options
  const [isLocked, setIsLocked] = useState(false)
  const target = useRef<HTMLElement | null>(null)
  const originalStyle = useRef<OriginalStyle | null>(null)
  let scrollPosition = 0

  const lockScroll = () => {
    if (target.current) {
      const { overflow, paddingRight, position, top } = target.current.style

      // Save the original styles
      originalStyle.current = {
        overflow,
        paddingRight,
        position,
        top,
        scrollPosition,
      }
      originalStyle.current.scrollPosition = window.scrollY
      // Prevent width reflow
      if (widthReflow) {
        // Use window inner width if body is the target as global scrollbar isn't part of the document
        const offsetWidth =
          target.current === document.body
            ? window.innerWidth
            : target.current.offsetWidth
        // Get current computed padding right in pixels
        const currentPaddingRight =
          parseInt(window.getComputedStyle(target.current).paddingRight, 10) ||
          0
        const scrollbarWidth = offsetWidth - target.current.scrollWidth
        target.current.style.paddingRight = `${scrollbarWidth + currentPaddingRight}px`
      }
      // Lock the scroll
      target.current.style.overflow = "hidden"
      target.current.style.position = "fixed"
      document.body.style.top = `-${originalStyle.current.scrollPosition}px`
      setIsLocked(true)
    }
  }

  const unlockScroll = () => {
    if (target.current && originalStyle.current) {
      target.current.style.overflow = originalStyle.current.overflow
      target.current.style.top = originalStyle.current.top
      target.current.style.position = originalStyle.current.position

      // Only reset padding right if we changed it
      if (widthReflow) {
        target.current.style.paddingRight = originalStyle.current.paddingRight
      }
      //Scroll to original scroll position
      window.scrollTo({
        left: 0,
        top: originalStyle.current.scrollPosition,
        behavior: "instant",
      })
    }

    setIsLocked(false)
  }

  useIsomorphicLayoutEffect(() => {
    if (IS_SERVER) return

    if (lockTarget) {
      target.current =
        typeof lockTarget === "string"
          ? document.querySelector(lockTarget)
          : lockTarget
    }

    if (!target.current) {
      target.current = document.body
    }

    if (autoLock) {
      lockScroll()
    }

    return () => {
      unlockScroll()
    }
  }, [autoLock, lockTarget, widthReflow])

  return { isLocked, lockScroll, unlockScroll }
}