Table of Contents

"use client";

import * as React from "react";

import { cn } from "@/lib/utils";

export type TocItem = {
  id: string;
  text: string;
  level: number;
  children?: TocItem[];
};

type ThrottledFn = (() => void) & { cancel: () => void };

function throttle(fn: () => void, limit: number): ThrottledFn {
  let lastRan = 0;
  let timer: ReturnType<typeof setTimeout> | null = null;

  const throttled = (() => {
    const now = Date.now();
    const remaining = limit - (now - lastRan);

    if (remaining <= 0) {
      if (timer) {
        clearTimeout(timer);
        timer = null;
      }
      lastRan = now;
      fn();
    } else if (!timer) {
      timer = setTimeout(() => {
        lastRan = Date.now();
        timer = null;
        fn();
      }, remaining);
    }
  }) as ThrottledFn;

  throttled.cancel = () => {
    if (timer) {
      clearTimeout(timer);
      timer = null;
    }
  };

  return throttled;
}

function flattenTocItems(items: TocItem[]): TocItem[] {
  const out: TocItem[] = [];
  const walk = (list: TocItem[]) => {
    for (const item of list) {
      out.push(item);
      if (item.children) walk(item.children);
    }
  };
  walk(items);
  return out;
}

function prefersReducedMotion() {
  if (typeof window === "undefined" || !window.matchMedia) return false;
  return window.matchMedia("(prefers-reduced-motion: reduce)").matches;
}

const TOP_OFFSET = 56;

function useActiveHeading(items: TocItem[]) {
  const [activeId, setActiveId] = React.useState<string | null>(null);

  React.useEffect(() => {
    const headings = flattenTocItems(items);

    const update = () => {
      if (window.scrollY === 0) {
        setActiveId(null);
        return;
      }

      const boxes = headings
        .map(({ id }) => {
          const el = document.getElementById(id);
          if (!el) return null;
          return { id, box: el.getBoundingClientRect() };
        })
        .filter(
          (entry): entry is { id: string; box: DOMRect } => entry !== null,
        );

      let current = boxes.find(
        ({ box }) => box.bottom > TOP_OFFSET && box.top < window.innerHeight,
      );

      if (!current) {
        current = [...boxes]
          .reverse()
          .find(({ box }) => box.bottom < TOP_OFFSET);
      }

      setActiveId(current ? current.id : null);
    };

    const onScroll = throttle(update, 200);

    update();
    window.addEventListener("scroll", onScroll, { passive: true });
    window.addEventListener("resize", onScroll, { passive: true });

    return () => {
      onScroll.cancel();
      window.removeEventListener("scroll", onScroll);
      window.removeEventListener("resize", onScroll);
    };
  }, [items]);

  return activeId;
}

type TocContextValue = {
  activeId: string | null;
};

const TocContext = React.createContext<TocContextValue | null>(null);

function useTocContext(component: string) {
  const ctx = React.useContext(TocContext);
  if (!ctx) {
    throw new Error(`${component} must be used within <TableOfContents>.`);
  }
  return ctx;
}

type TableOfContentsProps = Omit<React.ComponentProps<"nav">, "children"> & {
  items?: TocItem[];
  activeId?: string | null;
  label?: React.ReactNode;
  children?: React.ReactNode;
};

function TableOfContents({
  items,
  activeId: controlledActiveId,
  label = "On this page",
  className,
  children,
  "aria-label": ariaLabel,
  ...props
}: TableOfContentsProps) {
  const trackedActiveId = useActiveHeading(
    controlledActiveId ? [] : (items ?? []),
  );
  const activeId =
    controlledActiveId !== undefined ? controlledActiveId : trackedActiveId;

  const content =
    children ??
    (items && items.length > 0 ? (
      <>
        {label ? <TableOfContentsLabel>{label}</TableOfContentsLabel> : null}
        <TableOfContentsList items={items} />
      </>
    ) : null);

  if (!content) return null;

  return (
    <TocContext.Provider value={{ activeId }}>
      <nav
        data-slot="toc"
        aria-label={ariaLabel ?? "On this page"}
        className={cn("flex flex-col gap-3", className)}
        {...props}
      >
        {content}
      </nav>
    </TocContext.Provider>
  );
}

type TableOfContentsLabelProps = React.ComponentProps<"p">;

function TableOfContentsLabel({
  className,
  ...props
}: TableOfContentsLabelProps) {
  return (
    <p
      data-slot="toc-label"
      className={cn(
        "font-mono text-[10px] uppercase tracking-[0.14em] text-muted-foreground",
        className,
      )}
      {...props}
    />
  );
}

type TableOfContentsListProps = Omit<React.ComponentProps<"ul">, "children"> & {
  items?: TocItem[];
  level?: number;
  children?: React.ReactNode;
};

function TableOfContentsList({
  items,
  level = 0,
  className,
  children,
  ...props
}: TableOfContentsListProps) {
  if (children === undefined && (!items || items.length === 0)) return null;

  return (
    <ul
      data-slot="toc-list"
      data-level={level}
      className={cn(
        "flex flex-col gap-1",
        level === 0 ? "border-s border-border" : "ms-4 mt-1",
        className,
      )}
      {...props}
    >
      {children ??
        items?.map((item) => (
          <TableOfContentsItem key={item.id} item={item} level={level} />
        ))}
    </ul>
  );
}

type TableOfContentsItemProps = Omit<React.ComponentProps<"li">, "children"> & {
  item: TocItem;
  level?: number;
};

function TableOfContentsItem({
  item,
  level = 0,
  className,
  ...props
}: TableOfContentsItemProps) {
  return (
    <li data-slot="toc-item" className={className} {...props}>
      <TableOfContentsLink href={`#${item.id}`} level={item.level}>
        {item.text}
      </TableOfContentsLink>
      {item.children && item.children.length > 0 ? (
        <TableOfContentsList items={item.children} level={level + 1} />
      ) : null}
    </li>
  );
}

type TableOfContentsLinkProps = React.ComponentProps<"a"> & {
  level?: number;
};

function TableOfContentsLink({
  href = "",
  level = 2,
  className,
  onClick,
  children,
  ...props
}: TableOfContentsLinkProps) {
  const { activeId } = useTocContext("TableOfContentsLink");
  const id = href.startsWith("#") ? href.slice(1) : href;
  const isActive = id.length > 0 && activeId === id;

  const handleClick = (event: React.MouseEvent<HTMLAnchorElement>) => {
    onClick?.(event);
    if (event.defaultPrevented) return;
    if (!id || !href.startsWith("#")) return;

    const target = document.getElementById(id);
    if (!target) return;

    event.preventDefault();
    target.scrollIntoView({
      behavior: prefersReducedMotion() ? "auto" : "smooth",
    });
    window.history.replaceState(null, "", `#${id}`);
  };

  return (
    <a
      data-slot="toc-link"
      data-active={isActive ? "" : undefined}
      href={href}
      aria-current={isActive ? "true" : undefined}
      onClick={handleClick}
      className={cn(
        "-ms-px block border-s py-1 ps-4 text-sm transition-colors duration-150 ease-out",
        isActive
          ? "border-foreground font-medium text-foreground"
          : "border-transparent text-muted-foreground hover:text-foreground",
        level >= 4 && "text-xs",
        className,
      )}
      {...props}
    >
      {children}
    </a>
  );
}

export {
  TableOfContents,
  TableOfContentsLabel,
  TableOfContentsList,
  TableOfContentsItem,
  TableOfContentsLink,
};