From 7f894aecafad429558e60aa2d23a8877f09a06d4 Mon Sep 17 00:00:00 2001
From: Filip Hrisafov <filip.hrisafov@gmail.com>
Date: Fri, 15 Dec 2023 22:53:03 +0100
Subject: [PATCH] More fixes

---
 .../WrongUsageOfMappersFactoryInspection.java | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/mapstruct/intellij/inspection/WrongUsageOfMappersFactoryInspection.java b/src/main/java/org/mapstruct/intellij/inspection/WrongUsageOfMappersFactoryInspection.java
index ec058f0..2746646 100644
--- a/src/main/java/org/mapstruct/intellij/inspection/WrongUsageOfMappersFactoryInspection.java
+++ b/src/main/java/org/mapstruct/intellij/inspection/WrongUsageOfMappersFactoryInspection.java
@@ -7,6 +7,8 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Collection;
 
 import com.intellij.codeInsight.AnnotationUtil;
 import com.intellij.codeInsight.daemon.impl.quickfix.RemoveUnusedVariableFix;
@@ -100,15 +102,22 @@ public void visitMethodCallExpression(PsiMethodCallExpression expression) {
                 PsiClass mapperClass = (PsiClass) mapperElement;
                 PsiAnnotation mapperAnnotation = mapperClass.getAnnotation( MapstructUtil.MAPPER_ANNOTATION_FQN );
                 if ( mapperAnnotation == null ) {
-                    problemsHolder.registerProblem(
-                        expression,
-                        MapStructBundle.message( "inspection.wrong.usage.mappers.factory.non.mapstruct" ),
+                    Collection<LocalQuickFix> fixes = new ArrayList<>( 2 );
+                    fixes.add(
                         new AddAnnotationPsiFix(
                             MapstructUtil.MAPPER_ANNOTATION_FQN,
                             mapperClass,
                             PsiNameValuePair.EMPTY_ARRAY
-                        ),
-                        createRemoveMappersFix( expression )
+                        )
+                    );
+                    LocalQuickFix removeMappersFix = createRemoveMappersFix( expression );
+                    if ( removeMappersFix != null ) {
+                        fixes.add( removeMappersFix );
+                    }
+                    problemsHolder.registerProblem(
+                        expression,
+                        MapStructBundle.message( "inspection.wrong.usage.mappers.factory.non.mapstruct" ),
+                        fixes.toArray( LocalQuickFix[]::new )
                     );
                 }
                 else {