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

type ResolveRoleMenuSubmitIdsInput = {
  menuTree: RoleResourceTreeNode[];
  baselineIds: string[];
  dirtyIds: string[];
  checkedIds: string[];
  halfCheckedIds: string[];
};

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

export function resolveRoleMenuSubmitIds(input: ResolveRoleMenuSubmitIdsInput) {
  const baselineIds = normalizeIds(input.baselineIds);

  if (!input.dirtyIds.length) {
    return baselineIds;
  }

  const treeIndex = buildTreeIndex(input.menuTree);
  const affectedIds = collectAffectedIds(treeIndex, baselineIds, normalizeIds(input.dirtyIds));

  if (!affectedIds.size) {
    return baselineIds;
  }

  const nextIdSet = new Set<string>();

  baselineIds.forEach(id => {
    if (!affectedIds.has(id)) {
      nextIdSet.add(id);
    }
  });

  normalizeIds([...input.checkedIds, ...input.halfCheckedIds]).forEach(id => {
    if (affectedIds.has(id)) {
      nextIdSet.add(id);
    }
  });

  return sortIdsByTreeOrder(treeIndex.orderedIds, nextIdSet);
}

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

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

  const walk = (items: RoleResourceTreeNode[], parentId: string | null) => {
    items.forEach(item => {
      orderedIds.push(item.id);
      parentById.set(item.id, parentId);

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

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

  walk(nodes, null);

  return {
    orderedIds,
    parentById,
    subtreeIdsById
  } satisfies TreeIndex;
}

function collectAffectedIds(treeIndex: TreeIndex, baselineIds: string[], dirtyIds: string[]) {
  const affectedIds = new Set<string>();
  const baselineIdSet = new Set(baselineIds);

  dirtyIds.forEach(dirtyId => {
    const subtreeIds = treeIndex.subtreeIdsById.get(dirtyId) ?? [dirtyId];
    subtreeIds.forEach(id => affectedIds.add(id));

    const ancestors = collectAncestorIds(treeIndex.parentById, dirtyId);

    ancestors.forEach(ancestorId => {
      if (!baselineIdSet.has(ancestorId)) {
        return;
      }

      const ancestorSubtreeIds = treeIndex.subtreeIdsById.get(ancestorId) ?? [ancestorId];
      ancestorSubtreeIds.forEach(id => affectedIds.add(id));
    });
  });

  return affectedIds;
}

function collectAncestorIds(parentById: Map<string, string | null>, nodeId: string) {
  const ancestorIds: string[] = [];

  let currentId: string | null | undefined = nodeId;

  while (currentId) {
    ancestorIds.push(currentId);
    currentId = parentById.get(currentId) ?? null;
  }

  return ancestorIds;
}

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;
}