'use client';

import * as React from 'react';

import { ChevronDownIcon, XIcon } from '@/icons';
import { cn } from '@/lib/utils';
import { focusRing } from '@/lib/cva-presets';
import { useControllableState } from '@/hooks/use-controllable-state';
import { Popover, PopoverAnchor, PopoverContent } from '@/components/popover/popover';
import { Tree, type TreeNode } from '@/components/tree/tree';

export interface TreeSelectNode extends Omit<TreeNode, 'key' | 'children'> {
  value: React.Key;
  children?: TreeSelectNode[];
  /** Only used for searching and for the chip label; `title` is what renders. */
  label?: string;
}

export interface TreeSelectProps {
  treeData: TreeSelectNode[];
  value?: React.Key | React.Key[] | null;
  defaultValue?: React.Key | React.Key[] | null;
  onChange?: (value: React.Key | React.Key[] | null) => void;
  /** Allow several selections, shown as removable chips. */
  multiple?: boolean;
  /** Checkboxes in the panel. Implies `multiple`. */
  treeCheckable?: boolean;
  treeDefaultExpandAll?: boolean;
  /** Filter the tree from a search box above it. */
  showSearch?: boolean;
  searchPlaceholder?: string;
  placeholder?: string;
  emptyText?: React.ReactNode;
  disabled?: boolean;
  allowClear?: boolean;
  className?: string;
  contentClassName?: string;
  id?: string;
  'aria-label'?: string;
  'aria-labelledby'?: string;
}

/** The tree the panel renders, keyed the way `Tree` expects. */
const toTreeNodes = (nodes: TreeSelectNode[]): TreeNode[] =>
  nodes.map((node) => ({
    key: node.value,
    title: node.title,
    disabled: node.disabled,
    disableCheckbox: node.disableCheckbox,
    isLeaf: node.isLeaf,
    children: node.children ? toTreeNodes(node.children) : undefined,
  }));

/** Text a node is searched and labelled by. */
const labelOf = (node: TreeSelectNode) =>
  node.label ?? (typeof node.title === 'string' ? node.title : String(node.value));

const indexByValue = (nodes: TreeSelectNode[], out = new Map<React.Key, TreeSelectNode>()) => {
  for (const node of nodes) {
    out.set(node.value, node);
    if (node.children) indexByValue(node.children, out);
  }
  return out;
};

/** Every value that has children — what "expand all" means for this tree. */
const branchValues = (nodes: TreeSelectNode[], out: React.Key[] = []) => {
  for (const node of nodes) {
    if (node.children?.length) {
      out.push(node.value);
      branchValues(node.children, out);
    }
  }
  return out;
};

/**
 * Keeps a branch whose own label matches, or which still has a matching
 * descendant — dropping a parent would orphan its matching children.
 */
const filterTree = (nodes: TreeSelectNode[], query: string): TreeSelectNode[] => {
  const needle = query.trim().toLowerCase();
  if (!needle) return nodes;

  return nodes.reduce<TreeSelectNode[]>((kept, node) => {
    const children = node.children ? filterTree(node.children, needle) : undefined;
    const hit = labelOf(node).toLowerCase().includes(needle);

    if (hit || children?.length) kept.push({ ...node, children: children ?? node.children });
    return kept;
  }, []);
};

/**
 * Select whose options are a tree — categories, org units, folder paths.
 *
 * Combobox flattens its options, which loses the parent/child relationship
 * that makes a taxonomy navigable. This keeps the hierarchy, and with
 * `treeCheckable` adds branches whose checked state propagates both ways.
 *
 * ```tsx
 * <TreeSelect
 *   treeCheckable
 *   placeholder="Pick a team"
 *   treeData={[{ value: 'eng', title: 'Engineering', children: [
 *     { value: 'fe', title: 'Frontend' },
 *   ] }]}
 * />
 * ```
 */
function TreeSelect({
  treeData,
  value,
  defaultValue = null,
  onChange,
  multiple,
  treeCheckable,
  treeDefaultExpandAll,
  showSearch,
  searchPlaceholder = 'Search…',
  placeholder = 'Select…',
  emptyText = 'No results found.',
  disabled,
  allowClear = true,
  className,
  contentClassName,
  id,
  'aria-label': ariaLabel,
  'aria-labelledby': ariaLabelledBy,
}: TreeSelectProps) {
  const panelId = React.useId();
  const isMulti = Boolean(multiple || treeCheckable);

  const [open, setOpen] = React.useState(false);
  const [search, setSearch] = React.useState('');
  const [current, commit] = useControllableState<React.Key | React.Key[] | null>({
    value,
    defaultValue,
    onChange,
  });

  const byValue = React.useMemo(() => indexByValue(treeData), [treeData]);
  const visible = React.useMemo(
    () => (showSearch ? filterTree(treeData, search) : treeData),
    [treeData, search, showSearch]
  );

  /**
   * Expansion is owned here rather than left to `Tree`, because searching has
   * to open the branches holding the matches. `Tree`'s `defaultExpandAll` is a
   * *default* — it is read once on mount and would never react to a query
   * typed afterwards.
   */
  const [expandedKeys, setExpandedKeys] = React.useState<React.Key[]>(() =>
    treeDefaultExpandAll ? branchValues(treeData) : []
  );
  const searching = Boolean(showSearch && search.trim());
  const effectiveExpanded = searching ? branchValues(visible) : expandedKeys;

  const selectedValues: React.Key[] =
    current == null ? [] : Array.isArray(current) ? current : [current];

  const chips = selectedValues
    .map((entry) => byValue.get(entry))
    .filter((node): node is TreeSelectNode => Boolean(node));

  /* The trigger carries `role="combobox"`, and that role takes no name from
     its content — without one of these three it is announced unnamed. The
     placeholder is the last resort, being the only wording guaranteed present. */
  const triggerLabel = ariaLabel ?? (ariaLabelledBy ? undefined : placeholder);

  return (
    <Popover
      open={open}
      onOpenChange={(next) => {
        setOpen(next);
        if (!next) setSearch('');
      }}
    >
      <PopoverAnchor asChild>
        <div
          className={cn(
            'flex min-h-9 w-full items-center gap-2 rounded-md border border-input bg-transparent px-3 py-1 text-sm shadow-xs dark:bg-input/30',
            'transition-[color,box-shadow] duration-(--ui-duration-fast) ease-(--ui-ease-standard)',
            'focus-within:border-ring focus-within:ring-[3px] focus-within:ring-ring/50',
            disabled && 'cursor-not-allowed opacity-50',
            className
          )}
          data-slot="tree-select"
        >
          <button
            type="button"
            id={id}
            role="combobox"
            aria-expanded={open}
            aria-haspopup="tree"
            aria-controls={open ? panelId : undefined}
            aria-label={triggerLabel}
            aria-labelledby={ariaLabelledBy}
            disabled={disabled}
            data-slot="tree-select-trigger"
            onClick={() => setOpen((was) => !was)}
            className={cn(
              'flex flex-1 flex-wrap items-center gap-1 rounded-sm py-1 text-start outline-none',
              focusRing
            )}
          >
            {chips.length === 0 ? (
              <span className="text-muted-foreground">{placeholder}</span>
            ) : isMulti ? (
              chips.map((node) => (
                <span
                  key={node.value}
                  data-slot="tree-select-chip"
                  className="inline-flex items-center gap-1 rounded-sm bg-secondary px-1.5 py-0.5 text-xs text-secondary-foreground"
                >
                  {labelOf(node)}
                  {/* A span, not a button: this sits inside the trigger button,
                      where nesting one would be invalid HTML. */}
                  <span
                    role="button"
                    aria-label={`Remove ${labelOf(node)}`}
                    tabIndex={-1}
                    onClick={(event) => {
                      event.stopPropagation();
                      commit(selectedValues.filter((entry) => entry !== node.value));
                    }}
                    className="cursor-pointer opacity-60 hover:opacity-100"
                  >
                    <XIcon className="size-3" />
                  </span>
                </span>
              ))
            ) : (
              <span className="truncate">{labelOf(chips[0])}</span>
            )}
          </button>

          {allowClear && chips.length > 0 && !disabled ? (
            <button
              type="button"
              aria-label="Clear selection"
              data-slot="tree-select-clear"
              onClick={() => commit(isMulti ? [] : null)}
              className={cn(
                'shrink-0 rounded-sm text-muted-foreground hover:text-foreground',
                focusRing
              )}
            >
              <XIcon className="size-3.5" />
            </button>
          ) : null}
          <ChevronDownIcon className="size-4 shrink-0 opacity-50" />
        </div>
      </PopoverAnchor>

      <PopoverContent
        id={panelId}
        align="start"
        className={cn('w-(--radix-popover-trigger-width) p-0', contentClassName)}
        data-slot="tree-select-panel"
      >
        {showSearch ? (
          <div className="border-b p-2">
            <input
              type="text"
              autoComplete="off"
              aria-label={searchPlaceholder}
              placeholder={searchPlaceholder}
              value={search}
              onChange={(event) => setSearch(event.target.value)}
              /* `focusRing` rather than the band's own `focus-within:`: this is
                 the first control a keyboard reaches once the panel opens, and
                 the ring belongs on the field, not on a full-bleed strip. Same
                 preset the trigger above uses, so the two match. */
              className={cn(
                'w-full rounded-sm bg-transparent text-sm placeholder:text-muted-foreground',
                focusRing
              )}
            />
          </div>
        ) : null}

        {visible.length === 0 ? (
          <div className="py-6 text-center text-sm text-muted-foreground">{emptyText}</div>
        ) : (
          <Tree
            className="max-h-64 overflow-y-auto p-1"
            aria-label={triggerLabel ?? placeholder}
            treeData={toTreeNodes(visible)}
            checkable={treeCheckable}
            expandedKeys={effectiveExpanded}
            onExpand={setExpandedKeys}
            selectedKeys={treeCheckable ? [] : selectedValues}
            checkedKeys={treeCheckable ? selectedValues : undefined}
            onCheck={(keys) => commit(keys)}
            onSelect={(keys) => {
              if (treeCheckable) return;

              if (!isMulti) {
                commit(keys[0] ?? null);
                setOpen(false);
                return;
              }

              const picked = keys[0];
              if (picked === undefined) return;
              commit(
                selectedValues.includes(picked)
                  ? selectedValues.filter((entry) => entry !== picked)
                  : [...selectedValues, picked]
              );
            }}
          />
        )}
      </PopoverContent>
    </Popover>
  );
}

export { TreeSelect };