// Copyright (c) Meta Platforms, Inc. and affiliates.

'use client';

/**
 * @file useTableTreeState.tsx
 * @input React
 * @output Exports useTableTreeState hook and config/result types
 * @position Tree state helper; owns the expanded set + flattens nested data.
 *   Pairs with useTableTreeData (the headless tree plugin).
 *
 * Modeled after useTableSortableState — a convenience layer that owns state
 * (controlled or uncontrolled) and produces both the visible flattened rows
 * and a ready-to-use config for the headless plugin. Data shaping stays
 * outside the render pipeline: Table always receives exactly the rows it
 * renders, so collapsed subtrees are unmounted, not hidden.
 */
import { useCallback, useMemo, useRef, useState } from 'react';
import type { TableTreeRowMeta, UseTableTreeDataConfig } from "./useTableTreeData.js";

// =============================================================================
// Config
// =============================================================================

export interface UseTableTreeStateConfig<T extends Record<string, unknown>> {
  /** Nested data: rows may carry child rows under `childrenKey`. */
  data: T[];
  /** Row ID accessor: property name, or a function returning a unique id. */
  idKey: (keyof T & string) | ((item: T) => string | number);
  /**
   * Property holding each row's children array.
   * @default 'children'
   */
  childrenKey?: string;
  /** Initial expanded row ids for uncontrolled mode. Ignored when `expandedIds` is provided. */
  defaultExpandedIds?: Iterable<string>;
  /**
   * Controlled set of expanded row ids. When provided, the hook uses this
   * instead of internal state. Pair with `onExpandedIdsChange`.
   */
  expandedIds?: ReadonlySet<string>;
  /** Called with the next expanded set whenever expansion changes. */
  onExpandedIdsChange?: (ids: ReadonlySet<string>) => void;
  /**
   * Should this row show an expander? Overrides the default
   * "has a non-empty children array" check — use for lazy loading, where a
   * row is expandable before its children have been fetched.
   */
  isItemExpandable?: (item: T) => boolean;
  /**
   * Sort each sibling group independently during flattening. Children always
   * stay directly under their parent — sorting never crosses levels. Pass
   * `applySort` from `useTableSortableState` to compose with column sorting.
   */
  sortSiblings?: (siblings: T[]) => T[];
  /** Indent step per level, forwarded to useTableTreeData. */
  indent?: 'sm' | 'md' | 'lg';
  /** Column that carries the tree affordance, forwarded to useTableTreeData. */
  treeColumnKey?: string;
}

// =============================================================================
// Result
// =============================================================================

export interface UseTableTreeStateResult<T extends Record<string, unknown>> {
  /** The flattened, currently-visible rows. Pass to `<Table data>`. */
  visibleData: T[];
  /** Ready-to-use config for useTableTreeData. */
  treeConfig: UseTableTreeDataConfig<T>;
  /** The current expanded row ids. */
  expandedIds: ReadonlySet<string>;
  /**
   * Aggregate expansion state across every expandable row: `true` when all are
   * expanded, `false` when none are, `'indeterminate'` when some but not all
   * are. Drives the header expand-all control. Reports `false` for flat data.
   */
  isAllExpanded: boolean | 'indeterminate';
  /** Expand every expandable row in the tree. */
  expandAll: () => void;
  /** Collapse every row. */
  collapseAll: () => void;
}

// =============================================================================
// Hook
// =============================================================================

export function useTableTreeState<T extends Record<string, unknown>>(config: UseTableTreeStateConfig<T>): UseTableTreeStateResult<T> {
  const {
    data,
    idKey,
    childrenKey = 'children',
    defaultExpandedIds,
    expandedIds: controlledExpandedIds,
    onExpandedIdsChange,
    isItemExpandable,
    sortSiblings,
    indent,
    treeColumnKey
  } = config;
  const getId = useCallback((item: T): string => typeof idKey === 'function' ? String(idKey(item)) : String(item[idKey]), [idKey]);
  const getChildren = useCallback((item: T): T[] => {
    const children = item[childrenKey];
    return Array.isArray(children) ? children as T[] : [];
  }, [childrenKey]);
  const getIsExpandable = useCallback((item: T): boolean => isItemExpandable ? isItemExpandable(item) : getChildren(item).length > 0, [isItemExpandable, getChildren]);

  // Internal state (used in uncontrolled mode)
  const [internalExpandedIds, setInternalExpandedIds] = useState<ReadonlySet<string>>(() => new Set(defaultExpandedIds));

  // Resolve controlled vs uncontrolled
  const isControlled = controlledExpandedIds !== undefined;
  const expandedIds = isControlled ? controlledExpandedIds : internalExpandedIds;

  // Stable refs so callbacks don't churn on inline-prop identity changes
  const expandedIdsRef = useRef(expandedIds);
  expandedIdsRef.current = expandedIds;
  const onExpandedIdsChangeRef = useRef(onExpandedIdsChange);
  onExpandedIdsChangeRef.current = onExpandedIdsChange;
  const commitExpandedIds = useCallback((next: ReadonlySet<string>) => {
    // Advance the ref immediately so a second commit in the same React
    // batch (e.g. two onToggleItem calls in one event handler) builds
    // on this one instead of the pre-batch set. The ref re-syncs from
    // the resolved state on the next render either way.
    expandedIdsRef.current = next;
    if (!isControlled) {
      setInternalExpandedIds(next);
    }
    onExpandedIdsChangeRef.current?.(next);
  }, [isControlled]);

  // Flatten: depth-first walk emitting only visible rows, collecting per-row
  // meta ({level, hasChildren, isExpanded}) in the same pass. Children mount
  // only when their parent is expanded, so collapsed subtrees stay unmounted.
  // `path` holds the ids on the current ancestor chain: an edge pointing back
  // at an ancestor is skipped instead of recursing forever. Sibling arrays
  // are copied before sortSiblings so an in-place sorter can't reorder the
  // consumer's data.
  const {
    visibleData,
    metaMap
  } = useMemo(() => {
    const rows: T[] = [];
    const meta = new Map<string, TableTreeRowMeta>();
    const path = new Set<string>();
    const walk = (items: T[], level: number) => {
      const siblings = sortSiblings ? sortSiblings([...items]) : items;
      for (const item of siblings) {
        const id = getId(item);
        if (path.has(id)) {
          continue; // cyclic edge — this row is its own ancestor
        }
        const hasChildren = getIsExpandable(item);
        const isExpanded = hasChildren && expandedIds.has(id);
        rows.push(item);
        meta.set(id, {
          id,
          level,
          hasChildren,
          isExpanded
        });
        if (isExpanded) {
          path.add(id);
          walk(getChildren(item), level + 1);
          path.delete(id);
        }
      }
    };
    walk(data, 0);
    return {
      visibleData: rows,
      metaMap: meta
    };
  }, [data, expandedIds, getId, getChildren, getIsExpandable, sortSiblings]);

  // Every expandable row id across the whole tree, including collapsed
  // subtrees (drives expandAll and the migration no-op flag). Same
  // ancestor-chain guard as the flatten walk.
  const allExpandableIds = useMemo(() => {
    const idsAcc: string[] = [];
    const path = new Set<string>();
    const walk = (items: T[]) => {
      for (const item of items) {
        const id = getId(item);
        if (path.has(id)) {
          continue;
        }
        if (getIsExpandable(item)) {
          idsAcc.push(id);
        }
        path.add(id);
        walk(getChildren(item));
        path.delete(id);
      }
    };
    walk(data);
    return idsAcc;
  }, [data, getId, getChildren, getIsExpandable]);
  const hasExpandableRows = allExpandableIds.length > 0;

  // Aggregate state for the header expand-all control. Only expandable ids
  // count toward the total, so expanded ids that match a leaf or no row at
  // all never skew the tally.
  const isAllExpanded: boolean | 'indeterminate' = useMemo(() => {
    if (!hasExpandableRows) {
      return false;
    }
    const expandedCount = allExpandableIds.filter(id => expandedIds.has(id)).length;
    if (expandedCount === 0) {
      return false;
    }
    if (expandedCount === allExpandableIds.length) {
      return true;
    }
    return 'indeterminate';
  }, [hasExpandableRows, allExpandableIds, expandedIds]);
  const onToggleItem = useCallback((item: T) => {
    const id = getId(item);
    const next = new Set(expandedIdsRef.current);
    if (next.has(id)) {
      next.delete(id);
    } else {
      next.add(id);
    }
    commitExpandedIds(next);
  }, [getId, commitExpandedIds]);
  const expandAll = useCallback(() => {
    commitExpandedIds(new Set(allExpandableIds));
  }, [commitExpandedIds, allExpandableIds]);
  const collapseAll = useCallback(() => {
    commitExpandedIds(new Set());
  }, [commitExpandedIds]);
  const getRowMeta = useCallback((item: T): TableTreeRowMeta | undefined => metaMap.get(getId(item)), [metaMap, getId]);
  const treeConfig = useMemo((): UseTableTreeDataConfig<T> => ({
    getRowMeta,
    onToggleItem,
    hasExpandableRows,
    isAllExpanded,
    onExpandAll: expandAll,
    onCollapseAll: collapseAll,
    indent,
    treeColumnKey
  }), [getRowMeta, onToggleItem, hasExpandableRows, isAllExpanded, expandAll, collapseAll, indent, treeColumnKey]);
  return {
    visibleData,
    treeConfig,
    expandedIds,
    isAllExpanded,
    expandAll,
    collapseAll
  };
}