'use client';

import * as React from 'react';

import { CheckIcon, ChevronRightIcon } from '@/icons';
import { cn } from '@/lib/utils';
import { focusRing } from '@/lib/cva-presets';
import { useControllableState } from '@/hooks/use-controllable-state';
import { branchKeys, flatten, indexTree, isBranch, type FlatNode, type TreeNode } from './tree-nodes';
import { useTreeKeyboard } from './use-tree-keyboard';

export type { TreeNode };

export interface TreeCheckInfo {
  node: TreeNode;
  checked: boolean;
  /** Branches with some, but not all, of their descendants checked. */
  halfCheckedKeys: React.Key[];
}

export interface TreeProps extends Omit<React.ComponentProps<'div'>, 'onSelect'> {
  treeData: TreeNode[];
  /** Show checkboxes, with checked state propagating up and down. */
  checkable?: boolean;
  expandedKeys?: readonly React.Key[];
  defaultExpandedKeys?: readonly React.Key[];
  defaultExpandAll?: boolean;
  onExpand?: (keys: React.Key[], info: { node: TreeNode; expanded: boolean }) => void;
  selectedKeys?: readonly React.Key[];
  defaultSelectedKeys?: readonly React.Key[];
  onSelect?: (keys: React.Key[], info: { node: TreeNode; selected: boolean }) => void;
  checkedKeys?: readonly React.Key[];
  defaultCheckedKeys?: readonly React.Key[];
  onCheck?: (keys: React.Key[], info: TreeCheckInfo) => void;
  /** Window the rows. Needs `height`; worth it past a few hundred nodes. */
  virtual?: boolean;
  height?: number;
  /** Row height, only used when `virtual` is set. */
  itemHeight?: number;
  /** Override the expand/collapse affordance. */
  switcherIcon?: (node: { isLeaf: boolean; expanded: boolean }) => React.ReactNode;
  'aria-label'?: string;
}

/**
 * Hierarchical list with expand, select and optional checkboxes.
 *
 * The keyboard contract is the demanding part and is implemented in full —
 * arrows that move *and* expand, Home/End, `*` to open a whole level,
 * type-ahead, and a single tab stop so Tab leaves the tree rather than walking
 * every node in it. It lives in `./use-tree-keyboard`; the walks over the data
 * live in `./tree-nodes`.
 *
 * Checkbox state propagates both ways — checking a branch checks everything
 * under it, and a branch reads as indeterminate when only some of its
 * descendants are checked. `onCheck` reports the checked **leaves** plus the
 * half-checked branch keys, which is the shape a server usually wants.
 *
 * ```tsx
 * <Tree
 *   checkable
 *   defaultExpandAll
 *   treeData={[{ key: 'src', title: 'src', children: [{ key: 'app', title: 'app' }] }]}
 * />
 * ```
 */
function Tree({
  treeData,
  checkable,
  expandedKeys,
  defaultExpandedKeys,
  defaultExpandAll,
  onExpand,
  selectedKeys,
  defaultSelectedKeys,
  onSelect,
  checkedKeys,
  defaultCheckedKeys,
  onCheck,
  virtual,
  height,
  itemHeight = 28,
  switcherIcon,
  className,
  ...props
}: TreeProps) {
  const { byKey, leavesOf } = React.useMemo(() => indexTree(treeData), [treeData]);

  /* Each controlled prop arrives as an array and is used as a set. Building
     that set inside a memo matters here: without one it is a fresh object every
     render, and everything downstream — the flattened rows, the check-state
     lookup — recomputes for a tree that may hold thousands of nodes.
     `undefined` is what tells the hook the caller is not controlling it, so
     these three memos are the whole conversion.
     None of the three passes `onChange`: this component's callbacks report the
     keys *and* the node the change was made on, which is a signature the hook
     does not have. Each handler calls its own below. */
  const controlledExpanded = React.useMemo(
    () => (expandedKeys ? new Set(expandedKeys) : undefined),
    [expandedKeys]
  );
  const [expanded, setExpanded] = useControllableState<Set<React.Key>>({
    value: controlledExpanded,
    defaultValue: () =>
      new Set(defaultExpandAll ? branchKeys(treeData) : (defaultExpandedKeys ?? [])),
  });

  const controlledSelected = React.useMemo(
    () => (selectedKeys ? new Set(selectedKeys) : undefined),
    [selectedKeys]
  );
  const [selected, setSelected] = useControllableState<Set<React.Key>>({
    value: controlledSelected,
    defaultValue: () => new Set(defaultSelectedKeys ?? []),
  });

  const controlledChecked = React.useMemo(
    () => (checkedKeys ? new Set(checkedKeys) : undefined),
    [checkedKeys]
  );
  const [checked, setChecked] = useControllableState<Set<React.Key>>({
    value: controlledChecked,
    defaultValue: () => new Set(defaultCheckedKeys ?? []),
  });

  const rows = React.useMemo(() => flatten(treeData, expanded), [treeData, expanded]);

  /* One tab stop for the whole tree. Which node the arrows act on is tracked
     here — the roving-tabindex pattern every tree widget uses. */
  const [activeKey, setActiveKey] = React.useState<React.Key | undefined>(() => treeData[0]?.key);
  const active = rows.some((row) => row.node.key === activeKey) ? activeKey : rows[0]?.node.key;

  const treeRef = React.useRef<HTMLDivElement>(null);

  /**
   * A branch's checked state is derived from its leaves rather than stored.
   * Storing it would mean keeping every ancestor in sync on each toggle, and
   * any missed path shows up as a checkbox that lies about its subtree.
   *
   * The leaves themselves come from the index, not from a fresh walk — this
   * runs for every visible row on every render.
   */
  const checkStateOf = React.useCallback(
    (node: TreeNode): 'checked' | 'indeterminate' | 'unchecked' => {
      const leaves = leavesOf.get(node.key) ?? [];
      const hits = leaves.filter((key) => checked.has(key)).length;

      if (hits === 0) return 'unchecked';
      return hits === leaves.length ? 'checked' : 'indeterminate';
    },
    [checked, leavesOf]
  );

  const setExpandedKey = (key: React.Key, open: boolean) => {
    const next = new Set(expanded);
    if (open) next.add(key);
    else next.delete(key);

    setExpanded(next);
    onExpand?.([...next], { node: byKey.get(key)!, expanded: open });
  };

  const selectKey = (key: React.Key) => {
    const node = byKey.get(key)!;
    if (node.disabled) return;

    const next = new Set<React.Key>(selected.has(key) ? [] : [key]);
    setSelected(next);
    onSelect?.([...next], { node, selected: next.has(key) });
  };

  const toggleCheck = (key: React.Key) => {
    const node = byKey.get(key)!;
    if (node.disabled || node.disableCheckbox) return;

    const next = new Set(checked);
    const turningOn = checkStateOf(node) !== 'checked';

    /* Only leaves are stored. A branch counts as checked when all of its leaves
       are, so writing the leaves *is* the whole update — ancestors follow. */
    for (const leaf of leavesOf.get(node.key) ?? []) {
      const leafNode = byKey.get(leaf);
      if (leafNode?.disabled || leafNode?.disableCheckbox) continue;
      if (turningOn) next.add(leaf);
      else next.delete(leaf);
    }

    setChecked(next);

    const halfCheckedKeys = [...byKey.values()]
      .filter((candidate) => {
        if (!isBranch(candidate)) return false;
        const leaves = leavesOf.get(candidate.key) ?? [];
        const hits = leaves.filter((leaf) => next.has(leaf)).length;
        return hits > 0 && hits < leaves.length;
      })
      .map((candidate) => candidate.key);

    onCheck?.([...next], { node, checked: turningOn, halfCheckedKeys });
  };

  const focusRow = (key: React.Key) => {
    setActiveKey(key);
    treeRef.current
      ?.querySelector<HTMLElement>(`[data-tree-key="${CSS.escape(String(key))}"]`)
      ?.focus();
  };

  const handleKeyDown = useTreeKeyboard({
    rows,
    active,
    byKey,
    expanded,
    checkable,
    focusRow,
    setExpandedKey,
    setExpanded,
    onExpand,
    selectKey,
    toggleCheck,
  });

  const renderRow = ({ node, depth }: FlatNode) => {
    const branch = isBranch(node);
    const isOpen = expanded.has(node.key);
    const state = checkable ? checkStateOf(node) : 'unchecked';

    return (
      <div
        key={node.key}
        role="treeitem"
        data-tree-key={node.key}
        aria-expanded={branch ? isOpen : undefined}
        aria-selected={selected.has(node.key)}
        aria-checked={
          checkable ? (state === 'indeterminate' ? 'mixed' : state === 'checked') : undefined
        }
        aria-disabled={node.disabled || undefined}
        aria-level={depth + 1}
        tabIndex={node.key === active ? 0 : -1}
        data-slot="tree-item"
        onFocus={() => setActiveKey(node.key)}
        onClick={() => selectKey(node.key)}
        style={{ paddingInlineStart: `${depth * 1.25 + 0.25}rem` }}
        className={cn(
          'flex cursor-pointer items-center gap-1 rounded-md py-1 pe-2 text-sm',
          'transition-colors duration-(--ui-duration-fast) ease-(--ui-ease-standard)',
          focusRing,
          node.disabled
            ? 'cursor-not-allowed text-muted-foreground/50'
            : 'hover:bg-accent hover:text-accent-foreground',
          selected.has(node.key) && 'bg-accent text-accent-foreground'
        )}
      >
        <span
          data-slot="tree-switcher"
          aria-hidden
          onClick={(event) => {
            /* Toggling is not selecting — a click on the chevron must not also
               pick the node it belongs to. */
            event.stopPropagation();
            if (branch) setExpandedKey(node.key, !isOpen);
          }}
          className={cn(
            'inline-flex size-4 shrink-0 items-center justify-center',
            branch && 'cursor-pointer'
          )}
        >
          {branch
            ? (switcherIcon?.({ isLeaf: false, expanded: isOpen }) ?? (
                <ChevronRightIcon
                  className={cn(
                    'size-4 text-muted-foreground transition-transform duration-(--ui-duration-fast)',
                    isOpen && 'rotate-90'
                  )}
                />
              ))
            : switcherIcon?.({ isLeaf: true, expanded: false })}
        </span>

        {checkable ? (
          <span
            data-slot="tree-checkbox"
            aria-hidden
            onClick={(event) => {
              event.stopPropagation();
              toggleCheck(node.key);
            }}
            className={cn(
              'inline-flex size-4 shrink-0 items-center justify-center rounded-sm border',
              state === 'unchecked'
                ? 'border-input'
                : 'border-primary bg-primary text-primary-foreground',
              (node.disabled || node.disableCheckbox) && 'opacity-50'
            )}
          >
            {state === 'checked' ? <CheckIcon className="size-3" /> : null}
            {state === 'indeterminate' ? (
              <span className="h-0.5 w-2 rounded-full bg-primary-foreground" />
            ) : null}
          </span>
        ) : null}

        <span className="truncate">{node.title}</span>
      </div>
    );
  };

  return (
    <div
      ref={treeRef}
      role="tree"
      aria-multiselectable={checkable || undefined}
      data-slot="tree"
      onKeyDown={handleKeyDown}
      className={cn('text-sm', className)}
      {...props}
    >
      {virtual && height ? (
        <VirtualRows rows={rows} height={height} itemHeight={itemHeight} render={renderRow} />
      ) : (
        rows.map(renderRow)
      )}
    </div>
  );
}

/** Windowed rows, using the same fixed-height technique as `VirtualList`. */
function VirtualRows({
  rows,
  height,
  itemHeight,
  render,
}: {
  rows: FlatNode[];
  height: number;
  itemHeight: number;
  render: (row: FlatNode) => React.ReactNode;
}) {
  const [scrollTop, setScrollTop] = React.useState(0);

  const first = Math.max(0, Math.floor(scrollTop / itemHeight) - 3);
  const last = Math.min(rows.length, Math.floor((scrollTop + height - 1) / itemHeight) + 4);

  return (
    <div
      style={{ height }}
      className="overflow-y-auto"
      onScroll={(event) => setScrollTop(event.currentTarget.scrollTop)}
    >
      <div style={{ height: rows.length * itemHeight }} className="relative">
        <div style={{ transform: `translateY(${first * itemHeight}px)` }}>
          {rows.slice(first, last).map(render)}
        </div>
      </div>
    </div>
  );
}

export { Tree };