export type RoleResourceTreeNode = {
  id: string;
  children?: RoleResourceTreeNode[] | null;
};

type NormalizeRoleMenuCheckedIdsInput = {
  menuTree: RoleResourceTreeNode[];
  checkedIds: string[];
};

type TreeIndex = {
  orderedIds: string[];
  subtreeIdsById: Map<string, string[]>;
};

export function normalizeRoleMenuCheckedIds(input: NormalizeRoleMenuCheckedIdsInput) {
  const checkedIds = normalizeIds(input.checkedIds);

  if (!checkedIds.length) {
    return checkedIds;
  }

  const treeIndex = buildTreeIndex(input.menuTree);
  const normalizedIdSet = new Set(checkedIds);

  treeIndex.orderedIds.forEach(id => {
    if (!normalizedIdSet.has(id)) {
      return;
    }

    const descendantIds = (treeIndex.subtreeIdsById.get(id) ?? [id]).slice(1);

    if (!descendantIds.length) {
      return;
    }

    const selectedDescendantCount = descendantIds.reduce((count, descendantId) => {
      return normalizedIdSet.has(descendantId) ? count + 1 : count;
    }, 0);

    if (selectedDescendantCount > 0 && selectedDescendantCount < descendantIds.length) {
      normalizedIdSet.delete(id);
    }
  });

  return sortIdsByTreeOrder(treeIndex.orderedIds, normalizedIdSet);
}

function normalizeIds(ids: string[]) {
  return [...new Set(ids.map(id => String(id).trim()).filter(Boolean))];
}

function buildTreeIndex(nodes: RoleResourceTreeNode[]) {
  const orderedIds: string[] = [];
  const subtreeIdsById = new Map<string, string[]>();

  const walk = (items: RoleResourceTreeNode[]) => {
    items.forEach(item => {
      orderedIds.push(item.id);

      const childIds = item.children?.length ? walk(item.children) : [];
      subtreeIdsById.set(item.id, [item.id, ...childIds]);
    });

    return items.flatMap(item => subtreeIdsById.get(item.id) ?? [item.id]);
  };

  walk(nodes);

  return {
    orderedIds,
    subtreeIdsById
  } satisfies TreeIndex;
}

function sortIdsByTreeOrder(orderedIds: string[], idSet: Set<string>) {
  const sortedIds: string[] = [];

  orderedIds.forEach(id => {
    if (idSet.has(id)) {
      sortedIds.push(id);
      idSet.delete(id);
    }
  });

  idSet.forEach(id => {
    sortedIds.push(id);
  });

  return sortedIds;
}
