package mr.intellij.plugin.autofactory.actions;

import com.google.auto.factory.AutoFactory;
import com.intellij.openapi.editor.Editor;
import com.intellij.openapi.project.Project;
import com.intellij.psi.PsiClass;
import com.intellij.psi.PsiFile;
import org.jetbrains.annotations.NotNull;

import java.util.Optional;

import static com.intellij.codeInsight.generation.OverrideImplementUtil.getContextClass;
import static mr.intellij.plugin.autofactory.utils.PsiUtils.createFormattedAnnotation;
import static mr.intellij.plugin.autofactory.utils.PsiUtils.hasAutoFactory;

/**
 * Provides a way to add {@code @AutoFactory} annotation to a class.
 */
public class AutoFactoryOnClassAction extends BaseAction {

    public AutoFactoryOnClassAction() {
        super(new BaseCodeInsightActionHandler() {

            @Override
            protected boolean doInvoke(@NotNull Project project, @NotNull Editor editor, @NotNull PsiFile file) {
                return Optional.ofNullable(getContextClass(project, editor, file, false))
                        .map(this::addAnnotation)
                        .isPresent();
            }

            @SuppressWarnings("ConstantConditions")
            private boolean addAnnotation(@NotNull PsiClass psiClass) {
                if (hasAutoFactory(psiClass)) {
                    return false;
                }

                psiClass.getModifierList().addAfter(createFormattedAnnotation(psiClass, AutoFactory.class), null);
                return true;
            }
        });
    }
}

@@ -288,36 +288,8 @@
 l;%0A%0A
-import java.util.Optional;%0A%0A
 impo
@@ -437,109 +437,9 @@
 ils.
-createFormattedAnnotation;%0Aimport static mr.intellij.plugin.autofactory.utils.PsiUtils.hasAutoFactory
+*
 ;%0A%0A/
@@ -829,27 +829,21 @@
 urn 
-Optional.ofNullable
+addAnnotation
 (get
@@ -889,95 +889,8 @@
 se))
-%0A                        .map(this::addAnnotation)%0A                        .isPresent()
 ;%0A  
@@ -996,25 +996,16 @@
 otation(
-@NotNull 
 PsiClass
@@ -1037,16 +1037,36 @@
     if (
+psiClass == null %7C%7C 
 hasAutoF

Refactor AutoFactoryOnClassAction