// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package com.siyeh.ig.psiutils;

import com.intellij.codeInsight.BlockUtils;
import com.intellij.openapi.diagnostic.Attachment;
import com.intellij.openapi.diagnostic.RuntimeExceptionWithAttachments;
import com.intellij.openapi.util.Key;
import com.intellij.openapi.util.NlsSafe;
import com.intellij.psi.*;
import com.intellij.psi.tree.IElementType;
import com.intellij.psi.util.JavaPsiPatternUtil;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.ArrayUtil;
import com.intellij.util.ObjectUtils;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.NotNull;

import java.util.*;

public final class StatementExtractor {
  private static final Node EMPTY = new Node(null) {
    @Override
    public Node prepend(Node node) {
      return node;
    }

    @Override
    public String toString() {
      return "";
    }
  };

  /**
   * Generate statements from subexpressions of root expression which must be kept.
   *
   * @param expressionsToKeep list of expressions to keep. Each expression must be a descendant of root, they must be ordered inside list
   *                          in program order and cannot be ancestors of each other.
   * @param root a root expression
   * @return an array of non-physical statements which represent the same logic as passed expressions
   */
  public static PsiStatement @NotNull [] generateStatements(List<? extends PsiExpression> expressionsToKeep, PsiExpression root) {
    String statementsCode = generateStatementsText(expressionsToKeep, root);
    if(statementsCode.isEmpty()) return PsiStatement.EMPTY_ARRAY;
    PsiElementFactory factory = JavaPsiFacade.getElementFactory(root.getProject());
    PsiCodeBlock codeBlock = factory.createCodeBlockFromText("{" + statementsCode + "}", root);
    return codeBlock.getStatements();
  }

  public static String generateStatementsText(List<? extends PsiExpression> expressionsToKeep, PsiExpression root) {
    Node result = StreamEx.ofReversed(expressionsToKeep).map(expression -> createNode(expression, root)).foldLeft(EMPTY, Node::prepend);
    return result.toString();
  }

  private static @NotNull Node createNode(@NotNull PsiExpression expression, @NotNull PsiExpression root) {
    Node result = new Expr(expression);
    PsiExpression parent;
    while (expression != root) {
      PsiElement parentElement = expression.getParent();
      if (parentElement instanceof PsiExpressionList || parentElement instanceof PsiTemplate) {
        parentElement = parentElement.getParent();
      }
      if (parentElement instanceof PsiStatement) {
        PsiSwitchExpression switchExpression =
          PsiTreeUtil.getParentOfType(parentElement, PsiSwitchExpression.class, true, PsiMember.class, PsiLambdaExpression.class);
        if (switchExpression != null && PsiTreeUtil.isAncestor(root, switchExpression, false)) {
          boolean isYield =
            parentElement instanceof PsiYieldStatement && ((PsiYieldStatement)parentElement).findEnclosingExpression() == switchExpression;
          boolean isRuleExpression =
            parentElement instanceof PsiExpressionStatement && parentElement.getParent() instanceof PsiSwitchLabeledRuleStatement &&
            ((PsiSwitchLabeledRuleStatement)parentElement.getParent()).getEnclosingSwitchBlock() == switchExpression;
          if (isYield || isRuleExpression) {
            result = new Switch(switchExpression, Collections.singletonMap((PsiStatement)parentElement, result));
          }
          else {
            result = new Switch(switchExpression, Collections.emptyMap());
          }
          expression = switchExpression;
          continue;
        }
      }
      parent = ObjectUtils.tryCast(parentElement, PsiExpression.class);
      if (parent == null) {
        String message = PsiTreeUtil.isAncestor(root, expression, false) ?
                         "Expected to have expression parent, got " + parentElement.getClass() :
                         "Supplied root is not the expression ancestor";
        throw new RuntimeExceptionWithAttachments(message,
                                                  new Attachment("expression.txt", expression.getText()),
                                                  new Attachment("root.txt", root.getText()));
      }
      result = foldNode(result, expression, parent);
      expression = parent;
    }
    return result;
  }

  private static @NotNull Node foldNode(@NotNull Node node, @NotNull PsiExpression expression, @NotNull PsiExpression parent) {
    if (parent instanceof PsiPolyadicExpression polyadic) {
      IElementType type = polyadic.getOperationTokenType();
      boolean and;
      if (type == JavaTokenType.ANDAND) {
        and = true;
      }
      else if (type == JavaTokenType.OROR) {
        and = false;
      }
      else {
        return node;
      }
      PsiExpression[] operands = polyadic.getOperands();
      int index = ArrayUtil.indexOf(operands, expression);
      if (index == 0) return node;
      return new Cond(parent, parent, index, and ? node : EMPTY, and ? EMPTY : node);
    }
    if (parent instanceof PsiConditionalExpression ternary) {
      if (expression == ternary.getThenExpression()) {
        return new Cond(ternary, ternary.getCondition(), -1, node, EMPTY);
      }
      if (expression == ternary.getElseExpression()) {
        return new Cond(ternary, ternary.getCondition(), -1, EMPTY, node);
      }
    }
    return node;
  }

  private abstract static class Node {
    final PsiExpression myAnchor;

    protected Node(PsiExpression anchor) {
      myAnchor = anchor;
    }

    public abstract Node prepend(Node node);

    @Override
    public abstract String toString();
  }

  private static final class Cond extends Node {
    private final @NotNull PsiExpression myCondition;
    private final @NotNull Node myThenBranch;
    private final @NotNull Node myElseBranch;
    private final int myLimit;

    private Cond(@NotNull PsiExpression anchor,
                 @NotNull PsiExpression condition,
                 int limit,
                 @NotNull Node thenBranch,
                 @NotNull Node elseBranch) {
      super(anchor);
      myCondition = condition;
      myLimit = limit;
      assert limit < 0 || condition instanceof PsiPolyadicExpression;
      myThenBranch = thenBranch;
      myElseBranch = elseBranch;
    }

    private String getCondition(boolean invert) {
      if (myLimit < 0) {
        return invert ? BoolUtils.getNegatedExpressionText(myCondition) : myCondition.getText();
      }
      PsiPolyadicExpression condition = (PsiPolyadicExpression)myCondition;
      PsiExpression[] operands = condition.getOperands();
      String joiner = (condition.getOperationTokenType() == JavaTokenType.ANDAND) != invert ? "&&" : "||";
      return StreamEx.of(operands, 0, myLimit).map(invert ? BoolUtils::getNegatedExpressionText
                                                          : PsiExpression::getText)
        .joining(joiner);
    }

    @Override
    public @NlsSafe String toString() {
      if (myThenBranch == EMPTY) {
        return "if(" + getCondition(true) + ") {" + myElseBranch + "}";
      }
      return "if(" + getCondition(false) + ") {" + myThenBranch + "}" + (myElseBranch == EMPTY ? "" : "else {" + myElseBranch + "}");
    }

    @Override
    public Node prepend(Node node) {
      PsiExpression thatAnchor = node.myAnchor;
      if(thatAnchor == null) return this;
      if(thatAnchor == myAnchor) {
        assert node instanceof Cond;
        Cond cond = (Cond)node;
        assert myCondition == cond.myCondition;
        if(myLimit == cond.myLimit) {
          return new Cond(myAnchor, myCondition, myLimit, myThenBranch.prepend(cond.myThenBranch), myElseBranch.prepend(cond.myElseBranch));
        }
        assert myLimit > cond.myLimit;
        return this;
      }
      if(PsiTreeUtil.isAncestor(myCondition, thatAnchor, false)) {
        return this;
      }
      return new Cons(node, this);
    }
  }

  private static final class Expr extends Node {
    private Expr(@NotNull PsiExpression expression) {
      super(expression);
    }

    @Override
    public Node prepend(Node node) {
      return node.myAnchor == null ? this : new Cons(node, this);
    }

    @Override
    public String toString() {
      if (myAnchor instanceof PsiInstanceOfExpression) {
        var patternVariableWrappers = JavaPsiPatternUtil.collectPatternVariableWrappers(myAnchor);
        StringBuilder sb = new StringBuilder();
        for (var patternVariableWrapper : patternVariableWrappers) {
          String initializer = patternVariableWrapper.getEffectiveInitializerText();
          if (initializer != null) {
            PsiPatternVariable variable = patternVariableWrapper.getVariable();
            sb.append(variable.getTypeElement().getText()).append(" ").append(variable.getName()).append("=")
              .append(initializer).append(";");
          }
        }
        return sb.toString();
      }
      return myAnchor.getText() + ";";
    }
  }

  private static final class Switch extends Node {
    private static final Key<Node> NODE_KEY = Key.create("SwitchNode");

    private final @NotNull Map<PsiStatement, Node> myReturns;

    private Switch(@NotNull PsiSwitchExpression expression, @NotNull Map<PsiStatement, Node> sideEffectReturns) {
      super(expression);
      myReturns = sideEffectReturns;
    }

    @Override
    public Node prepend(Node node) {
      if (node.myAnchor == null) return this;
      if (node instanceof Switch && node.myAnchor == myAnchor) {
        if (myReturns.isEmpty()) return node;
        if (((Switch)node).myReturns.isEmpty()) return this;
        Map<PsiStatement, Node> newMap = new HashMap<>(myReturns);
        ((Switch)node).myReturns.forEach((statement, n) -> newMap.merge(statement, n, Node::prepend));
        return new Switch((PsiSwitchExpression)myAnchor, newMap);
      }
      return new Cons(node, this);
    }

    @Override
    public String toString() {
      myReturns.forEach((statement, node) -> statement.putCopyableUserData(NODE_KEY, node));
      PsiSwitchExpression copy = (PsiSwitchExpression)myAnchor.copy();
      Map<PsiStatement, PsiStatement[]> replacementMap = new HashMap<>();
      PsiElementFactory factory = JavaPsiFacade.getElementFactory(myAnchor.getProject());
      PsiCodeBlock body = Objects.requireNonNull(copy.getBody());
      body.accept(new JavaRecursiveElementWalkingVisitor() {
        @Override
        public void visitExpressionStatement(@NotNull PsiExpressionStatement statement) {
          if (statement.getParent() instanceof PsiSwitchLabeledRuleStatement &&
              ((PsiSwitchLabeledRuleStatement)statement.getParent()).getEnclosingSwitchBlock() == copy) {
            process(statement);
          }
        }

        @Override
        public void visitYieldStatement(@NotNull PsiYieldStatement statement) {
          if (statement.getExpression() != null && statement.findEnclosingExpression() == copy) {
            process(statement);
          }
        }

        @Override
        public void visitExpression(@NotNull PsiExpression expression) {
          // Do not go into any expressions
        }

        private void process(PsiStatement statement) {
          Node data = statement.getCopyableUserData(NODE_KEY);
          if (data == null) {
            replacementMap.put(statement, PsiStatement.EMPTY_ARRAY);
          }
          else {
            replacementMap.put(statement, factory.createCodeBlockFromText("{" + data + "}", statement).getStatements());
          }
        }
      });
      replacementMap.forEach((statement, replacements) -> {
        boolean keep = statement instanceof PsiYieldStatement && shouldKeepBreak(statement);
        if (!keep && replacements.length == 1) {
          statement.replace(replacements[0]);
        }
        else {
          if (!keep || replacements.length > 0) {
            if (!(statement.getParent() instanceof PsiCodeBlock)) {
              statement = BlockUtils.expandSingleStatementToBlockStatement(statement);
            }
            PsiElement parent = statement.getParent();
            for (PsiStatement replacement : replacements) {
              parent.addBefore(replacement, statement);
            }
          }
          if (!keep) {
            statement.delete();
          }
          else {
            statement.replace(factory.createStatementFromText("break;", null));
          }
        }
      });
      return copy.getText();
    }

    public boolean shouldKeepBreak(PsiStatement statement) {
      if (PsiTreeUtil.skipWhitespacesAndCommentsForward(statement) instanceof PsiStatement) {
        return true;
      }
      PsiElement parent = statement.getParent();
      if (parent instanceof PsiCodeBlock) {
        PsiElement gParent = parent.getParent();
        if (gParent instanceof PsiBlockStatement) {
          return shouldKeepBreak((PsiStatement)gParent);
        }
      }
      else if (parent instanceof PsiLabeledStatement || parent instanceof PsiIfStatement) {
        return shouldKeepBreak((PsiStatement)parent);
      }
      else if (parent instanceof PsiSwitchLabeledRuleStatement) {
        return false;
      }
      return true;
    }
  }

  private static final class Cons extends Node {
    private final @NotNull Node myHead;
    private final @NotNull Node myTail;

    private Cons(@NotNull Node head, @NotNull Node tail) {
      super(head.myAnchor);
      assert !(head instanceof Cons);
      myHead = head;
      myTail = tail;
    }

    @Override
    public Node prepend(Node node) {
      if(node.myAnchor == null) return this;
      if(PsiTreeUtil.isAncestor(myHead.myAnchor, node.myAnchor, false)) {
        Node newHead = myHead.prepend(node);
        return new Cons(newHead, myTail);
      }
      return new Cons(node, this);
    }

    @Override
    public String toString() {
      return myHead.toString() + myTail;
    }
  }
}
