package com.jetbrains.python.inspections;

import com.intellij.codeInspection.CustomSuppressableInspectionTool;
import com.intellij.codeInspection.LocalInspectionTool;
import com.intellij.codeInspection.SuppressIntentionAction;
import com.intellij.codeInspection.SuppressionUtil;
import com.intellij.psi.PsiComment;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiWhiteSpace;
import com.intellij.psi.util.PsiTreeUtil;
import com.jetbrains.python.PyBundle;
import com.jetbrains.python.actions.PySuppressInspectionFix;
import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
import com.jetbrains.python.psi.*;
import org.jetbrains.annotations.Nls;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * @author yole
 */
public abstract class PyInspection extends LocalInspectionTool implements CustomSuppressableInspectionTool {
  @Nls
  @NotNull
  @Override
  public String getGroupDisplayName() {
    return PyBundle.message("INSP.GROUP.python");
  }

  @NotNull
  @Override
  public String getShortName() {
    return getClass().getSimpleName();
  }

  @Override
  public boolean isEnabledByDefault() {
    return true;
  }

  @Override
  public SuppressIntentionAction[] getSuppressActions(@Nullable PsiElement element) {
    List<SuppressIntentionAction> result = new ArrayList<SuppressIntentionAction>();
    if (element != null) {
      if (PsiTreeUtil.getParentOfType(element, PyStatementList.class, false, ScopeOwner.class) != null ||
          PsiTreeUtil.getParentOfType(element, PyFunction.class, PyClass.class) == null) {
        result.add(new PySuppressInspectionFix(getSuppressId(), "Suppress for statement", PyStatement.class));
      }
      if (PsiTreeUtil.getParentOfType(element, PyFunction.class) != null) {
        result.add(new PySuppressInspectionFix(getSuppressId(), "Suppress for function", PyFunction.class));
      }
      if (PsiTreeUtil.getParentOfType(element, PyClass.class) != null) {
        result.add(new PySuppressInspectionFix(getSuppressId(), "Suppress for class", PyClass.class));
      }
    }
    return result.toArray(new SuppressIntentionAction[result.size()]);
  }

  @Override
  public boolean isSuppressedFor(PsiElement element) {
    return isSuppressedForParent(element, PyStatement.class) ||
           isSuppressedForParent(element, PyFunction.class) ||
           isSuppressedForParent(element, PyClass.class) ||
           isSuppressForCodeFragment(element);
  }

  private boolean isSuppressForCodeFragment(PsiElement element) {
    return isSuppressForCodeFragment() && PsiTreeUtil.getParentOfType(element, PyExpressionCodeFragment.class) != null;
  }

  protected boolean isSuppressForCodeFragment() {
    return false;
  }

  private boolean isSuppressedForParent(PsiElement element, final Class<? extends PyElement> parentClass) {
    PyElement parent = PsiTreeUtil.getParentOfType(element, parentClass, false);
    if (parent == null) {
      return false;
    }
    return isSuppressedForElement(parent);
  }

  private boolean isSuppressedForElement(PyElement stmt) {
    PsiElement prevSibling = stmt.getPrevSibling();
    if (prevSibling == null) {
      final PsiElement parent = stmt.getParent();
      if (parent != null) {
        prevSibling = parent.getPrevSibling();
      }
    }
    while (prevSibling instanceof PsiComment || prevSibling instanceof PsiWhiteSpace) {
      if (prevSibling instanceof PsiComment && isSuppressedInComment(prevSibling.getText().substring(1).trim())) {
        return true;
      }
      prevSibling = prevSibling.getPrevSibling();
    }
    return false;
  }

  private static final Pattern SUPPRESS_PATTERN = Pattern.compile(SuppressionUtil.COMMON_SUPPRESS_REGEXP);

  private boolean isSuppressedInComment(String commentText) {
    Matcher m = SUPPRESS_PATTERN.matcher(commentText);
    return m.matches() && SuppressionUtil.isInspectionToolIdMentioned(m.group(1), getSuppressId());
  }

  private String getSuppressId() {
    return getShortName().replace("Inspection", "");
  }
}

@@ -1372,16 +1372,22 @@
 ullable 
+final 
 PsiEleme
@@ -1493,31 +1493,191 @@
     
-if (element != null
+result.add(new PySuppressInspectionFix(getSuppressId(), %22Suppress for statement%22, PyStatement.class) %7B%0A      @Override%0A      protected PsiElement getContainer(PsiElement context
 ) %7B%0A
+  
     
@@ -1706,30 +1706,30 @@
 arentOfType(
-elemen
+contex
 t, PyStateme
@@ -1788,16 +1788,18 @@
         
+  
 PsiTreeU
@@ -1810,38 +1810,38 @@
 getParentOfType(
-elemen
+contex
 t, PyFunction.cl
@@ -1883,108 +1883,43 @@
     
-result.add(new PySuppressInspectionFix(getSuppressId(), %22Suppress for statement%22, PyStatement.class)
+  return super.getContainer(context
 );%0A 
@@ -1919,32 +1919,34 @@
 ext);%0A      
+ 
+ 
 %7D%0A      
 if (PsiTreeU
@@ -1937,82 +1937,39 @@
     
-if (PsiTreeUtil.getParentOfType(element, PyFunction.class) != null) %7B
+  return null;%0A      %7D
 %0A    
+%7D);%0A
     
@@ -2073,93 +2073,8 @@
 ));%0A
-      %7D%0A      if (PsiTreeUtil.getParentOfType(element, PyClass.class) != null) %7B%0A    
     
@@ -2168,30 +2168,16 @@
 lass));%0A
-      %7D%0A    %7D%0A
     retu

allow suppressing PyCharm inspections from toolwindow (PY-2103)