// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package com.jetbrains.python.formatter;

import com.intellij.formatting.*;
import com.intellij.lang.ASTNode;
import com.intellij.lang.injection.InjectedLanguageManager;
import com.intellij.openapi.util.TextRange;
import com.intellij.psi.*;
import com.intellij.psi.codeStyle.CommonCodeStyleSettings;
import com.intellij.psi.impl.source.tree.LeafElement;
import com.intellij.psi.impl.source.tree.TreeUtil;
import com.intellij.psi.tree.IElementType;
import com.intellij.psi.tree.TokenSet;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.ArrayUtil;
import com.intellij.util.ObjectUtils;
import com.jetbrains.python.PyElementTypes;
import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.PythonCodeStyleService;
import com.jetbrains.python.PythonDialectsTokenSetProvider;
import com.jetbrains.python.ast.*;
import com.jetbrains.python.ast.impl.PyPsiUtilsCore;
import com.jetbrains.python.ast.impl.PyUtilCore;
import com.jetbrains.python.pyi.PyiUtilCore;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

import static com.jetbrains.python.formatter.PyCodeStyleSettings.DICT_ALIGNMENT_ON_COLON;
import static com.jetbrains.python.formatter.PyCodeStyleSettings.DICT_ALIGNMENT_ON_VALUE;


public class PyBlock implements ASTBlock {
  private static final TokenSet STATEMENT_OR_DECLARATION = PythonDialectsTokenSetProvider.getInstance().getStatementTokens();


  private static final TokenSet ourListElementTypes = TokenSet.create(PyElementTypes.LIST_LITERAL_EXPRESSION,
                                                                      PyElementTypes.LIST_COMP_EXPRESSION,
                                                                      PyElementTypes.DICT_LITERAL_EXPRESSION,
                                                                      PyElementTypes.DICT_COMP_EXPRESSION,
                                                                      PyElementTypes.SET_LITERAL_EXPRESSION,
                                                                      PyElementTypes.SET_COMP_EXPRESSION,
                                                                      PyElementTypes.ARGUMENT_LIST,
                                                                      PyElementTypes.PARAMETER_LIST,
                                                                      PyElementTypes.TUPLE_EXPRESSION,
                                                                      PyElementTypes.PARENTHESIZED_EXPRESSION,
                                                                      PyElementTypes.SUBSCRIPTION_EXPRESSION,
                                                                      PyElementTypes.GENERATOR_EXPRESSION,
                                                                      PyElementTypes.SEQUENCE_PATTERN,
                                                                      PyElementTypes.MAPPING_PATTERN,
                                                                      PyElementTypes.PATTERN_ARGUMENT_LIST,
                                                                      PyElementTypes.TYPE_PARAMETER_LIST);

  private static final TokenSet ourCollectionLiteralTypes = TokenSet.create(PyElementTypes.LIST_LITERAL_EXPRESSION,
                                                                            PyElementTypes.LIST_COMP_EXPRESSION,
                                                                            PyElementTypes.DICT_LITERAL_EXPRESSION,
                                                                            PyElementTypes.DICT_COMP_EXPRESSION,
                                                                            PyElementTypes.SET_LITERAL_EXPRESSION,
                                                                            PyElementTypes.SET_COMP_EXPRESSION);

  private static final TokenSet ourHangingIndentOwners = TokenSet.create(PyElementTypes.LIST_LITERAL_EXPRESSION,
                                                                         PyElementTypes.LIST_COMP_EXPRESSION,
                                                                         PyElementTypes.DICT_LITERAL_EXPRESSION,
                                                                         PyElementTypes.DICT_COMP_EXPRESSION,
                                                                         PyElementTypes.SET_LITERAL_EXPRESSION,
                                                                         PyElementTypes.SET_COMP_EXPRESSION,
                                                                         PyElementTypes.ARGUMENT_LIST,
                                                                         PyElementTypes.PARAMETER_LIST,
                                                                         PyElementTypes.TUPLE_EXPRESSION,
                                                                         PyElementTypes.PARENTHESIZED_EXPRESSION,
                                                                         PyElementTypes.GENERATOR_EXPRESSION,
                                                                         PyElementTypes.FUNCTION_DECLARATION,
                                                                         PyElementTypes.CALL_EXPRESSION,
                                                                         PyElementTypes.FROM_IMPORT_STATEMENT,
                                                                         PyElementTypes.SEQUENCE_PATTERN,
                                                                         PyElementTypes.MAPPING_PATTERN,
                                                                         PyElementTypes.PATTERN_ARGUMENT_LIST,
                                                                         PyElementTypes.WITH_STATEMENT,
                                                                         PyElementTypes.TYPE_PARAMETER_LIST);

  private static final boolean ALIGN_CONDITIONS_WITHOUT_PARENTHESES = false;

  private final PyBlock myParent;
  private final Alignment myAlignment;
  private final Indent myIndent;
  private final ASTNode myNode;
  private final Wrap myWrap;
  private final PyBlockContext myContext;
  private List<PyBlock> mySubBlocks = null;
  private Map<ASTNode, PyBlock> mySubBlockByNode = null;
  private final boolean myEmptySequence;

  // Shared among multiple children sub-blocks
  private Alignment myChildAlignment = null;
  private Alignment myDictAlignment = null;
  private Wrap myDictWrapping = null;
  private Wrap myListWrapping = null;
  private Wrap mySetWrapping = null;
  private Wrap myTupleWrapping = null;
  private Wrap myFromImportWrapping = null;
  private Wrap myParameterListWrapping = null;
  private Wrap myArgumentListWrapping = null;

  public PyBlock(@Nullable PyBlock parent,
                 @NotNull ASTNode node,
                 @Nullable Alignment alignment,
                 @NotNull Indent indent,
                 @Nullable Wrap wrap,
                 @NotNull PyBlockContext context) {
    myParent = parent;
    myAlignment = alignment;
    myIndent = indent;
    myNode = node;
    myWrap = wrap;
    myContext = context;
    myEmptySequence = isEmptySequence(node);

    final CommonCodeStyleSettings settings = myContext.getSettings();
    final PyCodeStyleSettings pySettings = myContext.getPySettings();
    if (node.getElementType() == PyElementTypes.DICT_LITERAL_EXPRESSION) {
      myDictAlignment = Alignment.createAlignment(true);
      myDictWrapping = Wrap.createWrap(pySettings.DICT_WRAPPING, true);
    }
    else if (node.getElementType() == PyElementTypes.LIST_LITERAL_EXPRESSION) {
      myListWrapping = Wrap.createWrap(pySettings.LIST_WRAPPING, pySettings.LIST_NEW_LINE_AFTER_LEFT_BRACKET);
    }
    else if (node.getElementType() == PyElementTypes.SET_LITERAL_EXPRESSION) {
      mySetWrapping = Wrap.createWrap(pySettings.SET_WRAPPING, pySettings.SET_NEW_LINE_AFTER_LEFT_BRACE);
    }
    else if (node.getElementType() == PyElementTypes.TUPLE_EXPRESSION &&
             node.getTreeParent().getElementType() == PyElementTypes.PARENTHESIZED_EXPRESSION) {
      myTupleWrapping = Wrap.createWrap(pySettings.TUPLE_WRAPPING, pySettings.TUPLE_NEW_LINE_AFTER_LEFT_PARENTHESIS);
    }
    else if (node.getElementType() == PyElementTypes.FROM_IMPORT_STATEMENT) {
      myFromImportWrapping = Wrap.createWrap(pySettings.FROM_IMPORT_WRAPPING, false);
    }
    else if (node.getElementType() == PyElementTypes.PARAMETER_LIST) {
      myParameterListWrapping = Wrap.createWrap(settings.METHOD_PARAMETERS_WRAP, settings.METHOD_PARAMETERS_LPAREN_ON_NEXT_LINE);
    }
    else if (node.getElementType() == PyElementTypes.ARGUMENT_LIST) {
      myArgumentListWrapping = Wrap.createWrap(settings.CALL_PARAMETERS_WRAP, settings.CALL_PARAMETERS_LPAREN_ON_NEXT_LINE);
    }
  }

  @Override
  public @NotNull ASTNode getNode() {
    return myNode;
  }

  @Override
  public @NotNull TextRange getTextRange() {
    return myNode.getTextRange();
  }

  private Alignment getAlignmentForChildren() {
    if (myChildAlignment == null) {
      myChildAlignment = Alignment.createAlignment();
    }
    return myChildAlignment;
  }

  @Override
  public @NotNull List<Block> getSubBlocks() {
    if (mySubBlocks == null) {
      mySubBlockByNode = buildSubBlocks();
      mySubBlocks = new ArrayList<>(mySubBlockByNode.values());
    }
    return Collections.unmodifiableList(mySubBlocks);
  }

  private @Nullable PyBlock getSubBlockByNode(@NotNull ASTNode node) {
    return mySubBlockByNode.get(node);
  }

  private @Nullable PyBlock getSubBlockByIndex(int index) {
    return mySubBlocks.get(index);
  }

  private @NotNull Map<ASTNode, PyBlock> buildSubBlocks() {
    final Map<ASTNode, PyBlock> blocks = new LinkedHashMap<>();
    for (ASTNode child: getSubBlockNodes()) {

      final IElementType childType = child.getElementType();

      if (child.getTextRange().isEmpty()) continue;

      if (childType == TokenType.WHITE_SPACE) {
        continue;
      }

      blocks.put(child, buildSubBlock(child));
    }
    return Collections.unmodifiableMap(blocks);
  }

  protected @NotNull Iterable<ASTNode> getSubBlockNodes() {
    if (myNode.getElementType() == PyElementTypes.BINARY_EXPRESSION) {
      final ArrayList<ASTNode> result = new ArrayList<>();
      collectChildrenOperatorAndOperandNodes(myNode, result);
      return result;
    }
    return Arrays.asList(myNode.getChildren(null));
  }

  private static void collectChildrenOperatorAndOperandNodes(@NotNull ASTNode node, @NotNull List<ASTNode> result) {
    if (node.getElementType() == PyElementTypes.BINARY_EXPRESSION) {
      for (ASTNode child : node.getChildren(null)) {
        collectChildrenOperatorAndOperandNodes(child, result);
      }
    }
    else {
      result.add(node);
    }
  }

  protected @NotNull PyBlock buildSubBlock(@NotNull ASTNode child) {
    final IElementType parentType = myNode.getElementType();

    final ASTNode grandParentNode = myNode.getTreeParent();
    final IElementType grandparentType = grandParentNode == null ? null : grandParentNode.getElementType();

    final IElementType childType = child.getElementType();
    Wrap childWrap = null;
    Indent childIndent = Indent.getNoneIndent();
    Alignment childAlignment = null;

    final PyCodeStyleSettings settings = myContext.getPySettings();

    if (childType == PyElementTypes.STATEMENT_LIST) {
      if (hasLineBreaksBeforeInSameParent(child, 1) || needLineBreakInStatement()) {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (childType == PyElementTypes.CASE_CLAUSE) {
      childIndent = Indent.getNormalIndent();
    }
    else if (childType == PyElementTypes.IMPORT_ELEMENT) {
      if (parentType == PyElementTypes.FROM_IMPORT_STATEMENT) {
        childWrap = myFromImportWrapping;
      }
      else {
        childWrap = Wrap.createWrap(WrapType.NORMAL, true);
      }
      childIndent = Indent.getNormalIndent();
    }
    if (childType == PyTokenTypes.END_OF_LINE_COMMENT && (parentType == PyElementTypes.FROM_IMPORT_STATEMENT ||
                                                          parentType == PyElementTypes.MATCH_STATEMENT)) {
      childIndent = Indent.getNormalIndent();
    }

    // TODO merge it with the following check of needListAlignment()
    if (ourListElementTypes.contains(parentType)) {
      // wrapping in non-parenthesized tuple expression is not allowed (PY-1792)
      if (!(childType == PyElementTypes.TUPLE_EXPRESSION && parentType == PyElementTypes.PARENTHESIZED_EXPRESSION) &&
          (parentType != PyElementTypes.TUPLE_EXPRESSION || grandparentType == PyElementTypes.PARENTHESIZED_EXPRESSION) &&
          !PyTokenTypes.ALL_BRACES.contains(childType) &&
          childType != PyTokenTypes.COMMA &&
          !isSliceOperand(child) /*&& !isSubscriptionOperand(child)*/) {
        childWrap = Wrap.createWrap(WrapType.NORMAL, true);
      }
      if (needListAlignment(child) && !myEmptySequence) {
        childAlignment = getAlignmentForChildren();
      }
      if (childType == PyTokenTypes.END_OF_LINE_COMMENT) {
        childIndent = Indent.getNormalIndent();
      }
    }

    if (parentType == PyElementTypes.BINARY_EXPRESSION) {
      if (childType != PyElementTypes.BINARY_EXPRESSION) {
        final PyBlock topmostBinary = findTopmostBinaryExpressionBlock(child);
        assert topmostBinary != null;
        final PyBlock binaryParentBlock = topmostBinary.myParent;
        final ASTNode binaryParentNode = binaryParentBlock.myNode;
        final IElementType binaryParentType = binaryParentNode.getElementType();
        // TODO check for comprehensions explicitly here
        if (ourListElementTypes.contains(binaryParentType) && needListAlignment(child) && !myEmptySequence) {
          childAlignment = binaryParentBlock.getChildAlignment();
        }
        final boolean parenthesised = binaryParentType == PyElementTypes.PARENTHESIZED_EXPRESSION;
        if (binaryParentType != PyElementTypes.RETURN_STATEMENT &&
            binaryParentType != PyElementTypes.YIELD_EXPRESSION) {
          if (childAlignment == null && topmostBinary != null &&
              !(parenthesised && isIfCondition(binaryParentNode)) &&
              !(isCondition(topmostBinary.myNode) && !ALIGN_CONDITIONS_WITHOUT_PARENTHESES)) {
            childAlignment = topmostBinary.getAlignmentForChildren();
          }
          // We omit indentation for the binary expression itself in this case (similarly to PyTupleExpression inside
          // PyParenthesisedExpression) because we indent individual operands and operators inside rather than
          // the whole contained expression.
          childIndent = parenthesised ? Indent.getContinuationIndent() : Indent.getContinuationWithoutFirstIndent();
        }
        else {
          childIndent = Indent.getNormalIndent();
        }
      }
    }
    else if (parentType == PyElementTypes.OR_PATTERN) {
      childAlignment = getAlignmentForChildren();
    }
    else if (parentType == PyElementTypes.SEQUENCE_PATTERN || parentType == PyElementTypes.MAPPING_PATTERN) {
      if (PyTokenTypes.CLOSE_BRACES.contains(childType) && !settings.HANG_CLOSING_BRACKETS ||
          PyTokenTypes.OPEN_BRACES.contains(childType)) {
        childIndent = Indent.getNoneIndent();
      }
      else {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (ourCollectionLiteralTypes.contains(parentType) ||
             parentType == PyElementTypes.TUPLE_EXPRESSION && grandparentType == PyElementTypes.PARENTHESIZED_EXPRESSION) {
      if ((PyTokenTypes.CLOSE_BRACES.contains(childType) && !settings.HANG_CLOSING_BRACKETS) || 
          PyTokenTypes.OPEN_BRACES.contains(childType)) {
        childIndent = Indent.getNoneIndent();
      }
      else if (settings.USE_CONTINUATION_INDENT_FOR_COLLECTION_AND_COMPREHENSIONS) {
        childIndent = Indent.getContinuationIndent();
      }
      else {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (parentType == PyElementTypes.STRING_LITERAL_EXPRESSION) {
      if (PyTokenTypes.STRING_NODES.contains(childType) || childType == PyElementTypes.FSTRING_NODE) {
        if (childType == PyTokenTypes.TRIPLE_QUOTED_STRING &&
            myContext.getPySettings().FORMAT_INJECTED_FRAGMENTS &&
            isInsideInjection(child)) {
          return createInjectedBlock(child, myContext.getPySettings().ADD_INDENT_INSIDE_INJECTIONS
                                            ? Indent.getNormalIndent()
                                            : Indent.getNoneIndent());
        }
        childAlignment = getAlignmentForChildren();
      }
    }
    else if (parentType == PyElementTypes.FROM_IMPORT_STATEMENT) {
      if (myNode.findChildByType(PyTokenTypes.LPAR) != null) {
        if (childType == PyElementTypes.IMPORT_ELEMENT) {
          if (settings.ALIGN_MULTILINE_IMPORTS) {
            childAlignment = getAlignmentForChildren();
          }
          else {
            childIndent = Indent.getNormalIndent();
          }
        }
        if (childType == PyTokenTypes.RPAR) {
          childIndent = Indent.getNoneIndent();
          // Don't have hanging indent and is not going to have it due to the setting about opening parenthesis
          if (!hasHangingIndent(myNode.getPsi()) && !settings.FROM_IMPORT_NEW_LINE_AFTER_LEFT_PARENTHESIS) {
            childAlignment = getAlignmentForChildren();
          }
          else if (settings.HANG_CLOSING_BRACKETS) {
            childIndent = Indent.getNormalIndent();
          }
        }
      }
    }
    else if (isValueOfKeyValuePair(child)) {
      childIndent = Indent.getNormalIndent();
    }
    // TODO: merge with needChildAlignment()
    //Align elements vertically if there is an argument in the first line of parenthesized expression
    else if (!hasHangingIndent(myNode.getPsi()) &&
             (parentType == PyElementTypes.PARENTHESIZED_EXPRESSION ||
              (parentType == PyElementTypes.ARGUMENT_LIST && myContext.getSettings().ALIGN_MULTILINE_PARAMETERS_IN_CALLS) ||
              (parentType == PyElementTypes.PARAMETER_LIST && myContext.getSettings().ALIGN_MULTILINE_PARAMETERS)) &&
             !isIndentNext(child) &&
             !hasLineBreaksBeforeInSameParent(myNode.getFirstChildNode(), 1) &&
             !ourListElementTypes.contains(childType)) {

      if (!PyTokenTypes.ALL_BRACES.contains(childType)) {
        childAlignment = getAlignmentForChildren();
        childIndent = Indent.getNormalIndent();
      }
      else if (childType == PyTokenTypes.RPAR) {
        if (settings.HANG_CLOSING_BRACKETS) {
          childIndent = settings.USE_CONTINUATION_INDENT_FOR_COLLECTION_AND_COMPREHENSIONS ? Indent.getContinuationIndent() : Indent.getNormalIndent();
        } else {
          childIndent = Indent.getNoneIndent();
        }
      }
    }
    // Note that colons are aligned together with bounds and stride
    else if (myNode.getElementType() == PyElementTypes.SLICE_ITEM) {
      childAlignment = getChildAlignment();
    }
    else if (parentType == PyElementTypes.GENERATOR_EXPRESSION || parentType == PyElementTypes.PARENTHESIZED_EXPRESSION) {
      if (childType == PyElementTypes.TUPLE_EXPRESSION) {
        if (isIndentNext(child.getTreeParent())) {
          childIndent = Indent.getNormalIndent();
        }
      } else {
        final boolean tupleOrGenerator = parentType == PyElementTypes.GENERATOR_EXPRESSION ||
                                         myNode.getPsi(PyAstParenthesizedExpression.class)
                                           .getContainedExpression() instanceof PyAstTupleExpression;
        if ((childType == PyTokenTypes.RPAR && !(tupleOrGenerator && settings.HANG_CLOSING_BRACKETS)) ||
            !hasLineBreaksBeforeInSameParent(child, 1)) {
          childIndent = Indent.getNoneIndent();
        }
        // Operands and operators of a binary expression have their own indent, no need to increase it by indenting the expression itself
        else if (childType == PyElementTypes.BINARY_EXPRESSION && parentType == PyElementTypes.PARENTHESIZED_EXPRESSION) {
          childIndent = Indent.getNoneIndent();
        }
        else {
          final boolean useWiderIndent = isIndentNext(child) || settings.USE_CONTINUATION_INDENT_FOR_COLLECTION_AND_COMPREHENSIONS;
          childIndent = useWiderIndent ? Indent.getContinuationIndent() : Indent.getNormalIndent();
        }
      }
    }
    else if (parentType == PyElementTypes.ARGUMENT_LIST ||
             parentType == PyElementTypes.PATTERN_ARGUMENT_LIST ||
             parentType == PyElementTypes.PARAMETER_LIST) {
      if (childType == PyTokenTypes.RPAR && !settings.HANG_CLOSING_BRACKETS) {
        childIndent = Indent.getNoneIndent();
      }
      else if ((parentType == PyElementTypes.PARAMETER_LIST && settings.USE_CONTINUATION_INDENT_FOR_PARAMETERS) ||
               (parentType == PyElementTypes.ARGUMENT_LIST && settings.USE_CONTINUATION_INDENT_FOR_ARGUMENTS) ||
               argumentMayHaveSameIndentAsFollowingStatementList()) {
        childIndent = Indent.getContinuationIndent();
      }
      else {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (parentType == PyElementTypes.SUBSCRIPTION_EXPRESSION) {
      final PyAstExpression indexExpression = ((PyAstSubscriptionExpression)myNode.getPsi()).getIndexExpression();
      if (indexExpression != null && child == indexExpression.getNode()) {
        childIndent = Indent.getNormalIndent();
      }
    }
    else if (parentType == PyElementTypes.REFERENCE_EXPRESSION) {
      if (child != myNode.getFirstChildNode()) {
        childIndent = Indent.getNoneIndent();
        if (hasLineBreaksBeforeInSameParent(child, 1)) {
          if (isInControlStatement()) {
            childIndent = Indent.getContinuationIndent();
          }
          else {
            PyBlock b = myParent;
            while (b != null) {
              if (b.getNode().getPsi() instanceof PyAstParenthesizedExpression ||
                  b.getNode().getPsi() instanceof PyAstArgumentList ||
                  b.getNode().getPsi() instanceof PyAstParameterList) {
                childAlignment = getAlignmentOfChild(b, 1);
                break;
              }
              b = b.myParent;
            }
          }
        }
      }
    }
    childWrap = getSpecialWrapForContainers(child, childType, childWrap, parentType);

    if (isAfterStatementList(child) &&
        !hasLineBreaksBeforeInSameParent(child, 2) &&
        child.getElementType() != PyTokenTypes.END_OF_LINE_COMMENT) {
      // maybe enter was pressed and cut us from a previous (nested) statement list
      childIndent = Indent.getNormalIndent();
    }

    if (settings.DICT_ALIGNMENT == DICT_ALIGNMENT_ON_VALUE) {
      // Align not the whole value but its left bracket if it starts with it
      if (isValueOfKeyValuePairOfDictLiteral(child) && !isOpeningBracket(child.getFirstChildNode())) {
        childAlignment = myParent.myDictAlignment;
      }
      else if (isValueOfKeyValuePairOfDictLiteral(myNode) && isOpeningBracket(child)) {
        childAlignment = myParent.myParent.myDictAlignment;
      }
    }
    else if (myContext.getPySettings().DICT_ALIGNMENT == DICT_ALIGNMENT_ON_COLON) {
      if (isChildOfKeyValuePairOfDictLiteral(child) && childType == PyTokenTypes.COLON) {
        childAlignment = myParent.myDictAlignment;
      }
    }

    if (parentType == PyElementTypes.WITH_STATEMENT && isInsideWithStatementParentheses(myNode, child)) {
      if (needListAlignment(child)) {
        childAlignment = getAlignmentForChildren();
      }
      else {
        childIndent = Indent.getNormalIndent();
      }
      if (childType == PyTokenTypes.RPAR && !settings.HANG_CLOSING_BRACKETS) {
        childIndent = Indent.getNoneIndent();
      }
    }

    ASTNode prev = child.getTreePrev();
    while (prev != null && prev.getElementType() == TokenType.WHITE_SPACE) {
      if (prev.textContains('\\') &&
          !childIndent.equals(Indent.getContinuationIndent(false)) &&
          !childIndent.equals(Indent.getContinuationIndent(true))) {
        childIndent = isIndentNext(child) ? Indent.getContinuationIndent() : Indent.getNormalIndent();
        break;
      }
      prev = prev.getTreePrev();
    }

    // Don't wrap anything inside f-string fragments
    if (TreeUtil.findParent(child, PyElementTypes.FSTRING_FRAGMENT) != null) {
      childWrap = Wrap.createWrap(WrapType.NONE, false);
    }

    return new PyBlock(this, child, childAlignment, childIndent, childWrap, myContext);
  }

  private @Nullable Wrap getSpecialWrapForContainers(@NotNull ASTNode child,
                                           @NotNull IElementType childType,
                                           @Nullable Wrap childWrap,
                                           @Nullable IElementType parentType) {
    if (childType != PyTokenTypes.COMMA && childType != PyTokenTypes.END_OF_LINE_COMMENT) {
      // (...)
      if (childType != PyTokenTypes.LPAR && childType != PyTokenTypes.RPAR) {
        if (parentType == PyElementTypes.PARAMETER_LIST) return myParameterListWrapping;
        if (parentType == PyElementTypes.ARGUMENT_LIST) return myArgumentListWrapping;
        if (parentType == PyElementTypes.TUPLE_EXPRESSION && !isInsideSubscriptionExpression(child)) return myTupleWrapping;
      }
      // [...]
      if (childType != PyTokenTypes.LBRACKET && childType != PyTokenTypes.RBRACKET) {
        if (parentType == PyElementTypes.LIST_LITERAL_EXPRESSION && !isInsideSubscriptionExpression(child)) return myListWrapping;
      }
      // {...}
      if (childType != PyTokenTypes.LBRACE && childType != PyTokenTypes.RBRACE) {
        if (parentType == PyElementTypes.SET_LITERAL_EXPRESSION) return mySetWrapping;
      }
    }
    if (childType == PyElementTypes.KEY_VALUE_EXPRESSION && isChildOfDictLiteral(child)) {
      return myDictWrapping;
    }
    return childWrap;
  }

  private static boolean isInsideSubscriptionExpression(ASTNode child) {
    return PsiTreeUtil.getParentOfType(child.getPsi(), PyAstSubscriptionExpression.class) != null;
  }

  private static boolean isInsideWithStatementParentheses(@NotNull ASTNode withStatement, @NotNull ASTNode node) {
    ASTNode openingParenthesis = withStatement.findChildByType(PyTokenTypes.LPAR);
    if (openingParenthesis == null) {
      return false;
    }
    if (node.getStartOffset() < openingParenthesis.getStartOffset()) {
      return false;
    }
    ASTNode closingParenthesis = withStatement.findChildByType(PyTokenTypes.RPAR);
    if (closingParenthesis != null) {
      return node.getStartOffset() <= closingParenthesis.getStartOffset();
    }
    ASTNode afterParentheses = ObjectUtils.chooseNotNull(withStatement.findChildByType(PyTokenTypes.COLON),
                                                         withStatement.findChildByType(PyElementTypes.STATEMENT_LIST));
    return afterParentheses == null || node.getStartOffset() < afterParentheses.getStartOffset();
  }

  private static boolean isIfCondition(@NotNull ASTNode node) {
    @NotNull PsiElement element = node.getPsi();
    final PyAstIfPart ifPart = ObjectUtils.tryCast(element.getParent(), PyAstIfPart.class);
    return ifPart != null && ifPart.getCondition() == element && !ifPart.isElif();
  }

  private static boolean isCondition(@NotNull ASTNode node) {
    final PsiElement element = node.getPsi();
    final PyAstConditionalStatementPart conditionalStatement =
      ObjectUtils.tryCast(element.getParent(), PyAstConditionalStatementPart.class);
    return conditionalStatement != null && conditionalStatement.getCondition() == element;
  }

  private @Nullable PyBlock findTopmostBinaryExpressionBlock(@NotNull ASTNode child) {
    assert child.getElementType() != PyElementTypes.BINARY_EXPRESSION;
    PyBlock parentBlock = this;
    PyBlock alignmentOwner = null;
    while (parentBlock != null && parentBlock.myNode.getElementType() == PyElementTypes.BINARY_EXPRESSION) {
      alignmentOwner = parentBlock;
      parentBlock = parentBlock.myParent;
    }
    return alignmentOwner;
  }

  private static boolean isOpeningBracket(@Nullable ASTNode node) {
    return node != null && PyTokenTypes.OPEN_BRACES.contains(node.getElementType()) && node == node.getTreeParent().getFirstChildNode();
  }

  private static boolean isValueOfKeyValuePairOfDictLiteral(@NotNull ASTNode node) {
    return isValueOfKeyValuePair(node) && isChildOfDictLiteral(node.getTreeParent());
  }

  private static boolean isChildOfKeyValuePairOfDictLiteral(@NotNull ASTNode node) {
    return isChildOfKeyValuePair(node) && isChildOfDictLiteral(node.getTreeParent());
  }

  private static boolean isChildOfDictLiteral(@NotNull ASTNode node) {
    final ASTNode nodeParent = node.getTreeParent();
    return nodeParent != null && nodeParent.getElementType() == PyElementTypes.DICT_LITERAL_EXPRESSION;
  }

  private static boolean isChildOfKeyValuePair(@NotNull ASTNode node) {
    final ASTNode nodeParent = node.getTreeParent();
    return nodeParent != null && nodeParent.getElementType() == PyElementTypes.KEY_VALUE_EXPRESSION;
  }

  private static boolean isValueOfKeyValuePair(@NotNull ASTNode node) {
    return isChildOfKeyValuePair(node) && node.getTreeParent().getPsi(PyAstKeyValueExpression.class).getValue() == node.getPsi();
  }

  private static boolean isEmptySequence(@NotNull ASTNode node) {
    return node.getPsi() instanceof PyAstSequenceExpression && ((PyAstSequenceExpression)node.getPsi()).isEmpty();
  }

  private boolean argumentMayHaveSameIndentAsFollowingStatementList() {
    if (myNode.getElementType() != PyElementTypes.ARGUMENT_LIST) {
      return false;
    }
    // This check is supposed to prevent PEP8's error: Continuation line with the same indent as next logical line
    final PsiElement header = getControlStatementHeader(myNode);
    if (header instanceof PyAstStatementListContainer) {
      final PyAstStatementList statementList = ((PyAstStatementListContainer)header).getStatementList();
      return PyUtilCore.onSameLine(header, myNode.getPsi()) && !PyUtilCore.onSameLine(header, statementList);
    }
    return false;
  }

  // Check https://www.python.org/dev/peps/pep-0008/#indentation
  private static boolean hasHangingIndent(@NotNull PsiElement elem) {
    if (elem instanceof PyAstCallExpression) {
      final PyAstArgumentList argumentList = ((PyAstCallExpression)elem).getArgumentList();
      return argumentList != null && hasHangingIndent(argumentList);
    }
    else if (elem instanceof PyAstFunction) {
      return hasHangingIndent(((PyAstFunction)elem).getParameterList());
    }

    final PsiElement firstChild;
    if (elem instanceof PyAstFromImportStatement) {
      firstChild = ((PyAstFromImportStatement)elem).getLeftParen();
    }
    else if (elem instanceof PyAstWithStatement) {
      firstChild = PyPsiUtilsCore.getFirstChildOfType(elem, PyTokenTypes.LPAR);
    }
    else {
      firstChild = elem.getFirstChild();
    }
    if (firstChild == null) {
      return false;
    }
    final IElementType elementType = elem.getNode().getElementType();
    final ASTNode firstChildNode = firstChild.getNode();
    if (ourHangingIndentOwners.contains(elementType) && PyTokenTypes.OPEN_BRACES.contains(firstChildNode.getElementType())) {
      if (hasLineBreakAfterIgnoringComments(firstChildNode)) {
        return true;
      }
      final PsiElement firstItem = getFirstItem(elem);
      if (firstItem == null) {
        return !PyTokenTypes.CLOSE_BRACES.contains(elem.getLastChild().getNode().getElementType());
      }
      else {
        if (firstItem instanceof PyAstNamedParameter) {
          final PyAstExpression defaultValue = ((PyAstNamedParameter)firstItem).getDefaultValue();
          return defaultValue != null && hasHangingIndent(defaultValue);
        }
        else if (firstItem instanceof PyAstKeywordArgument) {
          final PyAstExpression valueExpression = ((PyAstKeywordArgument)firstItem).getValueExpression();
          return valueExpression != null && hasHangingIndent(valueExpression);
        }
        else if (firstItem instanceof PyAstKeyValueExpression) {
          final PyAstExpression value = ((PyAstKeyValueExpression)firstItem).getValue();
          return value != null && hasHangingIndent(value);
        }
        else if (firstItem instanceof PyAstWithItem) {
          PyAstExpression contextExpression = ((PyAstWithItem)firstItem).getExpression();
          return hasHangingIndent(contextExpression);
        }
