'use client';

import * as React from 'react';

import { CheckIcon, ChevronRightIcon } from '@/icons';
import { cn } from '@/lib/utils';
import { focusRing } from '@/lib/cva-presets';

export interface TreeNode {
  key: React.Key;
  title: React.ReactNode;
  children?: TreeNode[];
  disabled?: boolean;
  /** Exclude just the checkbox, leaving the node selectable. */
  disableCheckbox?: boolean;
  /** Force leaf rendering for a branch whose children load lazily. */
  isLeaf?: boolean;
}

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;
}

/** A node plus where it sits, which is what rendering and navigation both need. */
interface FlatNode {
  node: TreeNode;
  depth: number;
  parentKey: React.Key | null;
}

const isBranch = (node: TreeNode) => !node.isLeaf && !!node.children?.length;

/** Walks the tree, emitting only what is currently visible. */
const flatten = (
  nodes: TreeNode[],
  expanded: Set<React.Key>,
  depth = 0,
  parentKey: React.Key | null = null,
  out: FlatNode[] = []
) => {
  for (const node of nodes) {
    out.push({ node, depth, parentKey });
    if (isBranch(node) && expanded.has(node.key)) {
      flatten(node.children!, expanded, depth + 1, node.key, out);
    }
  }
  return out;
};

const branchKeys = (nodes: TreeNode[], out: React.Key[] = []) => {
  for (const node of nodes) {
    if (isBranch(node)) {
      out.push(node.key);
      branchKeys(node.children!, out);
    }
  }
  return out;
};

/** Index of every node by key, plus its parent, for navigation and propagation. */
const indexTree = (nodes: TreeNode[]) => {
  const byKey = new Map<React.Key, TreeNode>();
  const parentOf = new Map<React.Key, React.Key>();

  const walk = (list: TreeNode[], parent: React.Key | null) => {
    for (const node of list) {
      byKey.set(node.key, node);
      if (parent !== null) parentOf.set(node.key, parent);
      if (node.children?.length) walk(node.children, node.key);
    }
  };

  walk(nodes, null);
  return { byKey, parentOf };
};

const leafKeysUnder = (node: TreeNode, out: React.Key[] = []) => {
  if (!node.children?.length) out.push(node.key);
  else for (const child of node.children) leafKeysUnder(child, out);
  return out;
};

/**
 * 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.
 *
 * 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 } = 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. */
  const [uncontrolledExpanded, setUncontrolledExpanded] = React.useState<Set<React.Key>>(
    () => new Set(defaultExpandAll ? branchKeys(treeData) : (defaultExpandedKeys ?? []))
  );
  const expanded = React.useMemo(
    () => (expandedKeys ? new Set(expandedKeys) : uncontrolledExpanded),
    [expandedKeys, uncontrolledExpanded]
  );

  const [uncontrolledSelected, setUncontrolledSelected] = React.useState<Set<React.Key>>(
    () => new Set(defaultSelectedKeys ?? [])
  );
  const selected = React.useMemo(
    () => (selectedKeys ? new Set(selectedKeys) : uncontrolledSelected),
    [selectedKeys, uncontrolledSelected]
  );

  const [uncontrolledChecked, setUncontrolledChecked] = React.useState<Set<React.Key>>(
    () => new Set(defaultCheckedKeys ?? [])
  );
  const checked = React.useMemo(
    () => (checkedKeys ? new Set(checkedKeys) : uncontrolledChecked),
    [checkedKeys, uncontrolledChecked]
  );

  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);
  const typeAhead = React.useRef({ text: '', at: 0 });

  /**
   * 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.
   */
  const checkStateOf = React.useCallback(
    (node: TreeNode): 'checked' | 'indeterminate' | 'unchecked' => {
      const leaves = leafKeysUnder(node);
      const hits = leaves.filter((key) => checked.has(key)).length;

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

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

    if (!expandedKeys) setUncontrolledExpanded(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]);
    if (!selectedKeys) setUncontrolledSelected(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 leafKeysUnder(node)) {
      const leafNode = byKey.get(leaf);
      if (leafNode?.disabled || leafNode?.disableCheckbox) continue;
      if (turningOn) next.add(leaf);
      else next.delete(leaf);
    }

    if (!checkedKeys) setUncontrolledChecked(next);

    const halfCheckedKeys = [...byKey.values()]
      .filter((candidate) => {
        if (!isBranch(candidate)) return false;
        const leaves = leafKeysUnder(candidate);
        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 move = (delta: 1 | -1) => {
    const at = rows.findIndex((row) => row.node.key === active);
    const next = rows[at + delta];
    if (next) focusRow(next.node.key);
  };

  const handleKeyDown = (event: React.KeyboardEvent) => {
    if (active === undefined) return;

    const node = byKey.get(active)!;
    const row = rows.find((entry) => entry.node.key === active);
    const branch = isBranch(node);
    const isOpen = expanded.has(active);

    switch (event.key) {
      case 'ArrowDown':
        event.preventDefault();
        move(1);
        break;
      case 'ArrowUp':
        event.preventDefault();
        move(-1);
        break;
      case 'ArrowRight':
        event.preventDefault();
        /* A closed branch opens; an open one steps into its first child. */
        if (branch && !isOpen) setExpandedKey(active, true);
        else if (branch && isOpen) move(1);
        break;
      case 'ArrowLeft':
        event.preventDefault();
        /* An open branch closes; anything else climbs to its parent. */
        if (branch && isOpen) setExpandedKey(active, false);
        else if (row?.parentKey != null) focusRow(row.parentKey);
        break;
      case 'Home':
        event.preventDefault();
        if (rows[0]) focusRow(rows[0].node.key);
        break;
      case 'End':
        event.preventDefault();
        if (rows.length) focusRow(rows[rows.length - 1].node.key);
        break;
      case 'Enter':
        event.preventDefault();
        selectKey(active);
        break;
      case ' ':
        event.preventDefault();
        if (checkable) toggleCheck(active);
        else selectKey(active);
        break;
      case '*': {
        event.preventDefault();

        /* Opens every sibling of the focused node, per the ARIA tree pattern.
           Built as one set rather than a loop of single toggles: each toggle
           would derive its result from the same pre-update snapshot, so only
           the last sibling would survive the batch. */
        const siblings = rows.filter(
          (entry) => entry.parentKey === row?.parentKey && isBranch(entry.node)
        );
        if (siblings.length === 0) break;

        const next = new Set(expanded);
        for (const sibling of siblings) next.add(sibling.node.key);

        if (!expandedKeys) setUncontrolledExpanded(next);
        onExpand?.([...next], { node, expanded: true });
        break;
      }
      default: {
        if (event.key.length !== 1 || event.metaKey || event.ctrlKey || event.altKey) return;

        /* Type-ahead: keys pressed within a second build one search string. */
        const now = Date.now();
        typeAhead.current.text =
          now - typeAhead.current.at > 1000 ? event.key : typeAhead.current.text + event.key;
        typeAhead.current.at = now;

        const needle = typeAhead.current.text.toLowerCase();
        const at = rows.findIndex((entry) => entry.node.key === active);
        /* Search forward from the current row and wrap, so repeated presses
           cycle through the matches rather than sticking on the first. */
        const order = [...rows.slice(at + 1), ...rows.slice(0, at + 1)];
        const hit = order.find((entry) =>
          String(entry.node.title).toLowerCase().startsWith(needle)
        );

        if (hit) {
          event.preventDefault();
          focusRow(hit.node.key);
        }
      }
    }
  };

  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 };