import React from 'react';
import type { Key, Selection, TreeProps as AriaTreeProps } from 'react-aria-components';
import {
  Collection,
  Tree as AriaTree,
  TreeItem as AriaTreeItem,
  TreeItemContent as AriaTreeItemContent,
} from 'react-aria-components';

import { useTranslations } from '@shoptet/ui-core-web';
import { useStableCallback } from '@shoptet/utils';

import { useScreenSize } from '../../hooks';
import { Button } from '../Button/Button';

import { TreeItem, type TreeItemRenderProps } from './TreeItem/TreeItem';
import { TreeItemLoader } from './TreeItem/TreeItemLoader';
import { dictionary } from './dictionary';
import type { TreeItemKey } from './useTreeState';

export interface TreeBaseProps<T extends object = object, V extends TreeItemKey = string> {
  /** The items to render in the tree. */
  items: T[];
  /** Custom render function for tree items. If not provided, uses default TreeItemContent. */
  renderItem?: (props: TreeItemRenderProps<T, V>) => React.ReactNode;
  /** A function that returns a unique key for an item object. */
  getItemValue: (item: T) => V;
  /** A function that returns the children for an item object. */
  getItemChildren: (item: T) => T[];
  /** A function that returns whether an item has children. Used for lazy loading items */
  getItemHasChildren?: (item: T) => boolean;
  /** A function that returns the text label for an item. */
  getItemTextLabel: (item: T) => string;
  /** A function that returns the text description for an item. */
  getItemTextDescription?: (item: T) => string | undefined;
  /** Handler called when children need to be loaded for an item. Returns a promise of child values. */
  onLoadChildren?: (values: V) => Promise<V[]>;
  /** Set of item keys that are currently loading their children. */
  loadingValues?: Set<V>;
  /** The currently selected item values (controlled). */
  value?: V[];
  /** The default selected item values (uncontrolled). */
  defaultValue?: V[];
  /** Handler called when the selection changes. */
  onChange?: (values: V[]) => void;
  /** The currently expanded item values (controlled). */
  expandedValues?: V[];
  /** The default expanded item values (uncontrolled). */
  defaultExpandedValues?: V[];
  /** Handler called when expanded items change. */
  onExpandedChange?: (expandedValues: V[]) => void;
  /** Accessible label for the tree. */
  label: string;
  /** Drag and drop hooks for enabling drag and drop functionality. */
  dragAndDropHooks?: AriaTreeProps<T>['dragAndDropHooks'];
  /** Display checkbox error state. */
  error?: boolean;
  /** Display checkbox warning state. */
  warning?: boolean;
  /** The id for the tree. */
  id?: string;
}
export function TreeBase<T extends object, V extends TreeItemKey = string>({
  items,
  renderItem,
  getItemValue,
  getItemChildren,
  getItemHasChildren,
  getItemTextLabel,
  getItemTextDescription,
  onLoadChildren,
  loadingValues,
  value: valueProp,
  defaultValue: defaultValueProp,
  onChange,
  expandedValues: expandedValuesProp,
  defaultExpandedValues: defaultExpandedValuesProp,
  onExpandedChange,
  label,
  dragAndDropHooks,
  error,
  warning,
  id,
  ...restProps
}: TreeBaseProps<T, V>) {
  restProps satisfies Record<PropertyKey, never>;

  const translations = useTranslations(dictionary);

  const [localValue, setLocalValue] = React.useState<V[]>(defaultValueProp || []);
  const value = valueProp === undefined ? localValue : valueProp;

  const handleChange = useStableCallback((selection: Selection) => {
    if (selection === 'all') {
      const allKeys = collectExpandableKeys(items, getItemValue, getItemChildren, getItemHasChildren);
      setLocalValue(allKeys);
      onChange?.(allKeys);
    } else {
      const selectedKeys = [...selection] as V[];
      setLocalValue(selectedKeys);
      onChange?.(selectedKeys);
    }
  });

  const [localExpandedValues, setLocalExpandedValues] = React.useState<V[]>(defaultExpandedValuesProp || []);
  const expandedValues = expandedValuesProp === undefined ? localExpandedValues : expandedValuesProp;

  const handleExpandedChange = useStableCallback((expanded: Set<Key> | Key[]) => {
    setLocalExpandedValues([...expanded] as V[]);
    onExpandedChange?.([...expanded] as V[]);
  });

  // Only root level items are considered -> If you close them, you don't see expanded children
  const shouldExpand = items.some(item => expandedValues.includes(getItemValue(item)));

  const toggleExpandAll = useStableCallback(() => {
    if (shouldExpand) {
      handleExpandedChange(new Set());
    } else {
      handleExpandedChange(collectExpandableKeys(items, getItemValue, getItemChildren, getItemHasChildren));
    }
  });

  const screenSize = useScreenSize();

  return (
    <div className='v2Tree' {...restProps}>
      {/* We do not have all the data to expand all the items when onLoadChildren is provided */}
      {onLoadChildren ? null : (
        <div className='v2Tree__header'>
          <Button
            iconRight={shouldExpand ? 'shp-minus' : 'shp-plus'}
            variant='secondary'
            size='md'
            onClick={toggleExpandAll}
          >
            {shouldExpand ? translations.collapseAll : translations.expandAll}
          </Button>
        </div>
      )}

      <div className='v2Tree__listContainer'>
        <AriaTree
          className='v2Tree__list'
          id={id}
          items={items}
          selectionMode={defaultValueProp || valueProp ? 'multiple' : 'none'}
          selectedKeys={value}
          defaultSelectedKeys={defaultValueProp}
          onSelectionChange={handleChange}
          expandedKeys={expandedValues}
          defaultExpandedKeys={defaultExpandedValuesProp}
          onExpandedChange={handleExpandedChange}
          aria-label={label}
          dragAndDropHooks={dragAndDropHooks}
          dependencies={[loadingValues, screenSize, error, warning]}
        >
          {function renderTreeItem(item) {
            const value = getItemValue(item);
            const label = getItemTextLabel(item);
            const description = getItemTextDescription ? getItemTextDescription(item) : undefined;

            const directChildrenItems = getItemChildren(item);

            const loading = loadingValues?.has(value) ?? false;

            return (
              <AriaTreeItem className='v2Tree__listItem' id={value} textValue={label}>
                <AriaTreeItemContent>
                  {state =>
                    renderItem ? (
                      renderItem({
                        value,
                        state,
                        item,
                        label,
                        description,
                        loading: loading,
                        screenSize,
                        error,
                        warning,
                      })
                    ) : (
                      <TreeItem
                        value={value}
                        state={state}
                        label={label}
                        description={description}
                        loading={loading}
                        screenSize={screenSize}
                        error={error}
                        warning={warning}
                      />
                    )
                  }
                </AriaTreeItemContent>

                <Collection items={directChildrenItems}>{renderTreeItem}</Collection>

                {onLoadChildren && getItemHasChildren?.(item) && directChildrenItems.length === 0 && (
                  <TreeItemLoader onLoad={onLoadChildren} loading={loading} />
                )}
              </AriaTreeItem>
            );
          }}
        </AriaTree>
      </div>
    </div>
  );
}

function collectExpandableKeys<T extends object, V extends TreeItemKey>(
  items: T[],
  getItemValue: (item: T) => V,
  getItemChildren: (item: T) => T[],
  getItemHasChildren?: (item: T) => boolean,
  accumulator: V[] = []
): V[] {
  items.forEach(item => {
    const key = getItemValue(item);

    if (getItemHasChildren?.(item)) {
      accumulator.push(key);
    }

    const children = getItemChildren(item);

    if (children.length > 0) {
      accumulator.push(key);
      collectExpandableKeys(children, getItemValue, getItemChildren, getItemHasChildren, accumulator);
    }
  });

  return accumulator;
}