private static @NotNull PsiElement getElementToChooseInPsiTree(@NotNull PsiElement element) {
    if (element.getFirstChild() != null) {
      return element;
    }
    var parent = element.getParent();
    if (parent == null) {
      return element;
    }

    var isIdentifierName = element.toString().toLowerCase().contains("identifier"); // heuristic
    if (isIdentifierName) {
      return parent;
    }

    if (parent instanceof PsiNameIdentifierOwner parentAsNameIdentifierOwner) {
      var parentNameIdentifier = parentAsNameIdentifierOwner.getNameIdentifier();
      if (Objects.equals(parentNameIdentifier, element)) {
        return parent;
      }
    }
    return element;
  }
private static @NotNull PsiElement getElementToChooseInPsiTree(@NotNull PsiElement element)