import { Tree } from 'antd';
import { uniq, remove } from 'lodash';
import React, { useState, useEffect } from 'react';

const getChildKeys = tree => {
  let childKeys = tree.children?.map(c => getChildKeys(c)).flat(Infinity) || [];
  return [tree.key, ...childKeys];
};

export default props => {
  const { expandedKeys, onExpand, onSelect, selectedKeys, ...rest } = props;
  const [epKeys, setEpKeys] = useState([expandedKeys]);

  const handleExpand = (keys, { expanded, node }) => {
    if (!expanded) {
      let childKeys = getChildKeys(node);
      remove(keys, k => childKeys.includes(k));
    }
    setEpKeys(keys);
  };

  const handleSelect = (keys, e) => {
    let result = [];
    if (keys.length > 0 && epKeys.includes(keys[0])) {
      remove(epKeys, k => keys.includes(k));
      result = [...epKeys];
    } else if (keys[0]) {
      result = [...epKeys, keys[0]];
    } else {
      if (epKeys.includes(e.node.key)) {
        remove(epKeys, k => k === e.node.key);
      } else {
        epKeys.push(e.node.key);
      }
      result = [...epKeys];
    }
    // eslint-disable-next-line no-unused-expressions
    onSelect && onSelect(keys, e);
    setEpKeys(result);
  };

  return (
    <Tree
      onSelect={handleSelect}
      onExpand={handleExpand}
      expandedKeys={epKeys}
      selectedKeys={selectedKeys}
      {...rest}
    />
  );
};
