Stack

'use client';

import {
  type HTMLAttributes,
  type ReactElement,
  type ReactNode,
  useEffect,
  useRef,
} from 'react';

import { debounce } from '@/lib/debounce';
import { sum } from '@/lib/math/sum';

const px = (value: number) => (Number.isNaN(value) ? undefined : `${value}px`);

interface StackProps extends HTMLAttributes<HTMLDivElement> {
  stick: 'top' | 'bottom';
  children: ReactElement<typeof StackItem> | ReactElement<typeof StackItem>[];
}

const Stack = ({ stick = 'top', children, ...rest }: StackProps) => {
  const rootRef = useRef<HTMLDivElement>(null);

  useEffect(() => {
    const root = rootRef.current;
    if (!root) return;

    const items = [
      ...root.querySelectorAll('[data-stack="item"]'),
    ] as HTMLElement[];

    const headers = items.map((item) =>
      item.querySelector('[data-stack="header"]')
    ) as HTMLElement[];

    const length = items.length;

    const heights = {
      viewport: 0,
      items: new Array(length).fill(0),
      headers: new Array(length).fill(0),
      contents: new Array(length).fill(0),
      update: () => {
        heights.viewport = window.innerHeight;

        for (let index = 0; index < length; index++) {
          heights.items[index] = items[index].getBoundingClientRect().height;
          heights.headers[index] =
            headers[index]?.getBoundingClientRect().height || 0;
          heights.contents[index] =
            heights.items[index] - heights.headers[index];
        }
      },
    };

    const computeAlignTopBottomMargins = () => {
      const values: number[] = [];

      for (let i = length - 1; i >= 0; i--) {
        values[i] =
          i === length - 1
            ? 0
            : heights.items[i + 1] -
              heights.items[i] +
              heights.headers[i] +
              values[i + 1];
      }

      return values;
    };

    const onResize = debounce(() => {
      heights.update();

      const preComputedMarginBottoms =
        stick === 'top' ? computeAlignTopBottomMargins() : [];

      for (let index = 0; index < length; index++) {
        const styles = {
          position: 'sticky' as const,
          top: undefined as string | undefined,
          bottom: undefined as string | undefined,
          marginTop: undefined as string | undefined,
          marginBottom: undefined as string | undefined,
        };

        // sum of all headers before the current item
        const preHeaderHeightSum = sum(...heights.headers.slice(0, index));

        if (stick === 'top') {
          styles.top = px(preHeaderHeightSum);
          styles.marginBottom = px(preComputedMarginBottoms[index]);
          styles.marginTop = px(
            -1 * (preComputedMarginBottoms[index - 1] || 0)
          );
        }

        if (stick === 'bottom') {
          // sum of all headers after the current item
          const subHeaderHeightSum = sum(
            ...heights.headers.slice(index + 1, length)
          );

          styles.bottom = px(
            -1 * (heights.contents[index] - subHeaderHeightSum)
          );
          styles.marginTop = px(preHeaderHeightSum);
          styles.marginBottom = px(
            -1 * (preHeaderHeightSum + heights.headers[index])
          );
        }

        Object.assign(items[index].style, styles);
      }

      if (stick === 'bottom') {
        // correct layout deficit create by negative margins
        root.style.paddingBottom = px(sum(...heights.headers)) as string;
      }
    });

    const resizeObserver = new ResizeObserver(onResize);
    resizeObserver.observe(root);

    return () => {
      resizeObserver.disconnect();

      // clear styles
      root.style.removeProperty('paddingBottom');
      items.forEach((item) => {
        item.style.removeProperty('position');
        item.style.removeProperty('bottom');
        item.style.removeProperty('top');
        item.style.removeProperty('marginTop');
        item.style.removeProperty('marginBottom');
      });
    };
  }, [stick]);

  return (
    <div {...rest} ref={rootRef} data-stack="root">
      {children}
    </div>
  );
};

interface StackItemProps extends HTMLAttributes<HTMLDivElement> {
  children: [ReactElement<typeof StackHeader>, ...ReactNode[]];
}

const StackItem = ({ children, ...rest }: StackItemProps) => {
  return (
    <div {...rest} data-stack="item">
      {children}
    </div>
  );
};

const StackHeader = ({ children, ...rest }: HTMLAttributes<HTMLDivElement>) => {
  return (
    <div {...rest} data-stack="header">
      {children}
    </div>
  );
};

export { Stack, StackHeader, StackItem };