
package com.jetbrains.python.codeInsight.intentions;

import com.intellij.codeInsight.CodeInsightUtilCore;
import com.intellij.codeInsight.template.*;
import com.intellij.openapi.editor.Document;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.fileEditor.FileEditorManager;
import com.intellij.openapi.fileEditor.OpenFileDescriptor;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.TextRange;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.psi.PsiDocumentManager;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiReference;
import com.intellij.psi.util.PsiTreeUtil;
import com.intellij.util.IncorrectOperationException;
import com.intellij.util.ObjectUtils;
import com.jetbrains.python.PyBundle;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.PyTokenTypes;
import com.jetbrains.python.debugger.PySignature;
import com.jetbrains.python.debugger.PySignatureCacheManager;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyPsiUtils;
import org.jetbrains.annotations.NotNull;

/**
 * User: ktisha
 * <p>
 * Helps to specify type  in annotations in python3
 */
public class SpecifyTypeInPy3AnnotationsIntention extends TypeIntention {
  private String myText = PyBundle.message("INTN.specify.type.in.annotation");

  @NotNull
  public String getText() {
    return myText;
  }

  @NotNull
  public String getFamilyName() {
    return PyBundle.message("INTN.specify.type.in.annotation");
  }

  @Override
  public boolean isAvailable(@NotNull Project project, Editor editor, PsiFile file) {
    if (!LanguageLevel.forElement(file).isPy3K()) return false;
    return super.isAvailable(project, editor, file);
  }

  @Override
  public void doInvoke(@NotNull Project project, Editor editor, PsiFile file) throws IncorrectOperationException {
    final PsiElement elementAt = PyUtil.findNonWhitespaceAtOffset(file, editor.getCaretModel().getOffset());
    final PyExpression problemElement = getProblemElement(elementAt);
    final PsiReference reference = problemElement == null ? null : problemElement.getReference();

    final PsiElement resolved = reference != null ? reference.resolve() : null;
    final PyNamedParameter parameter = getParameter(problemElement, resolved);

    if (parameter != null) {
      annotateParameter(project, editor, parameter, true);
    }
    else {
      PyCallable callable = getCallable(elementAt);
      if (callable instanceof PyFunction) {
        annotateReturnType(project, (PyFunction)callable, true);
      }
    }
  }

  static PyNamedParameter annotateParameter(Project project,
                                            Editor editor,
                                            @NotNull PyNamedParameter parameter,
                                            boolean createTemplate) {
    final PyExpression defaultParamValue = parameter.getDefaultValue();

    final String paramName = StringUtil.notNullize(parameter.getName());
    final PyElementGenerator elementGenerator = PyElementGenerator.getInstance(project);

    final String defaultParamText = defaultParamValue == null ? null : defaultParamValue.getText();

    String paramType = parameterType(parameter);


    final PyNamedParameter namedParameter = elementGenerator.createParameter(paramName, defaultParamText, paramType,
                                                                             LanguageLevel.forElement(parameter));
    assert namedParameter != null;
    parameter = (PyNamedParameter)parameter.replace(namedParameter);
    parameter = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(parameter);
    editor.getCaretModel().moveToOffset(parameter.getTextOffset());
    final PyAnnotation annotation = parameter.getAnnotation();
    if (annotation != null && createTemplate) {
      final PyExpression annotationValue = annotation.getValue();

      final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(parameter);
      assert annotationValue != null : "Generated parameter must have annotation";
      final int replacementStart = annotation.getStartOffsetInParent() + annotationValue.getStartOffsetInParent();
      builder.replaceRange(TextRange.create(replacementStart,
                                            replacementStart + annotationValue.getTextLength()), paramType);
      final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
      TemplateManager.getInstance(project).startTemplate(editor, template);
    }

    return parameter;
  }

  static String parameterType(PyParameter parameter) {
    String paramType = PyNames.OBJECT;

    PyFunction function = PsiTreeUtil.getParentOfType(parameter, PyFunction.class);
    if (function != null) {
      final PySignature signature = PySignatureCacheManager.getInstance(parameter.getProject()).findSignature(
        function);
      String parameterName = parameter.getName();
      if (signature != null && parameterName != null) {
        paramType = ObjectUtils.chooseNotNull(signature.getArgTypeQualifiedName(parameterName), paramType);
      }
    }
    return paramType;
  }


  static String returnType(@NotNull PyFunction function) {
    String returnType = PyNames.OBJECT;
    final PySignature signature = PySignatureCacheManager.getInstance(function.getProject()).findSignature(function);
    if (signature != null) {
      returnType = ObjectUtils.chooseNotNull(signature.getReturnTypeQualifiedName(), returnType);
    }
    return returnType;
  }

  public static PyExpression annotateReturnType(Project project, PyFunction function, boolean createTemplate) {
    String returnType = returnType(function);

    final String annotationText = "-> " + returnType;

    final PsiDocumentManager manager = PsiDocumentManager.getInstance(project);
    Document documentWithCallable = manager.getDocument(function.getContainingFile());
    if (documentWithCallable != null) {
      try {
        manager.doPostponedOperationsAndUnblockDocument(documentWithCallable);
        final PyAnnotation oldAnnotation = function.getAnnotation();
        if (oldAnnotation != null) {
          final TextRange oldRange = oldAnnotation.getTextRange();
          documentWithCallable.replaceString(oldRange.getStartOffset(), oldRange.getEndOffset(), annotationText);
        }
        else {
          final PsiElement prevElem = PyPsiUtils.getPrevNonCommentSibling(function.getStatementList(), true);
          assert prevElem != null;
          final TextRange range = prevElem.getTextRange();
          if (prevElem.getNode().getElementType() == PyTokenTypes.COLON) {
            documentWithCallable.insertString(range.getStartOffset(), " " + annotationText);
          }
          else {
            documentWithCallable.insertString(range.getEndOffset(), " " + annotationText + ":");
          }
        }
      }
      finally {
        manager.commitDocument(documentWithCallable);
      }
    }


    function = CodeInsightUtilCore.forcePsiPostprocessAndRestoreElement(function);
    final PyAnnotation annotation = function.getAnnotation();
    assert annotation != null;
    final PyExpression annotationValue = annotation.getValue();
    assert annotationValue != null : "Generated function must have annotation";

    if (createTemplate) {
      final int offset = annotationValue.getTextOffset();

      final TemplateBuilder builder = TemplateBuilderFactory.getInstance().createTemplateBuilder(annotationValue);
      builder.replaceRange(TextRange.create(0, returnType.length()), returnType);
      final Template template = ((TemplateBuilderImpl)builder).buildInlineTemplate();
      final OpenFileDescriptor descriptor = new OpenFileDescriptor(
        project,
        function.getContainingFile().getVirtualFile(),
        offset
      );
      final Editor targetEditor = FileEditorManager.getInstance(project).openTextEditor(descriptor, true);
      if (targetEditor != null) {
        targetEditor.getCaretModel().moveToOffset(offset);
        TemplateManager.getInstance(project).startTemplate(targetEditor, template);
      }
    }
    return annotationValue;
  }

  @Override
  protected boolean isParamTypeDefined(PyParameter parameter) {
    return isDefinedInAnnotation(parameter);
  }

  private static boolean isDefinedInAnnotation(PyParameter parameter) {
    if (LanguageLevel.forElement(parameter).isOlderThan(LanguageLevel.PYTHON30)) {
      return false;
    }
    if (parameter instanceof PyNamedParameter && (((PyNamedParameter)parameter).getAnnotation() != null)) return true;
    return false;
  }

  @Override
  protected boolean isReturnTypeDefined(@NotNull PyFunction function) {
    return function.getAnnotation() != null;
  }

  @Override
  protected void updateText(boolean isReturn) {
    myText = isReturn ? PyBundle.message("INTN.specify.return.type.in.annotation") : PyBundle.message("INTN.specify.type.in.annotation");
  }
}