// Adapted from Kibo UI (MIT), with owned React/CSS APIs. See THIRD_PARTY_NOTICES.md.
"use client";
import { useId, useRef, useState, type ComponentProps } from "react";
export interface TreeItem {
  id: string;
  label: string;
  children?: readonly TreeItem[];
  disabled?: boolean;
}
export type TreeProps = Omit<
  ComponentProps<"ul">,
  "onSelect" | "defaultValue"
> & {
  items: readonly TreeItem[];
  label: string;
  value?: string;
  defaultValue?: string;
  onValueChange?: (id: string) => void;
  defaultExpanded?: readonly string[];
};
export function Tree({
  items,
  label,
  value,
  defaultValue,
  onValueChange,
  defaultExpanded = [],
  className = "",
  ...props
}: TreeProps) {
  const ref = useRef<HTMLUListElement>(null);
  const treeId = useId();
  const [expanded, setExpanded] = useState(new Set(defaultExpanded));
  const [selected, setSelected] = useState(defaultValue);
  const [focused, setFocused] = useState<string>();
  const visible: { item: TreeItem; parent?: string }[] = [];
  const collect = (nodes: readonly TreeItem[], parent?: string) => {
    for (const item of nodes) {
      visible.push({ item, parent });
      if (expanded.has(item.id) && item.children)
        collect(item.children, item.id);
    }
  };
  collect(items);
  const current = value ?? selected;
  const focusable = visible.some((n) => n.item.id === focused)
    ? focused
    : visible.find((n) => !n.item.disabled)?.item.id;
  const focus = (id: string | undefined) => {
    if (!id) return;
    setFocused(id);
    requestAnimationFrame(() =>
      Array.from(
        ref.current?.querySelectorAll<HTMLElement>('[role="treeitem"]') ?? [],
      )
        .find((el) => el.dataset.id === id)
        ?.focus(),
    );
  };
  const toggle = (id: string, open: boolean) =>
    setExpanded((old) => {
      const next = new Set(old);
      if (open) next.add(id);
      else next.delete(id);
      return next;
    });
  const render = (nodes: readonly TreeItem[], level = 1) =>
    nodes.map((item, index) => (
      <li role="none" key={item.id}>
        <div
          role="treeitem"
          data-id={item.id}
          aria-level={level}
          aria-setsize={nodes.length}
          aria-posinset={index + 1}
          aria-owns={
            item.children?.length && expanded.has(item.id)
              ? `${treeId}-${encodeURIComponent(item.id)}`
              : undefined
          }
          aria-label={item.label}
          aria-selected={current === item.id}
          aria-disabled={item.disabled}
          aria-expanded={
            item.children?.length ? expanded.has(item.id) : undefined
          }
          tabIndex={focusable === item.id ? 0 : -1}
          onFocus={() => setFocused(item.id)}
          onClick={() => {
            focus(item.id);
            if (item.disabled) return;
            if (value === undefined) setSelected(item.id);
            onValueChange?.(item.id);
            if (item.children?.length) toggle(item.id, !expanded.has(item.id));
          }}
          onKeyDown={(event) => {
            const index = visible.findIndex((n) => n.item.id === item.id);
            if (
              [
                "ArrowDown",
                "ArrowUp",
                "ArrowLeft",
                "ArrowRight",
                "Home",
                "End",
                "Enter",
                " ",
              ].includes(event.key)
            )
              event.preventDefault();
            if (event.key === "ArrowDown") focus(visible[index + 1]?.item.id);
            if (event.key === "ArrowUp") focus(visible[index - 1]?.item.id);
            if (event.key === "Home") focus(visible[0]?.item.id);
            if (event.key === "End") focus(visible.at(-1)?.item.id);
            if (event.key === "ArrowRight" && item.children?.length) {
              if (expanded.has(item.id)) focus(item.children[0]?.id);
              else toggle(item.id, true);
            }
            if (event.key === "ArrowLeft") {
              if (expanded.has(item.id)) toggle(item.id, false);
              else focus(visible[index]?.parent);
            }
            if (
              (event.key === "Enter" || event.key === " ") &&
              !item.disabled
            ) {
              if (value === undefined) setSelected(item.id);
              onValueChange?.(item.id);
            }
            if (
              event.key.length === 1 &&
              event.key !== " " &&
              !event.ctrlKey &&
              !event.metaKey
            ) {
              const ordered = [
                ...visible.slice(index + 1),
                ...visible.slice(0, index + 1),
              ];
              focus(
                ordered.find((n) =>
                  n.item.label
                    .toLowerCase()
                    .startsWith(event.key.toLowerCase()),
                )?.item.id,
              );
            }
          }}
        >
          <span aria-hidden="true">
            {item.children?.length ? (expanded.has(item.id) ? "▾" : "▸") : "·"}
          </span>
          {item.label}
        </div>
        {item.children?.length && expanded.has(item.id) ? (
          <ul role="group" id={`${treeId}-${encodeURIComponent(item.id)}`}>
            {render(item.children, level + 1)}
          </ul>
        ) : null}
      </li>
    ));
  return (
    <ul
      role="tree"
      aria-label={label}
      ref={ref}
      className={`ynr-tree ${className}`}
      {...props}
    >
      {render(items)}
    </ul>
  );
}