Masonry Grid

"use client";

import * as React from "react";
import { AnimatePresence, motion } from "motion/react";

type Spring = { stiffness: number; damping: number; mass: number };

/** A fixed column count, or breakpoints keyed by minimum container width in px: { 0: 1, 480: 2, 720: 3 }. */
export type MasonryColumns = number | Record<number, number>;

export interface MasonryGridProps
  extends Omit<
    React.HTMLAttributes<HTMLElement>,
    "children" | "onDrag" | "onDragStart" | "onDragEnd" | "onAnimationStart" | "onAnimationEnd"
  > {
  /** Keyed elements, one per tile. Keys must be stable: they identify tiles across reorders. */
  children: React.ReactNode;
  /** Column count, or breakpoints by container width. */
  columns?: MasonryColumns;
  /** Space between tiles, in px (both axes). */
  gap?: number;
  /** Animate reflow, entering and leaving tiles. */
  animate?: boolean;
  /** Spring for tiles moving to new positions. */
  spring?: Spring;
  /** Render as a list (ul > li) for collections of similar items. */
  as?: "div" | "ul";
}

function subscribeReducedMotion(callback: () => void) {
  const mq = window.matchMedia("(prefers-reduced-motion: reduce)");
  mq.addEventListener("change", callback);
  return () => mq.removeEventListener("change", callback);
}

/** SSR-safe: hydrates as "no preference", then follows the media query. */
function usePrefersReducedMotion() {
  return React.useSyncExternalStore(
    subscribeReducedMotion,
    () => window.matchMedia("(prefers-reduced-motion: reduce)").matches,
    () => false,
  );
}

const DEFAULT_COLUMNS: Record<number, number> = { 0: 1, 480: 2, 720: 3, 960: 4 };

function columnCount(columns: MasonryColumns, width: number) {
  if (typeof columns === "number") return Math.max(1, Math.floor(columns));
  let count = 1;
  for (const [min, n] of Object.entries(columns).sort((a, b) => Number(a[0]) - Number(b[0]))) {
    if (width >= Number(min)) count = n;
  }
  return Math.max(1, Math.floor(count));
}

const round = (n: number) => Math.round(n * 100) / 100;

/**
 * A masonry layout measured in JavaScript: each tile goes into the currently shortest column, so the
 * placement is stable and DOM order stays the reading order. Tiles are absolutely positioned and spring
 * to new places when tiles are added, removed or change height (ResizeObserver, e.g. once an image loads).
 */
export function MasonryGrid({
  children,
  columns = DEFAULT_COLUMNS,
  gap = 16,
  animate = true,
  spring = { stiffness: 320, damping: 32, mass: 1 },
  as = "div",
  className = "",
  style,
  ...props
}: MasonryGridProps) {
  const reduceMotion = usePrefersReducedMotion();
  const containerRef = React.useRef<HTMLElement>(null);
  const tileRefs = React.useRef(new Map<string, HTMLElement>());
  const [width, setWidth] = React.useState(0);
  const [heights, setHeights] = React.useState<Record<string, number>>({});
  const [settled, setSettled] = React.useState(false);

  const tiles = React.Children.toArray(children).filter(React.isValidElement);
  const keys = tiles.map((tile, i) => String(tile.key ?? i));
  const keySignature = keys.join("|");

  const count = width > 0 ? columnCount(columns, width) : 1;
  const columnWidth = width > 0 ? (width - gap * (count - 1)) / count : 0;

  // Shortest-column placement in DOM order. Ties go to the leftmost column, so results are deterministic.
  const columnHeights = new Array<number>(count).fill(0);
  const positions = new Map<string, { x: number; y: number }>();
  for (const key of keys) {
    const h = heights[key];
    if (h === undefined) continue;
    let col = 0;
    for (let c = 1; c < count; c++) if (columnHeights[c] < columnHeights[col] - 0.5) col = c;
    positions.set(key, { x: round(col * (columnWidth + gap)), y: round(columnHeights[col]) });
    columnHeights[col] += h + gap;
  }
  const totalHeight = Math.max(0, Math.max(...columnHeights) - gap);
  const ready = width > 0;

  const readHeights = React.useCallback(() => {
    setHeights((prev) => {
      let changed = false;
      const next: Record<string, number> = {};
      tileRefs.current.forEach((el, key) => {
        const h = el.offsetHeight;
        next[key] = h;
        if (prev[key] !== h) changed = true;
      });
      if (Object.keys(prev).length !== Object.keys(next).length) changed = true;
      return changed ? next : prev;
    });
  }, []);

  // Container width drives the column count.
  React.useEffect(() => {
    const el = containerRef.current;
    if (!el) return;
    const ro = new ResizeObserver(([entry]) => {
      const w = Math.round(entry.contentRect.width * 100) / 100;
      setWidth((prev) => (Math.abs(prev - w) < 0.5 ? prev : w));
    });
    ro.observe(el);
    return () => ro.disconnect();
  }, []);

  // Tile heights change with column width, content and late-loading images.
  React.useEffect(() => {
    const ro = new ResizeObserver(() => readHeights());
    tileRefs.current.forEach((el) => ro.observe(el));
    return () => ro.disconnect();
  }, [keySignature, readHeights]);

  // Measure synchronously after a width or tile change so the first placement is painted in place.
  React.useLayoutEffect(() => {
    if (!ready) return;
    readHeights();
    // Layout can only be measured in the browser; enabling transitions after the first placement keeps it from animating in.
    // eslint-disable-next-line react-hooks/set-state-in-effect
    if (!settled) setSettled(true);
  }, [ready, columnWidth, keySignature, readHeights, settled]);

  const moving = animate && !reduceMotion && settled;
  const transition = moving ? { type: "spring" as const, ...spring } : { duration: 0 };
  const Root = as;
  const Tile = as === "ul" ? motion.li : motion.div;

  return (
    <Root
      ref={containerRef as React.Ref<HTMLDivElement & HTMLUListElement>}
      className={`relative w-full ${as === "ul" ? "list-none p-0" : ""} ${className}`}
      // Before the first measure (server render) tiles flow in CSS columns, so content is never hidden.
      style={ready ? { height: totalHeight, ...style } : { columnCount: typeof columns === "number" ? columns : 2, columnGap: gap, ...style }}
      {...props}
    >
      <AnimatePresence initial={false}>
        {tiles.map((tile, i) => {
          const key = keys[i];
          const pos = positions.get(key);
          return (
            <Tile
              key={key}
              ref={(node: HTMLElement | null) => {
                if (node) tileRefs.current.set(key, node);
                else tileRefs.current.delete(key);
              }}
              initial={moving ? { opacity: 0, scale: 0.94 } : false}
              animate={{ opacity: 1, scale: 1, x: pos?.x ?? 0, y: pos?.y ?? 0 }}
              exit={moving ? { opacity: 0, scale: 0.94, transition: { duration: 0.18 } } : { opacity: 0, transition: { duration: 0 } }}
              transition={transition}
              style={
                ready
                  ? { position: "absolute", top: 0, left: 0, width: columnWidth, visibility: pos ? "visible" : "hidden" }
                  : { breakInside: "avoid", marginBottom: gap }
              }
            >
              {tile}
            </Tile>
          );
        })}
      </AnimatePresence>
    </Root>
  );
}