/*
 * Copyright 2000-2017 JetBrains s.r.o.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.jetbrains.python.codeInsight.controlflow;

import com.google.common.collect.Lists;
import com.intellij.codeInsight.controlflow.ControlFlowBuilder;
import com.intellij.codeInsight.controlflow.Instruction;
import com.intellij.codeInsight.controlflow.TransparentInstruction;
import com.intellij.codeInsight.controlflow.impl.ConditionalInstructionImpl;
import com.intellij.codeInsight.controlflow.impl.TransparentInstructionImpl;
import com.intellij.openapi.util.Pair;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiNamedElement;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.psi.util.QualifiedName;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.*;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.List;

public class PyControlFlowBuilder extends PyRecursiveElementVisitor {

  private final ControlFlowBuilder myBuilder = new ControlFlowBuilder();
  private final @Nullable LanguageLevel myLanguageLevel;

  private @Nullable TrueFalseNodes myTrueFalseNodes;

  // see com.jetbrains.python.PyPatternTypeTest.testMatchClassPatternShadowingCapture
  private final @NotNull List<String> myPatternBindingNames = new ArrayList<>();

  public PyControlFlowBuilder(@Nullable LanguageLevel languageLevel) { myLanguageLevel = languageLevel; }

  private record TrueFalseNodes(@NotNull Instruction trueNode, @NotNull Instruction falseNode) {
  }

  public PyControlFlow buildControlFlow(final @NotNull ScopeOwner owner) {
    return new PyControlFlow(myBuilder.build(this, owner));
  }

  protected @NotNull ControlFlowBuilder getBuilder() {
    return myBuilder;
  }


  @Override
  public void visitPyFunction(final @NotNull PyFunction node) {
    // Create node and stop here
    myBuilder.startNode(node);

    visitParameterListExpressions(node.getParameterList());
    visitDecorators(node.getDecoratorList());
    if (node.getTypeParameterList() == null) {
      final PyAnnotation annotation = node.getAnnotation();
      if (annotation != null) {
        annotation.acceptChildren(this);
      }
    }

    final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, node, node.getName());
    myBuilder.addNodeAndCheckPending(instruction);
  }

  @Override
  public void visitPyDecoratorList(@NotNull PyDecoratorList node) {
  }

  private void visitDecorators(PyDecoratorList list) {
    if (list != null) {
      for (PyDecorator decorator : list.getDecorators()) {
        decorator.accept(this);
      }
    }
  }

  private void visitParameterListExpressions(PyParameterList parameterList) {
    ParamHelper.walkDownParamArray(parameterList.getParameters(), new ParamHelper.ParamVisitor() {
      @Override
      public void visitNamedParameter(PyNamedParameter param, boolean first, boolean last) {
        final PyExpression defaultValue = param.getDefaultValue();
        if (defaultValue != null) {
          defaultValue.accept(PyControlFlowBuilder.this);
        }
        if (parameterList.getParent() instanceof PyFunction function && function.getTypeParameterList() == null) {
          final PyAnnotation annotation = param.getAnnotation();
          if (annotation != null) {
            annotation.acceptChildren(PyControlFlowBuilder.this);
          }
        }
      }
    });
  }

  @Override
  public void visitPyClass(final @NotNull PyClass node) {
    // Create node and stop here
    myBuilder.startNode(node);

    for (PsiElement element : node.getSuperClassExpressions()) {
      element.accept(this);
    }
    visitDecorators(node.getDecoratorList());
    final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, node, node.getName());
    myBuilder.addNodeAndCheckPending(instruction);
  }

  @Override
  public void visitPyStatement(final @NotNull PyStatement node) {
    myBuilder.startNode(node);
    super.visitPyStatement(node);
  }

  @Override
  public void visitPyElement(@NotNull PyElement node) {
    if (node instanceof PsiNamedElement && !(node instanceof PyKeywordArgument)) {
      myBuilder.startNode(node);
      myBuilder.addNode(ReadWriteInstruction.newInstruction(myBuilder, node, node.getName(), ReadWriteInstruction.ACCESS.WRITE));
    }
    super.visitPyElement(node);
  }

  @Override
  public void visitPyCallExpression(final @NotNull PyCallExpression node) {
    super.visitPyCallExpression(node);

    var callInstruction = new CallInstruction(myBuilder, node);
    myBuilder.addNodeAndCheckPending(callInstruction);

    if (node.isCalleeText(PyNames.ASSERT_IS_INSTANCE)) {
      addTypeAssertionNodes(node, true);
    }
  }

  @Override
  public void visitPySubscriptionExpression(@NotNull PySubscriptionExpression node) {
    myBuilder.startNode(node);
    node.getOperand().accept(this);
    final PyExpression expression = node.getIndexExpression();
    if (expression != null) {
      expression.accept(this);
    }
  }

  @Override
  public void visitPyReferenceExpression(final @NotNull PyReferenceExpression node) {
    final PyExpression qualifier = node.getQualifier();
    if (qualifier != null) {
      qualifier.accept(this);
      return;
    }
    if (PyImportStatementNavigator.getImportStatementByElement(node) != null) {
      return;
    }

    final ReadWriteInstruction.ACCESS access = PyAugAssignmentStatementNavigator.getStatementByTarget(node) != null
                                               ? ReadWriteInstruction.ACCESS.READWRITE
                                               : ReadWriteInstruction.ACCESS.READ;
    final ReadWriteInstruction readWriteInstruction = ReadWriteInstruction.newInstruction(myBuilder, node, getName(node), access);
    myBuilder.addNodeAndCheckPending(readWriteInstruction);
  }

  @Override
  public void visitPyBoolLiteralExpression(@NotNull PyBoolLiteralExpression node) {
    final ReadWriteInstruction readWriteInstruction = ReadWriteInstruction.newInstruction(myBuilder, node, node.getText(),
                                                                                          ReadWriteInstruction.ACCESS.READ);
    myBuilder.addNodeAndCheckPending(readWriteInstruction);
  }

  @Override
  public void visitPyNoneLiteralExpression(@NotNull PyNoneLiteralExpression node) {
    final ReadWriteInstruction readWriteInstruction = ReadWriteInstruction.newInstruction(myBuilder, node, node.getText(),
                                                                                          ReadWriteInstruction.ACCESS.READ);
    myBuilder.addNodeAndCheckPending(readWriteInstruction);
  }

  @Override
  public void visitPyTypeDeclarationStatement(@NotNull PyTypeDeclarationStatement node) {
    myBuilder.startNode(node);
    final PyAnnotation annotation = node.getAnnotation();
    if (annotation != null) {
      annotation.accept(this);
    }
    node.getTarget().accept(this);
  }

  @Override
  public void visitPyAssignmentStatement(final @NotNull PyAssignmentStatement node) {
    myBuilder.startNode(node);
    final PyExpression value = node.getAssignedValue();
    if (value != null) {
      value.accept(this);
    }
    final PyAnnotation annotation = node.getAnnotation();
    if (annotation != null) {
      annotation.accept(this);
    }
    for (PyExpression expression : node.getRawTargets()) {
      expression.accept(this);
    }
  }

  @Override
  public void visitPyDelStatement(@NotNull PyDelStatement node) {
    myBuilder.startNode(node);
    for (PyExpression target : node.getTargets()) {
      if (target instanceof PyReferenceExpression expr) {
        PyExpression qualifier = expr.getQualifier();
        if (qualifier != null) {
          qualifier.accept(this);
        }
        else {
          myBuilder.addNode(ReadWriteInstruction.newInstruction(myBuilder, target, expr.getName(), ReadWriteInstruction.ACCESS.DELETE));
        }
      }
      else {
        target.accept(this);
      }
    }
  }

  @Override
  public void visitPyAugAssignmentStatement(final @NotNull PyAugAssignmentStatement node) {
    myBuilder.startNode(node);
    final PyExpression value = node.getValue();
    if (value != null) {
      value.accept(this);
    }
    node.getTarget().accept(this);
  }

  @Override
  public void visitPyTargetExpression(final @NotNull PyTargetExpression node) {
    final QualifiedName qName = node.asQualifiedName();
    if (qName != null) {
      final ReadWriteInstruction instruction = ReadWriteInstruction.newInstruction(myBuilder, node, qName.toString(),
                                                                                   ReadWriteInstruction.ACCESS.WRITE);
      myBuilder.addNodeAndCheckPending(instruction);
    }

    final PyExpression qualifier = node.getQualifier();
    if (qualifier != null) {
      qualifier.accept(this);
    }
  }

  @Override
  public void visitPyNamedParameter(final @NotNull PyNamedParameter node) {
    PyAnnotation annotation = node.getAnnotation();
    if (annotation != null) {
      annotation.accept(this);
    }
    final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, node, node.getName());
    myBuilder.addNodeAndCheckPending(instruction);
  }

  @Override
  public void visitPyAnnotation(@NotNull PyAnnotation node) {
    // Unless there is a type parameter list, return type and parameter annotations for functions are evaluated in their enclosing scope
    // and processed in visitPyFunction.
    // If there are type parameters, though, we need to put the corresponding instructions *inside* the function's scope to be able to
    // access them from annotations. 
    PyFunction function = PsiTreeUtil.getParentOfType(node, PyFunction.class, true, PyStatement.class);
    if (function == null || function.getTypeParameterList() != null) {
      super.visitPyAnnotation(node);
    }
  }

  @Override
  public void visitPyImportStatement(final @NotNull PyImportStatement node) {
    visitPyImportStatementBase(node);
  }

  @Override
  public void visitPyFromImportStatement(@NotNull PyFromImportStatement node) {
    visitPyImportStatementBase(node);
    final PyStarImportElement starImportElement = node.getStarImportElement();
    if (starImportElement != null) {
      starImportElement.accept(this);
    }
  }

  @Override
  public void visitPyStarImportElement(@NotNull PyStarImportElement node) {
    myBuilder.startNode(node);
  }

  private void visitPyImportStatementBase(PyImportStatementBase node) {
    myBuilder.startNode(node);
    for (PyImportElement importElement : node.getImportElements()) {
      final ReadWriteInstruction instruction = ReadWriteInstruction.write(myBuilder, importElement, importElement.getVisibleName());
      myBuilder.addNodeAndCheckPending(instruction);
    }
  }

  @Override
  public void visitPyConditionalExpression(@NotNull PyConditionalExpression node) {
    myBuilder.startNode(node);
    TransparentInstruction trueNode = addTransparentInstruction();
    TransparentInstruction falseNode = addTransparentInstruction();
    TransparentInstruction exitNode = addTransparentInstruction();
    final PyExpression condition = node.getCondition();
    if (condition != null) {
      visitCondition(condition, trueNode, falseNode);
    }
    final PyExpression truePart = node.getTruePart();
    final PyExpression falsePart = node.getFalsePart();
    if (truePart != null) {
      myBuilder.prevInstruction = trueNode;
      truePart.accept(this);
      myBuilder.addEdge(myBuilder.prevInstruction, exitNode);
    }
    if (falsePart != null) {
      myBuilder.prevInstruction = falseNode;
      falsePart.accept(this);
      myBuilder.addEdge(myBuilder.prevInstruction, exitNode);
    }
    myBuilder.prevInstruction = exitNode;
  }

  @Override
  public void visitPyMatchStatement(@NotNull PyMatchStatement matchStatement) {
    myBuilder.startNode(matchStatement);
    PyExpression subject = matchStatement.getSubject();
    if (subject != null) {
      subject.accept(this);
    }
    Instruction nextClause = myBuilder.prevInstruction;
    boolean unreachable = false;
    for (PyCaseClause clause : matchStatement.getCaseClauses()) {
      myBuilder.prevInstruction = nextClause;
      nextClause = addTransparentInstruction();

      myPatternBindingNames.clear();

      PyPattern pattern = clause.getPattern();
      if (pattern != null) {
        pattern.accept(this);
        addTypeAssertionNodes(clause, true, myPatternBindingNames);
      }

      PyExpression guard = clause.getGuardCondition();
      if (guard != null) {
        TransparentInstruction trueNode = addTransparentInstruction();
        visitCondition(guard, trueNode, nextClause);
        myBuilder.prevInstruction = trueNode;
        addTypeAssertionNodes(guard, true);
      }

      if (unreachable) {
        myBuilder.addNode(new PyUnreachableInstruction(myBuilder, false));
      }
      if (pattern != null && pattern.isIrrefutable() && (guard == null || PyEvaluator.evaluateAsBooleanNoResolve(guard, false))) {
        unreachable = true;
      }
      myBuilder.startNode(clause.getStatementList());
      clause.getStatementList().accept(this);
      myBuilder.addPendingEdge(matchStatement, myBuilder.prevInstruction);
      myBuilder.updatePendingElementScope(clause.getStatementList(), matchStatement);
    }
    myBuilder.prevInstruction = nextClause;
    myBuilder.addNodeAndCheckPending(new TransparentInstructionImpl(myBuilder, matchStatement, ""));
    if (!myBuilder.prevInstruction.allPred().isEmpty()) {
      addTypeAssertionNodes(matchStatement, false);
    }
    myBuilder.addPendingEdge(matchStatement, myBuilder.prevInstruction);
    myBuilder.prevInstruction = null;
  }

  @Override
  public void visitPyPattern(@NotNull PyPattern node) {
    boolean isRefutable = !node.isIrrefutable();
    if (isRefutable) {
      myBuilder.addNodeAndCheckPending(new RefutablePatternInstruction(myBuilder, node, false));
    }
    else {
      myBuilder.startNode(node);
    }
    myBuilder.addPendingEdge(node.getParent(), myBuilder.prevInstruction);

    node.acceptChildren(this);
    myBuilder.updatePendingElementScope(node, node.getParent());

    if (isRefutable) {
      myBuilder.addNode(new RefutablePatternInstruction(myBuilder, node, true));
    }
  }

  @Override
  public void visitPyOrPattern(@NotNull PyOrPattern node) {
    myBuilder.addNodeAndCheckPending(new RefutablePatternInstruction(myBuilder, node, false));

    TransparentInstruction onSuccess = new TransparentInstructionImpl(myBuilder, node, "onSuccess");
    List<PyPattern> alternatives = node.getAlternatives();
    PyPattern lastAlternative = ContainerUtil.getLastItem(alternatives);

    for (PyPattern alternative : alternatives) {
      alternative.accept(this);
      if (alternative != lastAlternative) {
        // Allow next alternative to handle the fail edge of this alternative
        myBuilder.updatePendingElementScope(node, alternative);
      }
      myBuilder.addEdge(myBuilder.prevInstruction, onSuccess);
      myBuilder.prevInstruction = null;
    }
    myBuilder.addNode(onSuccess);
    myBuilder.addNode(new RefutablePatternInstruction(myBuilder, node, true));
    myBuilder.updatePendingElementScope(node, node.getParent());
  }

  @Override
  public void visitPyClassPattern(@NotNull PyClassPattern node) {
    myBuilder.addNodeAndCheckPending(new RefutablePatternInstruction(myBuilder, node, false));

    node.getClassNameReference().accept(this);
    myBuilder.addPendingEdge(node.getParent(), myBuilder.prevInstruction);

    node.getArgumentList().acceptChildren(this);
    myBuilder.updatePendingElementScope(node, node.getParent());

    myBuilder.addNode(new RefutablePatternInstruction(myBuilder, node, true));
  }

  @Override
  public void visitPyValuePattern(@NotNull PyValuePattern node) {
    myBuilder.addNodeAndCheckPending(new RefutablePatternInstruction(myBuilder, node, false));

    node.getValue().accept(this);
    myBuilder.addPendingEdge(node.getParent(), myBuilder.prevInstruction);

    myBuilder.addNode(new RefutablePatternInstruction(myBuilder, node, true));
  }

  @Override
  public void visitPyAsPattern(@NotNull PyAsPattern node) {
    // AsPattern can't fail by itself – it fails only if its child fails.
    // So no need to create an additional fail edge
    myBuilder.startNode(node);
    node.acceptChildren(this);
    if (node.getTarget() != null) {
      myPatternBindingNames.add(node.getTarget().getName());
    }
    myBuilder.updatePendingElementScope(node, node.getParent());
  }

  @Override
  public void visitPyCapturePattern(@NotNull PyCapturePattern node) {
    node.acceptChildren(this);
    // Although capture pattern is irrefutable, I add fail edge
    // here to add some connection to the next case clause.
    // Perhaps this can be reworked and simplified later.
    myBuilder.addPendingEdge(node.getParent(), myBuilder.prevInstruction);
    myPatternBindingNames.add(node.getTarget().getName());
  }

  @Override
  public void visitPyGroupPattern(@NotNull PyGroupPattern node) {
    // GroupPattern can't fail by itself – it fails only if its child fails.
    // So no need to create an additional fail edge 
    // Also no need for a dedicated node for GroupPattern itself
    node.acceptChildren(this);
    myBuilder.updatePendingElementScope(node, node.getParent());
  }

  @Override
  public void visitPyIfStatement(final @NotNull PyIfStatement node) {
    myBuilder.startNode(node);

    List<Instruction> exitInstructions = new ArrayList<>();
    @Nullable PyExpression alwaysTrueCondition = null;
    for (PyIfPart ifPart : StreamEx.of(node.getIfPart()).append(node.getElifParts())) {
      TransparentInstruction thenNode = addTransparentInstruction();
      TransparentInstruction elseNode = addTransparentInstruction();
      PyExpression condition = ifPart.getCondition();
      if (condition != null) {
        visitCondition(condition, thenNode, elseNode);
      }
      myBuilder.prevInstruction = thenNode;

      Boolean conditionResult = PyEvaluator.evaluateAsBooleanNoResolve(condition, myLanguageLevel);
      boolean unreachable = alwaysTrueCondition != null || Boolean.FALSE.equals(conditionResult);
      if (unreachable) {
        boolean isUnreachableForTypeChecking;
        if (alwaysTrueCondition != null) {
          isUnreachableForTypeChecking = isTypeCheckingCondition(alwaysTrueCondition);
        }
        else {
          isUnreachableForTypeChecking = isTypeCheckingCondition(condition);
        }
        myBuilder.addNode(new PyUnreachableInstruction(myBuilder, isUnreachableForTypeChecking));
      }
      if (Boolean.TRUE.equals(conditionResult)) {
        alwaysTrueCondition = condition;
      }
      visitPyStatementPart(ifPart);

      if (!unreachable) {
        exitInstructions.add(myBuilder.prevInstruction);
      }
      myBuilder.prevInstruction = elseNode;
    }

    final PyElsePart elsePart = node.getElsePart();
    if (elsePart != null) {
      if (alwaysTrueCondition != null) {
        myBuilder.addNode(new PyUnreachableInstruction(myBuilder, isTypeCheckingCondition(alwaysTrueCondition)));
      }
      visitPyStatementPart(elsePart);
    }

    if (alwaysTrueCondition == null) {
      exitInstructions.add(myBuilder.prevInstruction);
    }
    myBuilder.prevInstruction = addTransparentInstruction(node);

    for (Instruction exitInstruction : Lists.reverse(exitInstructions)) {
      myBuilder.addEdge(exitInstruction, myBuilder.prevInstruction);
    }
  }

  private static boolean isTypeCheckingCondition(@Nullable PyExpression expression) {
    return isTypeCheckingCheck(expression) || isVersionCheck(expression);
  }

  private static boolean isTypeCheckingCheck(@Nullable PyExpression expression) {
    expression = PyPsiUtils.flattenParens(expression);
    if (expression instanceof PyPrefixExpression prefixExpression && prefixExpression.getOperator() == PyTokenTypes.NOT_KEYWORD) {
      return isTypeCheckingCondition(prefixExpression.getOperand());
    }
    return PyEvaluator.isTypeCheckingExpression(expression);
  }

  private static boolean isVersionCheck(@Nullable PyExpression expression) {
    expression = PyPsiUtils.flattenParens(expression);
    if (expression instanceof PyPrefixExpression prefixExpression && prefixExpression.getOperator() == PyTokenTypes.NOT_KEYWORD) {
      return isVersionCheck(prefixExpression.getOperand());
    }
    if (expression instanceof PyBinaryExpression binaryExpression) {
      PyElementType op = binaryExpression.getOperator();
      if (PyTokenTypes.AND_KEYWORD.equals(op) || PyTokenTypes.OR_KEYWORD.equals(op)) {
        return isVersionCheck(binaryExpression.getLeftExpression()) &&
               isVersionCheck(binaryExpression.getRightExpression());
      }
      if (PyTokenTypes.RELATIONAL_OPERATIONS.contains(op)) {
        if (PyEvaluator.isSysVersionInfoExpression(binaryExpression.getLeftExpression()) &&
            PyPsiUtils.flattenParens(binaryExpression.getRightExpression()) instanceof PyTupleExpression) {
          return true;
        }
        if (PyPsiUtils.flattenParens(binaryExpression.getLeftExpression()) instanceof PyTupleExpression &&
            PyEvaluator.isSysVersionInfoExpression(binaryExpression.getRightExpression())) {
          return true;
        }
      }
    }
    return false;
  }

  @Override
  public void visitPyPrefixExpression(@NotNull PyPrefixExpression node) {
    if (myTrueFalseNodes != null && node.getOperator() == PyTokenTypes.NOT_KEYWORD) {
      PyExpression operand = node.getOperand();
      if (operand != null) {
        visitCondition(operand, myTrueFalseNodes.falseNode, myTrueFalseNodes.trueNode);
      }
    }
    else {
      super.visitPyPrefixExpression(node);
    }
  }

  @Override
  public void visitPyBinaryExpression(@NotNull PyBinaryExpression node) {
    PyElementType operator = node.getOperator();
    if (operator == PyTokenTypes.AND_KEYWORD || operator == PyTokenTypes.OR_KEYWORD) {
      final PyExpression left = node.getLeftExpression();
      final PyExpression right = node.getRightExpression();
      if (left == null || right == null) return;

      myBuilder.startNode(node);

      Instruction trueNode;
      Instruction falseNode;
      Instruction exitNode;
      if (myTrueFalseNodes == null) {
        trueNode = falseNode = exitNode = addTransparentInstruction();
      }
      else {
        trueNode = myTrueFalseNodes.trueNode;
        falseNode = myTrueFalseNodes.falseNode;
        exitNode = null;
      }

      Instruction rightNode = addTransparentInstruction();
      if (operator == PyTokenTypes.AND_KEYWORD) {
        visitCondition(left, rightNode, falseNode);
      }
      else {
        visitCondition(left, trueNode, rightNode);
      }

      myBuilder.prevInstruction = rightNode;
      visitCondition(right, trueNode, falseNode);

      if (exitNode != null) {
        myBuilder.prevInstruction = exitNode;
      }
    }
    else {
      super.visitPyBinaryExpression(node);
    }
  }

  @Override
  public void visitPyWhileStatement(final @NotNull PyWhileStatement node) {
    Instruction entryNode = myBuilder.startNode(node);

    TransparentInstruction thenNode = addTransparentInstruction();
    TransparentInstruction elseNode = addTransparentInstruction();

    final PyConditionalStatementPart whilePart = node.getWhilePart();
    final PyExpression condition = whilePart.getCondition();
    if (condition != null) {
      visitCondition(condition, thenNode, elseNode);
    }

    final Boolean conditionResult = PyEvaluator.evaluateAsBooleanNoResolve(condition);

    myBuilder.prevInstruction = Boolean.FALSE.equals(conditionResult) ? null : thenNode;
    visitPyStatementPart(whilePart);

    myBuilder.checkPending(entryNode);
    myBuilder.addEdge(myBuilder.prevInstruction, entryNode);

    myBuilder.prevInstruction = Boolean.TRUE.equals(conditionResult) ? null : elseNode;
    PyElsePart elsePart = node.getElsePart();
    if (elsePart != null) {
      visitPyStatementPart(elsePart);
    }

    collectInternalPendingEdges(node);
  }

  @Override
  public void visitPyForStatement(final @NotNull PyForStatement node) {
    myBuilder.startNode(node);
    final PyForPart forPart = node.getForPart();
    final PyExpression source = forPart.getSource();
    if (source != null) {
      source.accept(this);
    }
    final Instruction head = myBuilder.prevInstruction;
    final PyElsePart elsePart = node.getElsePart();
    if (elsePart == null && !loopHasAtLeastOneIteration(node)) {
      myBuilder.addPendingEdge(node, myBuilder.prevInstruction);
    }
    final PyStatementList list = forPart.getStatementList();
    final Instruction body;
    final PyExpression target = forPart.getTarget();
    if (target != null) {
      body = myBuilder.startNode(target);
      target.accept(this);
    }
    else {
      body = myBuilder.startNode(list);
    }
    list.accept(this);
    if (myBuilder.prevInstruction != null) {
      myBuilder.addEdge(myBuilder.prevInstruction, body);  //loop
      myBuilder.addPendingEdge(list, myBuilder.prevInstruction); // exit
    }
    myBuilder.processPending((pendingScope, instruction) -> {
      if (pendingScope != null && PsiTreeUtil.isAncestor(list, pendingScope, false)) {
        myBuilder.addEdge(instruction, body);  //loop
        myBuilder.addPendingEdge(list, instruction); // exit
      }
      else {
        myBuilder.addPendingEdge(pendingScope, instruction);
      }
    });
    myBuilder.prevInstruction = head;
    if (elsePart != null) {
      elsePart.accept(this);
      myBuilder.addPendingEdge(node, myBuilder.prevInstruction);
    }


    myBuilder.flowAbrupted();
    collectInternalPendingEdges(node);
  }

  private static boolean loopHasAtLeastOneIteration(@NotNull PyLoopStatement loopStatement) {
    final PyExpression expression = loopStatement instanceof PyForStatement
                                    ? ((PyForStatement)loopStatement).getForPart().getSource()
                                    : loopStatement instanceof PyWhileStatement
                                      ? ((PyWhileStatement)loopStatement).getWhilePart().getCondition()
                                      : null;

    return PyEvaluator.evaluateAsBooleanNoResolve(expression, false);
  }

  @Override
  public void visitPyBreakStatement(final @NotNull PyBreakStatement node) {
    myBuilder.startNode(node);
    final PyLoopStatement loop = node.getLoopStatement();
    if (loop != null) {
      myBuilder.addPendingEdge(loop, myBuilder.prevInstruction);
    }
    else {
      myBuilder.addPendingEdge(null, myBuilder.prevInstruction);
    }
    myBuilder.flowAbrupted();
  }

  @Override
  public void visitPyContinueStatement(final @NotNull PyContinueStatement node) {
    myBuilder.startNode(node);
    final PyLoopStatement loop = node.getLoopStatement();
    if (loop != null) {
      final Instruction instruction = myBuilder.findInstructionByElement(loop);
      if (instruction != null) {
        myBuilder.addEdge(myBuilder.prevInstruction, instruction);
      }
      else {
        myBuilder.addPendingEdge(null, null);
      }

      // There is no edge between loop statement and next after loop instruction
      // when loop has at least one iteration
      // so `continue` is marked as one more last instruction in the loop
      // see visitPyWhileStatement
      // see visitPyForStatement
      if (loopHasAtLeastOneIteration(loop)) {
        myBuilder.addPendingEdge(loop, myBuilder.prevInstruction);
      }
    }
    myBuilder.flowAbrupted();
  }

  @Override
  public void visitPyYieldExpression(@NotNull PyYieldExpression node) {
    myBuilder.startNode(node);
    final PyExpression expression = node.getExpression();
    if (expression != null) {
      expression.accept(this);
    }
  }

  @Override
  public void visitPyRaiseStatement(final @NotNull PyRaiseStatement node) {
    myBuilder.startNode(node);
    final PyExpression[] expressions = node.getExpressions();
    for (PyExpression expression : expressions) {
      expression.accept(this);
    }
    myBuilder.addNode(new PyRaiseInstruction(myBuilder, node));
    myBuilder.addPendingEdge(null, myBuilder.prevInstruction);
    myBuilder.flowAbrupted();
  }

  @Override
  public void visitPyReturnStatement(final @NotNull PyReturnStatement node) {
    myBuilder.startNode(node);
    final PyExpression expression = node.getExpression();
    if (expression != null) {
      expression.accept(this);
    }
    myBuilder.addPendingEdge(null, myBuilder.prevInstruction);
    myBuilder.flowAbrupted();
  }

  @Override
  public void visitPyTryExceptStatement(final @NotNull PyTryExceptStatement node) {
    myBuilder.startNode(node);

    // Process try part
    final PyTryPart tryPart = node.getTryPart();
    myBuilder.startNode(tryPart);
    tryPart.accept(this);

    // Goto else part after execution, or exit
    final PyElsePart elsePart = node.getElsePart();
    if (elsePart != null) {
      myBuilder.startNode(elsePart);
      elsePart.accept(this);
    }
    myBuilder.addPendingEdge(node, myBuilder.prevInstruction);

    // Process except parts
    final List<Instruction> exceptInstructions = new ArrayList<>();
    List<Pair<PsiElement, Instruction>> pendingBackup = new ArrayList<>();
    for (PyExceptPart exceptPart : node.getExceptParts()) {
      pendingBackup.addAll(myBuilder.pending);
      myBuilder.pending = new ArrayList<>();
      myBuilder.flowAbrupted();
      final Instruction exceptInstruction = myBuilder.startNode(exceptPart);
      exceptPart.accept(this);
      myBuilder.addPendingEdge(node, myBuilder.prevInstruction);
      exceptInstructions.add(exceptInstruction);
    }
    for (Pair<PsiElement, Instruction> pair : pendingBackup) {
      myBuilder.addPendingEdge(pair.first, pair.second);
    }

    final List<Pair<PsiElement, Instruction>> pendingNormalExits = new ArrayList<>();
    final PyFinallyPart finallyPart = node.getFinallyPart();
    final Instruction finallyFailInstruction;

    // Store pending normal exit instructions from try-except-else parts
    if (finallyPart != null) {
      myBuilder.processPending((pendingScope, instruction) -> {
        final PsiElement pendingElement = instruction.getElement();
        if (pendingElement != null) {
          final boolean isPending = PsiTreeUtil.isAncestor(node, pendingElement, false) &&
                                    !PsiTreeUtil.isAncestor(finallyPart, pendingElement, false);
          if (isPending && pendingScope != null) {
            pendingNormalExits.add(Pair.createNonNull(pendingScope, instruction));
          }
          else {
            myBuilder.addPendingEdge(pendingScope, instruction);
          }
        }
      });
    }

    // Finally-fail part handling
    if (finallyPart != null) {
      myBuilder.flowAbrupted();
      finallyFailInstruction = myBuilder.startNode(finallyPart);
      finallyPart.accept(this);
      myBuilder.addNodeAndCheckPending(new PyFinallyFailExitInstruction(myBuilder, finallyFailInstruction));
      myBuilder.addPendingEdge(null, myBuilder.prevInstruction);
      myBuilder.flowAbrupted();
    }
    else {
      finallyFailInstruction = null;
    }

    // Create exception edges
    for (Instruction instruction : myBuilder.instructions) {
      final PsiElement e = instruction.getElement();
      if (e == null || !canRaiseExceptions(instruction)) {
        continue;
      }
      // All instructions inside the try part have edges to except and finally parts
      if (PsiTreeUtil.getParentOfType(e, PyTryPart.class, false) == tryPart) {
        for (Instruction inst : exceptInstructions) {
          myBuilder.addEdge(instruction, inst);
        }
        if (finallyPart != null) {
          myBuilder.addEdge(instruction, finallyFailInstruction);
        }
      }
      if (finallyPart != null) {
        // All instructions inside except parts have edges to the finally part
        for (PyExceptPart exceptPart : node.getExceptParts()) {
