// Copyright 2000-2018 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.inspections.quickfix;

import com.intellij.modcommand.ModPsiUpdater;
import com.intellij.modcommand.PsiUpdateModCommandQuickFix;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.Couple;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.psi.PsiElement;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.PyPsiBundle;
import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import com.jetbrains.python.refactoring.PyPsiRefactoringUtil;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.*;

/**
 * For:
 * class B(A):
 * def __init__(self):
 * A.__init__(self)           #  inserted
 * print "Constructor B was called"
 * <p/>
 * User: catherine
 */
public class AddCallSuperQuickFix extends PsiUpdateModCommandQuickFix {

  @Override
  public @NotNull String getFamilyName() {
    return PyPsiBundle.message("QFIX.add.super");
  }

  @Override
  public void applyFix(@NotNull Project project, @NotNull PsiElement element, @NotNull ModPsiUpdater updater) {
    final PyFunction problemFunction = PsiTreeUtil.getParentOfType(element, PyFunction.class);
    if (problemFunction == null) return;
    final StringBuilder superCall = new StringBuilder();
    final PyClass klass = problemFunction.getContainingClass();
    if (klass == null) return;
    final PyClass[] superClasses = klass.getSuperClasses(null);
    if (superClasses.length == 0) return;

    final PyClass superClass = superClasses[0];
    final PyFunction superInit = superClass.findMethodByName(PyNames.INIT, true, null);
    if (superInit == null) return;

    final ParametersInfo origInfo = new ParametersInfo(problemFunction.getParameterList());
    final ParametersInfo superInfo = new ParametersInfo(superInit.getParameterList());
    final boolean addSelfToCall;

    if (klass.isNewStyleClass(null)) {
      addSelfToCall = false;
      if (LanguageLevel.forElement(klass).isPy3K()) {
        superCall.append("super().__init__(");
      }
      else {
        superCall.append("super(").append(klass.getName()).append(", ").append(getSelfParameterName(origInfo)).append(").__init__(");
      }
    }
    else {
      addSelfToCall = true;
      superCall.append(superClass.getName()).append(".__init__(");
    }

    final Couple<List<String>> couple = buildNewFunctionParamsAndSuperInitCallArgs(origInfo, superInfo, addSelfToCall);
    final StringBuilder newParameters = new StringBuilder("(");
    StringUtil.join(couple.getFirst(), ", ", newParameters);
    newParameters.append(")");

    StringUtil.join(couple.getSecond(), ", ", superCall);
    superCall.append(")");

    final PyElementGenerator generator = PyElementGenerator.getInstance(project);
    final LanguageLevel languageLevel = LanguageLevel.forElement(problemFunction);
    final PyStatement callSuperStatement = generator.createFromText(languageLevel, PyStatement.class, superCall.toString());
    final PyParameterList newParameterList = generator.createParameterList(languageLevel, newParameters.toString());
    problemFunction.getParameterList().replace(newParameterList);
    final PyStatementList statementList = problemFunction.getStatementList();
    PyPsiRefactoringUtil.addElementToStatementList(callSuperStatement, statementList, true);
    PyPsiUtils.removeRedundantPass(statementList);
  }


  private static @NotNull String getSelfParameterName(@NotNull ParametersInfo info) {
    final PyParameter selfParameter = info.getSelfParameter();
    if (selfParameter == null) {
      return PyNames.CANONICAL_SELF;
    }
    return StringUtil.defaultIfEmpty(selfParameter.getName(), PyNames.CANONICAL_SELF);
  }

  private static @NotNull Couple<List<String>> buildNewFunctionParamsAndSuperInitCallArgs(@NotNull ParametersInfo origInfo,
                                                                                          @NotNull ParametersInfo superInfo,
                                                                                          boolean addSelfToCall) {
    final List<String> newFunctionParams = new ArrayList<>();
    final List<String> superCallArgs = new ArrayList<>();

    final PyParameter selfParameter = origInfo.getSelfParameter();
    if (selfParameter != null && StringUtil.isNotEmpty(selfParameter.getName())) {
      newFunctionParams.add(selfParameter.getText());
    }
    else {
      newFunctionParams.add(PyNames.CANONICAL_SELF);
    }

    if (addSelfToCall) {
      superCallArgs.add(getSelfParameterName(origInfo));
    }

    // Required parameters (not-keyword)
    for (PyParameter param : origInfo.getRequiredParameters()) {
      newFunctionParams.add(param.getText());
    }
    for (PyParameter param : superInfo.getRequiredParameters()) {
      // Special case as if base class has constructor __init__((a, b), c) and
      // subclass has constructor __init__(a, (b, c))
      final PyTupleParameter tupleParam = param.getAsTuple();
      if (tupleParam != null) {
        final List<String> uniqueNames = collectParameterNames(tupleParam);
        final boolean hasDuplicates = uniqueNames.removeAll(origInfo.getAllParameterNames());
        if (hasDuplicates) {
          newFunctionParams.addAll(uniqueNames);
        }
        else {
          newFunctionParams.add(param.getText());
        }
        // Retain original structure of tuple parameter.
        // Note that tuple parameters cannot have annotations or nested default values, so it's syntactically safe
        superCallArgs.add(param.getText());
      }
      else {
        if (!origInfo.getAllParameterNames().contains(param.getName())) {
          newFunctionParams.add(param.getText());
        }
        superCallArgs.add(param.getName());
      }
    }

    // Optional parameters (not-keyword)
    for (PyParameter param : origInfo.getOptionalParameters()) {
      newFunctionParams.add(param.getText());
    }

    // Pass parameters with default values to super class constructor, only if both functions contain them
    for (PyParameter param : superInfo.getOptionalParameters()) {
      final PyTupleParameter tupleParam = param.getAsTuple();
      if (tupleParam != null) {
        if (origInfo.getAllParameterNames().containsAll(collectParameterNames(tupleParam))) {
          final String paramText = tupleParam.getText();
          final PsiElement equalSign = PyPsiUtils.getFirstChildOfType(param, PyTokenTypes.EQ);
          if (equalSign != null) {
            superCallArgs.add(paramText.substring(0, equalSign.getStartOffsetInParent()).trim());
          }
        }
      }
      else {
        if (origInfo.getAllParameterNames().contains(param.getName())) {
          superCallArgs.add(param.getName());
        }
      }
    }

    // Positional vararg
    PyParameter starredParam = null;
    if (origInfo.getPositionalContainerParameter() != null) {
      starredParam = origInfo.getPositionalContainerParameter();
    }
    else if (superInfo.getPositionalContainerParameter() != null) {
      starredParam = superInfo.getPositionalContainerParameter();
    }
    else if (origInfo.getSingleStarParameter() != null) {
      starredParam = origInfo.getSingleStarParameter();
    }
    else if (superInfo.getSingleStarParameter() != null) {
      starredParam = superInfo.getSingleStarParameter();
    }
    if (starredParam != null) {
      newFunctionParams.add(starredParam.getText());
      if (superInfo.getPositionalContainerParameter() != null) {
        superCallArgs.add("*" + starredParam.getName());
      }
    }

    // Required keyword-only parameters
    boolean newSignatureContainsKeywordParams = false;
    for (PyParameter param : origInfo.getRequiredKeywordOnlyParameters()) {
      newFunctionParams.add(param.getText());
      newSignatureContainsKeywordParams = true;
    }
    for (PyParameter param : superInfo.getRequiredKeywordOnlyParameters()) {
      if (!origInfo.getAllParameterNames().contains(param.getName())) {
        newFunctionParams.add(param.getText());
        newSignatureContainsKeywordParams = true;
      }
      superCallArgs.add(param.getName() + "=" + param.getName());
    }

    // Optional keyword-only parameters
    for (PyParameter param : origInfo.getOptionalKeywordOnlyParameters()) {
      newFunctionParams.add(param.getText());
      newSignatureContainsKeywordParams = true;
    }

    // If '*' param is followed by nothing in result signature, remove it altogether
    if (starredParam instanceof PySingleStarParameter && !newSignatureContainsKeywordParams) {
      newFunctionParams.remove(newFunctionParams.size() - 1);
    }

    for (PyParameter param : superInfo.getOptionalKeywordOnlyParameters()) {
      if (origInfo.getAllParameterNames().contains(param.getName())) {
        superCallArgs.add(param.getName() + "=" + param.getName());
      }
    }

    // Keyword vararg
    PyParameter doubleStarredParam = null;
    if (origInfo.getKeywordContainerParameter() != null) {
      doubleStarredParam = origInfo.getKeywordContainerParameter();
    }
    else if (superInfo.getKeywordContainerParameter() != null) {
      doubleStarredParam = superInfo.getKeywordContainerParameter();
    }
    if (doubleStarredParam != null) {
      newFunctionParams.add(doubleStarredParam.getText());
      if (superInfo.getKeywordContainerParameter() != null) {
        superCallArgs.add("**" + doubleStarredParam.getName());
      }
    }
    return Couple.of(newFunctionParams, superCallArgs);
  }

  private static class ParametersInfo {

    private final PyParameter mySelfParam;
    /**
     * Parameters without default value that come before first "*..." parameter.
     */
    private final List<PyParameter> myRequiredParams = new ArrayList<>();
    /**
     * Parameters with default value that come before first "*..." parameter.
     */
    private final List<PyParameter> myOptionalParams = new ArrayList<>();
    /**
     * Parameter of form "*args" (positional vararg), not the same as single "*".
     */
    private final PyParameter myPositionalContainerParam;
    /**
     * Parameter "*", that is used to delimit normal and keyword-only parameters.
     */
    private final PyParameter mySingleStarParam;
    /**
     * Parameters without default value that come after first "*..." parameter.
     */
    private final List<PyParameter> myRequiredKwOnlyParams = new ArrayList<>();
    /**
     * Parameters with default value that come after first "*..." parameter.
     */
    private final List<PyParameter> myOptionalKwOnlyParams = new ArrayList<>();
    /**
     * Parameter of form "**kwargs" (keyword vararg).
     */
    private final PyParameter myKeywordContainerParam;

    private final Set<String> myAllParameterNames = new LinkedHashSet<>();

    ParametersInfo(@NotNull PyParameterList parameterList) {
      PyParameter positionalContainer = null;
      PyParameter singleStarParam = null;
      PyParameter keywordContainer = null;
      PyParameter selfParam = null;

      for (PyParameter param : parameterList.getParameters()) {
        myAllParameterNames.addAll(collectParameterNames(param));

        if (param.isSelf()) {
          selfParam = param;
        }
        else if (param instanceof PySingleStarParameter) {
          singleStarParam = param;
        }
        else if (param.getAsNamed() != null && param.getAsNamed().isKeywordContainer()) {
          keywordContainer = param;
        }
        else if (param.getAsNamed() != null && param.getAsNamed().isPositionalContainer()) {
          positionalContainer = param;
        }
        else if (param.getAsNamed() == null || !param.getAsNamed().isKeywordOnly()) {
          if (param.hasDefaultValue()) {
            myOptionalParams.add(param);
          }
          else {
            myRequiredParams.add(param);
          }
        }
        else {
          if (param.hasDefaultValue()) {
            myOptionalKwOnlyParams.add(param);
          }
          else {
            myRequiredKwOnlyParams.add(param);
          }
        }
      }

      mySelfParam = selfParam;
      myPositionalContainerParam = positionalContainer;
      mySingleStarParam = singleStarParam;
      myKeywordContainerParam = keywordContainer;
    }

    public @Nullable PyParameter getSelfParameter() {
      return mySelfParam;
    }

    public @NotNull List<PyParameter> getRequiredParameters() {
      return Collections.unmodifiableList(myRequiredParams);
    }

    public @NotNull List<PyParameter> getOptionalParameters() {
      return Collections.unmodifiableList(myOptionalParams);
    }

    public @Nullable PyParameter getPositionalContainerParameter() {
      return myPositionalContainerParam;
    }

    public @Nullable PyParameter getSingleStarParameter() {
      return mySingleStarParam;
    }

    public @NotNull List<PyParameter> getRequiredKeywordOnlyParameters() {
      return Collections.unmodifiableList(myRequiredKwOnlyParams);
    }

    public @NotNull List<PyParameter> getOptionalKeywordOnlyParameters() {
      return Collections.unmodifiableList(myOptionalKwOnlyParams);
    }

    public @Nullable PyParameter getKeywordContainerParameter() {
      return myKeywordContainerParam;
    }

    public @NotNull Set<String> getAllParameterNames() {
      return Collections.unmodifiableSet(myAllParameterNames);
    }
  }

  private static @NotNull List<String> collectParameterNames(@NotNull PyParameter param) {
    final List<String> result = new ArrayList<>();
    collectParameterNames(param, result);
    return result;
  }


  private static void collectParameterNames(@NotNull PyParameter param, @NotNull Collection<String> acc) {
    final PyTupleParameter tupleParam = param.getAsTuple();
    if (tupleParam != null) {
      for (PyParameter subParam : tupleParam.getContents()) {
        collectParameterNames(subParam, acc);
      }
    }
    else {
      ContainerUtil.addIfNotNull(acc, param.getName());
    }
  }
}
