import { useModel } from '@umijs/max';
import { Tree } from 'antd';
import type { DataNode } from 'antd/es/tree';
import { useEffect, useState } from 'react';

interface TreeGroupProps {
  groupIds?: string | string[] | null;
  multiple?: boolean;
  allowedGroupIds?: string[];
  disable?: boolean;
  onSelected?: (value: string | string[] | null) => void;
  titleRender?: (nodeData: DataNode) => React.ReactNode;
}

const TreeGroup = ({
  groupIds,
  multiple = false,
  allowedGroupIds,
  disable = false,
  onSelected,
  titleRender,
}: TreeGroupProps) => {
  const { initialState } = useModel('@@initialState');
  const { currentUserProfile } = initialState || {};
  const [treeData, setTreeData] = useState<DataNode[]>([]);
  const [selectedKeys, setSelectedKeys] = useState<React.Key[]>([]);
  const { groups, getGroups } = useModel('master.useGroups');

  const filterGroupsByAllowedIds = (
    groups: MasterModel.GroupNode[],
    allowedIds: Set<string>,
  ): MasterModel.GroupNode[] => {
    return groups
      .map((group) => {
        const filteredChildren = group.children
          ? filterGroupsByAllowedIds(group.children, allowedIds)
          : [];

        if (allowedIds.has(group.id) || filteredChildren.length > 0) {
          return {
            ...group,
            children: filteredChildren,
          };
        }

        return null;
      })
      .filter(Boolean) as MasterModel.GroupNode[];
  };

  const transformGroupNodeToTreeData = (
    groups: MasterModel.GroupNode[],
  ): DataNode[] =>
    groups.map((group) => ({
      title: group.name,
      key: group.id,
      children: group.children?.length
        ? transformGroupNodeToTreeData(group.children)
        : undefined,
    }));

  useEffect(() => {
    if (!groups) return;

    let displayGroups = groups;

    if (allowedGroupIds && allowedGroupIds.length > 0) {
      const allowedSet = new Set(allowedGroupIds);
      displayGroups = filterGroupsByAllowedIds(groups, allowedSet);
    }

    setTreeData(transformGroupNodeToTreeData(displayGroups));
  }, [groups, allowedGroupIds]);

  useEffect(() => {
    if (!groups) {
      getGroups();
    }
  }, [groups]);

  useEffect(() => {
    const loadTreeData = async () => {
      const transformedData = transformGroupNodeToTreeData(groups || []);
      setTreeData(transformedData);
    };

    loadTreeData();
  }, [currentUserProfile, groups]);

  useEffect(() => {
    // Set selected keys from groupIds prop
    if (groupIds) {
      setSelectedKeys(Array.isArray(groupIds) ? groupIds : [groupIds]);
    } else {
      setSelectedKeys([]);
    }
  }, [groupIds]);

  const handleSelect = (selectedKeys: React.Key[]) => {
    const keys = selectedKeys.map((k) => String(k));
    setSelectedKeys(selectedKeys);

    if (multiple) {
      onSelected?.(keys.length > 0 ? keys : null);
    } else {
      onSelected?.(keys.length > 0 ? keys[0] : null);
    }
  };

  const handleCheck = (
    checkedKeys:
      | React.Key[]
      | { checked: React.Key[]; halfChecked: React.Key[] },
  ) => {
    const keys = Array.isArray(checkedKeys) ? checkedKeys : checkedKeys.checked;
    const stringKeys = keys.map((k) => String(k));
    setSelectedKeys(keys);

    onSelected?.(stringKeys.length > 0 ? stringKeys : null);
  };

  return (
    <Tree
      disabled={disable}
      checkable={multiple}
      treeData={treeData}
      checkedKeys={multiple ? selectedKeys : undefined}
      selectedKeys={!multiple ? selectedKeys : undefined}
      onCheck={multiple ? handleCheck : undefined}
      onSelect={!multiple ? handleSelect : undefined}
      defaultExpandAll={true}
      showLine
      autoExpandParent
      {...(titleRender && { titleRender })}
    />
  );
};

export default TreeGroup;
