"use client";
import React, {
  useRef,
  useLayoutEffect,
  useEffect,
  useState,
  useCallback,
  forwardRef,
  useImperativeHandle,
  Ref,
} from "react";
import { useLenis } from "@studio-freight/react-lenis";

type MaybeElement = Window | HTMLElement | null;

export type CustomHeadroomHandle = {
  show: () => void;
  hide: () => void;
  toggle: () => void;
  getPinned: () => boolean;
};

export type CustomHeadroomProps = {
  children: React.ReactElement;
  /**
   * Pixel offset before headroom behavior starts (e.g. keep header visible for first X px)
   */
  pinStart?: number;
  /**
   * How many pixels of downward scroll are required to trigger hide
   */
  downTolerance?: number;
  /**
   * How many pixels of upward scroll are required to trigger show
   */
  upTolerance?: number;
  /**
   * Transition duration in ms for CSS transform
   */
  transitionDuration?: number;
  /**
   * If true, component will not attach scroll listeners during SSR (Next.js friendly)
   */
  disableOnSSR?: boolean;
  /**
   * Optional scroll container element to observe. If omitted, window is used.
   * Note: If useLenis is true, this will be ignored and Lenis scroll position will be used instead.
   */
  scrollContainer?: () => MaybeElement;
  /**
   * If true (default), a spacer div is rendered to preserve layout space
   * equal to the header height so content does not appear under the header.
   * Set to false when you want the header to overlay the content, e.g. over a hero.
   */
  preserveSpace?: boolean;
  /**
   * If true, uses Lenis smooth scroll position instead of native scroll events.
   * This is recommended when using Lenis for smooth scrolling.
   */
  useLenis?: boolean;
  /**
   * If true, headroom starts disabled (always visible)
   */
  disabled?: boolean;
  /**
   * Callback hooks
   */
  onPin?: () => void;
  onUnpin?: () => void;
  className?: string;
  style?: React.CSSProperties;
};

const isBrowser = typeof window !== "undefined";

function getScrollY(container: MaybeElement) {
  if (!container) return 0;
  if (container === window) return window.scrollY || window.pageYOffset || 0;
  return (container as HTMLElement).scrollTop || 0;
}

const defaultProps = {
  pinStart: 0,
  downTolerance: 5,
  upTolerance: 5,
  transitionDuration: 220,
  disableOnSSR: true,
  disabled: false,
};

const CustomHeadroom = forwardRef(function CustomHeadroom(
  props: CustomHeadroomProps,
  ref: Ref<CustomHeadroomHandle>
) {
  const {
    children,
    pinStart = defaultProps.pinStart,
    downTolerance = defaultProps.downTolerance,
    upTolerance = defaultProps.upTolerance,
    transitionDuration = defaultProps.transitionDuration,
    disableOnSSR = defaultProps.disableOnSSR,
    scrollContainer,
    preserveSpace = true,
    useLenis: useLenisScroll = false,
    disabled = defaultProps.disabled,
    onPin,
    onUnpin,
    className,
    style,
  } = props;

  // Always call useLenis hook (React rules), but only use it if useLenisScroll is enabled
  const lenisInstance = useLenis();
  const lenis = useLenisScroll ? lenisInstance : null;

  // refs and state
  const containerRef = useRef<HTMLDivElement | null>(null); // wrapper we transform
  const spacerRef = useRef<HTMLDivElement | null>(null);
  const headerHeightRef = useRef<number>(0);
  const lastScrollYRef = useRef<number>(0);
  const pinnedRef = useRef<boolean>(true); // visible initially
  const tickingRef = useRef<number | null>(null);
  const ignoreScrollRef = useRef(false); // used during programmatic show/hide
  const scrollTimeoutRef = useRef<number | null>(null);

  const [hydrated, setHydrated] = useState<boolean>(isBrowser ? true : false);
  const [visible, setVisible] = useState<boolean>(true); // UI state for render
  const [measured, setMeasured] = useState<boolean>(false); // header measured

  // Resolve scroll container
  const resolvedContainerRef = useRef<MaybeElement>(isBrowser ? window : null);
  useEffect(() => {
    if (!isBrowser) {
      resolvedContainerRef.current = null;
      return;
    }
    resolvedContainerRef.current = scrollContainer
      ? scrollContainer() ?? window
      : window;
  }, [scrollContainer]);

  // expose imperative API
  useImperativeHandle(
    ref,
    (): CustomHeadroomHandle => ({
      show: () => {
        ignoreScrollRef.current = true;
        pinnedRef.current = true;
        setVisible(true);
        scheduleTransform(0);
        onPin?.();
        // re-enable scroll after a brief tick
        window.setTimeout(() => {
          ignoreScrollRef.current = false;
        }, 50);
      },
      hide: () => {
        ignoreScrollRef.current = true;
        pinnedRef.current = false;
        setVisible(false);
        scheduleTransform(-headerHeightRef.current || -999);
        onUnpin?.();
        window.setTimeout(() => {
          ignoreScrollRef.current = false;
        }, 50);
      },
      toggle: () => {
        if (pinnedRef.current) {
          const handle = ref as React.RefObject<CustomHeadroomHandle>;
          handle?.current?.hide();
        } else {
          const handle = ref as React.RefObject<CustomHeadroomHandle>;
          handle?.current?.show();
        }
      },
      getPinned: () => pinnedRef.current,
    })
  );

  // measure header height and map children
  useLayoutEffect(() => {
    if (!isBrowser) return;
    // get the child DOM node: prefer forwarded ref, else query
    const wrapper = containerRef.current;
    if (!wrapper) return;
    const childEl = wrapper.firstElementChild as HTMLElement | null;
    const measure = () => {
      const rect = childEl?.getBoundingClientRect();
      const h = rect?.height ?? 0;
      headerHeightRef.current = Math.round(h);
      if (spacerRef.current)
        spacerRef.current.style.height = preserveSpace
          ? `${headerHeightRef.current}px`
          : "0px";
      setMeasured(true);
    };
    measure();
    // watch for resize (header could be responsive)
    const RO = window.ResizeObserver;
    let ro: ResizeObserver | null = null;

    if (RO) {
      ro = new RO(() => measure());
      if (childEl && ro) ro.observe(childEl);
    }

    window.addEventListener("resize", measure);

    return () => {
      if (ro && childEl) ro.unobserve(childEl);
      window.removeEventListener("resize", measure);
    };
  }, [children, preserveSpace]);

  // helper to set transform via rAF
  const applyTransform = useCallback(
    (y: number) => {
      const el = containerRef.current;
      if (!el) return;
      el.style.willChange = "transform";
      el.style.transform = `translateY(${y}px)`;
      el.style.transition = `transform ${transitionDuration}ms cubic-bezier(.2,.9,.2,1)`;
    },
    [transitionDuration]
  );

  // schedule rAF
  const scheduleTransform = useCallback(
    (y: number) => {
      if (tickingRef.current != null) {
        cancelAnimationFrame(tickingRef.current);
        tickingRef.current = null;
      }
      tickingRef.current = requestAnimationFrame(() => {
        applyTransform(y);
        tickingRef.current = null;
      });
    },
    [applyTransform]
  );

  // scroll handler core
  const handleScroll = useCallback(() => {
    if (!isBrowser) return;
    if (disableOnSSR && !hydrated) return;
    if (disabled) return;
    if (ignoreScrollRef.current) return;

    // Use Lenis scroll position if enabled, otherwise use native scroll
    let currentY: number;
    if (useLenisScroll && lenis) {
      currentY = lenis.scroll;
    } else {
      const container = resolvedContainerRef.current;
      currentY = getScrollY(container);
    }

    const lastY = lastScrollYRef.current;
    const delta = currentY - lastY;
    lastScrollYRef.current = currentY;

    // small guard
    if (!measured) {
      scheduleTransform(0);
      return;
    }

    // if we are near top, always show
    if (currentY <= pinStart) {
      if (!pinnedRef.current) {
        pinnedRef.current = true;
        setVisible(true);
        onPin?.();
      }
      scheduleTransform(0);
      return;
    }

    // scrolling down
    if (delta > downTolerance) {
      // hide
      if (pinnedRef.current) {
        pinnedRef.current = false;
        setVisible(false);
        onUnpin?.();
      }
      scheduleTransform(-headerHeightRef.current);
      return;
    }

    // scrolling up
    if (delta < -upTolerance) {
      if (!pinnedRef.current) {
        pinnedRef.current = true;
        setVisible(true);
        onPin?.();
      }
      scheduleTransform(0);
      return;
    }

    // otherwise do nothing (within tolerance)
  }, [
    disableOnSSR,
    hydrated,
    disabled,
    measured,
    pinStart,
    downTolerance,
    upTolerance,
    scheduleTransform,
    onPin,
    onUnpin,
    useLenisScroll,
    lenis,
  ]);

  // touch / momentum protection: clears lastScrollY after small idle
  const handleScrollSettled = useCallback(() => {
    if (scrollTimeoutRef.current) {
      clearTimeout(scrollTimeoutRef.current);
      scrollTimeoutRef.current = null;
    }
    scrollTimeoutRef.current = window.setTimeout(() => {
      // stabilize lastScrollY to current
      lastScrollYRef.current = getScrollY(resolvedContainerRef.current);
    }, 100);
  }, []);

  // focus handling: if any element inside header receives focus, show header
  useEffect(() => {
    if (!isBrowser) return;
    const root = containerRef.current;
    if (!root) return;
    const onFocus = (e: FocusEvent) => {
      const tgt = e.target as Node | null;
      if (!tgt) return;
      if (root.contains(tgt)) {
        // immediate show
        if (!pinnedRef.current) {
          pinnedRef.current = true;
          setVisible(true);
          scheduleTransform(0);
          onPin?.();
        }
      }
    };
    document.addEventListener("focusin", onFocus);
    return () => document.removeEventListener("focusin", onFocus);
  }, [scheduleTransform, onPin]);

  // attach scroll listeners - either via Lenis or native scroll events
  useEffect(() => {
    if (!isBrowser) return;
    if (disableOnSSR && !hydrated) {
      // wait until hydrated
      setHydrated(true);
      return;
    }
    if (disabled) return;

    // If using Lenis, subscribe to Lenis scroll events
    if (useLenisScroll && lenis) {
      // Initialize lastScrollY with current Lenis scroll position
      lastScrollYRef.current = lenis.scroll;

      const onLenisScroll = () => {
        handleScroll();
        handleScrollSettled();
      };

      lenis.on("scroll", onLenisScroll);

      return () => {
        lenis.off("scroll", onLenisScroll);
      };
    }

    // Otherwise, use native scroll events
    const container = resolvedContainerRef.current ?? window;

    // Initialize lastScrollY
    lastScrollYRef.current = getScrollY(container);

    const onScroll = () => {
      handleScroll();
      handleScrollSettled();
    };

    // passive for performance
    container.addEventListener("scroll", onScroll, {
      passive: true,
    } as AddEventListenerOptions);
    // also listen to touchend in some mobile cases
    container.addEventListener("touchend", handleScrollSettled, {
      passive: true,
    });

    return () => {
      try {
        container.removeEventListener("scroll", onScroll as EventListener);
        container.removeEventListener(
          "touchend",
          handleScrollSettled as EventListener
        );
      } catch {
        // ignore
      }
    };
  }, [
    handleScroll,
    handleScrollSettled,
    disableOnSSR,
    hydrated,
    disabled,
    useLenisScroll,
    lenis,
  ]);

  // route popstate / history navigation: show header to avoid hidden nav after back
  useEffect(() => {
    if (!isBrowser) return;
    const onNav = () => {
      // show header on navigation so users always have access
      ignoreScrollRef.current = true;
      pinnedRef.current = true;
      setVisible(true);
      scheduleTransform(0);
      onPin?.();
      window.setTimeout(() => {
        ignoreScrollRef.current = false;
      }, 60);
    };
    window.addEventListener("popstate", onNav);
    return () => window.removeEventListener("popstate", onNav);
  }, [onPin, scheduleTransform]);

  // cleanup rAF on unmount
  useEffect(() => {
    return () => {
      if (tickingRef.current != null) cancelAnimationFrame(tickingRef.current);
      if (scrollTimeoutRef.current) clearTimeout(scrollTimeoutRef.current);
    };
  }, []);

  // aria-hidden when not visible
  const child = React.cloneElement(children, {
    ref: (node: HTMLElement | null) => {
      // try to keep user's ref if they provided one
      const existingRef = (
        children as React.ReactElement & { ref?: React.Ref<HTMLElement> }
      ).ref;
      if (typeof existingRef === "function") {
        existingRef(node);
      } else if (
        existingRef &&
        typeof existingRef === "object" &&
        "current" in existingRef
      ) {
        (existingRef as React.MutableRefObject<HTMLElement | null>).current =
          node;
      }
    },
    "aria-hidden": !visible ? "true" : undefined,
  } as React.HTMLAttributes<HTMLElement>);

  // spacer preserves layout (prevents content jump when header collapses)
  return (
    <div
      className={`custom-headroom-root ${className ?? ""}`}
      style={{
        ...style,
        width: "100%",
        maxWidth: "100vw",
        overflowX: "hidden",
      }}
    >
      {/* spacer keeps layout height */}
      <div
        ref={spacerRef}
        aria-hidden="true"
        style={{
          height:
            preserveSpace && measured ? `${headerHeightRef.current}px` : "0px",
        }}
      />
      {/* transform wrapper */}
      <div
        ref={containerRef}
        className="custom-headroom-container"
        style={{
          position: "fixed",
          top: 0,
          left: 0,
          right: 0,
          width: "100%",
          maxWidth: "100vw",
          zIndex: 60,
          transform: "translateY(0)",
          transition: `transform ${transitionDuration}ms cubic-bezier(.2,.9,.2,1)`,
          willChange: "transform",
        }}
      >
        {child}
      </div>
    </div>
  );
});

export default CustomHeadroom;