docs/components/SectionProvider.tsx

"use client";

import {
    type ReactNode,
    type RefObject,
    createContext,
    useContext,
    useEffect,
    useLayoutEffect,
    useState,
} from "react";
import { type StoreApi, createStore, useStore } from "zustand";

import { remToPx } from "../lib/remToPx";

export interface Section {
    id: string;
    title: string;
    offsetRem?: number;
    tag?: string;
    headingRef?: RefObject<HTMLHeadingElement>;
}

interface SectionState {
    sections: Section[];
    visibleSections: string[];
    setVisibleSections: (visibleSections: string[]) => void;
    registerHeading: ({
        id,
        ref,
        offsetRem,
    }: {
        id: string;
        ref: RefObject<HTMLHeadingElement>;
        offsetRem: number;
    }) => void;
}

function createSectionStore(sections: Section[]) {
    return createStore<SectionState>()((set) => ({
        sections,
        visibleSections: [],
        setVisibleSections: (visibleSections) =>
            set((state) =>
                state.visibleSections.join() === visibleSections.join()
                    ? {}
                    : { visibleSections },
            ),
        registerHeading: ({ id, ref, offsetRem }) =>
            set((state) => {
                return {
                    sections: state.sections.map((section) => {
                        if (section.id === id) {
                            return {
                                ...section,
                                headingRef: ref,
                                offsetRem,
                            };
                        }
                        return section;
                    }),
                };
            }),
    }));
}

function useVisibleSections(sectionStore: StoreApi<SectionState>) {
    const setVisibleSections = useStore(
        sectionStore,
        (s) => s.setVisibleSections,
    );
    const sections = useStore(sectionStore, (s) => s.sections);

    useEffect(() => {
        function checkVisibleSections() {
            const { innerHeight, scrollY } = window;
            const newVisibleSections: string[] = [];

            for (
                let sectionIndex = 0;
                sectionIndex < sections.length;
                sectionIndex++
            ) {
                const {
                    id,
                    headingRef,
                    offsetRem = 0,
                } = sections[sectionIndex];

                if (!headingRef?.current) {
                    continue;
                }

                const offset = remToPx(offsetRem);
                const top =
                    headingRef.current.getBoundingClientRect().top + scrollY;

                if (sectionIndex === 0 && top - offset > scrollY) {
                    newVisibleSections.push("_top");
                }

                const nextSection = sections[sectionIndex + 1];
                const bottom =
                    (nextSection?.headingRef?.current?.getBoundingClientRect()
                        .top ?? Number.POSITIVE_INFINITY) +
                    scrollY -
                    remToPx(nextSection?.offsetRem ?? 0);

                if (
                    (top > scrollY && top < scrollY + innerHeight) ||
                    (bottom > scrollY && bottom < scrollY + innerHeight) ||
                    (top <= scrollY && bottom >= scrollY + innerHeight)
                ) {
                    newVisibleSections.push(id);
                }
            }

            setVisibleSections(newVisibleSections);
        }

        const raf = window.requestAnimationFrame(() => checkVisibleSections());
        window.addEventListener("scroll", checkVisibleSections, {
            passive: true,
        });
        window.addEventListener("resize", checkVisibleSections);

        return () => {
            window.cancelAnimationFrame(raf);
            window.removeEventListener("scroll", checkVisibleSections);
            window.removeEventListener("resize", checkVisibleSections);
        };
    }, [setVisibleSections, sections]);
}

const SectionStoreContext = createContext<StoreApi<SectionState> | null>(null);

const useIsomorphicLayoutEffect =
    typeof window === "undefined" ? useEffect : useLayoutEffect;

export function SectionProvider({
    sections,
    children,
}: {
    sections: Section[];
    children: ReactNode;
}) {
    const [sectionStore] = useState(() => createSectionStore(sections));

    useVisibleSections(sectionStore);

    useIsomorphicLayoutEffect(() => {
        sectionStore.setState({ sections });
    }, [sectionStore, sections]);

    return (
        <SectionStoreContext.Provider value={sectionStore}>
            {children}
        </SectionStoreContext.Provider>
    );
}

export function useSectionStore<T>(selector: (state: SectionState) => T) {
    const store = useContext(SectionStoreContext);
    return useStore(store as NonNullable<typeof store>, selector);
}