refactor

file_combine/src/test/java/click/kamil/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/Main.java

@@ -72,7 +72,7 @@ public class Main {
     static void processEntryPoint(TypeDeclaration td, List<StateMachineExporter> outputs, Path outputDir, CodebaseContext context) throws IOException {
         String className = context.getFqn(td);
         System.out.println("Processing state machine config: " + className);
-        
+
         StateMachineAggregator aggregator = new StateMachineAggregator(context);
         List<Transition> transitions = aggregator.aggregateTransitions(td);
 
@@ -96,9 +96,9 @@ public class Main {
         String uniqueName = parentFqn + "#" + methodName;
 
         System.out.println("Processing state machine bean: " + uniqueName);
-        
+
-        List<Transition> transitions = AstTransitionParser.parseTransitions(m);
+        List<Transition> transitions = AstTransitionParser.parseTransitions(m, context);
-        
+
         Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
         Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
 

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/AstTransitionParser.java

@@ -2,11 +2,27 @@ package click.kamil.springstatemachineexporter.ast.app;
 
 import click.kamil.springstatemachineexporter.ast.app.domain.Action;
 import click.kamil.springstatemachineexporter.ast.app.domain.Guard;
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
-import org.eclipse.jdt.core.dom.*;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.eclipse.jdt.core.dom.ASTVisitor;
+import org.eclipse.jdt.core.dom.ClassInstanceCreation;
+import org.eclipse.jdt.core.dom.CompilationUnit;
+import org.eclipse.jdt.core.dom.Expression;
+import org.eclipse.jdt.core.dom.LambdaExpression;
+import org.eclipse.jdt.core.dom.MethodDeclaration;
+import org.eclipse.jdt.core.dom.MethodInvocation;
+import org.eclipse.jdt.core.dom.SimpleName;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 public class AstTransitionParser {
@@ -15,15 +31,16 @@ public class AstTransitionParser {
             .map(TransitionType::getMethodName)
             .collect(Collectors.toSet());
 
-    public static List<Transition> parseTransitions(MethodDeclaration method) {
+    public static List<Transition> parseTransitions(MethodDeclaration method, CodebaseContext context) {
-        if (method == null) return Collections.emptyList();
+        if (method == null)
+            return Collections.emptyList();
         List<Transition> transitions = new ArrayList<>();
 
         method.accept(new ASTVisitor() {
             @Override
             public boolean visit(MethodInvocation node) {
                 if (isEntryPoint(node)) {
-                    transitions.addAll(parseTransitionsFromExpression(node, new HashMap<>()));
+                    transitions.addAll(parseTransitionsFromExpression(node, new HashMap<>(), false, context));
                 }
                 return super.visit(node);
             }
@@ -59,12 +76,13 @@ public class AstTransitionParser {
         return chain;
     }
 
-    public static List<Transition> parseTransitionsFromExpression(MethodInvocation rootCall, Map<String, Expression> argsMap) {
+    public static List<Transition> parseTransitionsFromExpression(MethodInvocation rootCall, Map<String, Expression> argsMap, CodebaseContext context) {
-        return parseTransitionsFromExpression(rootCall, argsMap, false);
+        return parseTransitionsFromExpression(rootCall, argsMap, false, context);
     }
 
-    public static List<Transition> parseTransitionsFromExpression(MethodInvocation rootCall, Map<String, Expression> argsMap, boolean isContinuation) {
+    public static List<Transition> parseTransitionsFromExpression(MethodInvocation rootCall, Map<String, Expression> argsMap, boolean isContinuation, CodebaseContext context) {
         List<Transition> transitions = new ArrayList<>();
+        CompilationUnit cu = (CompilationUnit) rootCall.getRoot();
 
         // Step 1: Unroll method chain
         List<MethodInvocation> calls = new ArrayList<>();
@@ -87,12 +105,13 @@ public class AstTransitionParser {
                 currentSegment.add(call);
             }
         }
-        if (!currentSegment.isEmpty()) segments.add(currentSegment);
+        if (!currentSegment.isEmpty())
+            segments.add(currentSegment);
 
         // Step 3: Parse each segment
         for (int i = 0; i < segments.size(); i++) {
             List<MethodInvocation> segment = segments.get(i);
-            
+
             Optional<TransitionType> segmentType = segment.stream()
                     .map(mi -> TransitionType.fromMethodName(mi.getName().getIdentifier()))
                     .flatMap(Optional::stream)
@@ -102,16 +121,16 @@ public class AstTransitionParser {
             boolean isFirstSegmentOfContinuation = isContinuation && i == 0;
 
             if (segmentType.isPresent() && (segmentType.get() == TransitionType.CHOICE || segmentType.get() == TransitionType.JUNCTION)) {
-                List<Transition> choiceTransitions = parseChoiceOrJunctionTransitions(segment, segmentType.get(), argsMap);
+                List<Transition> choiceTransitions = parseChoiceOrJunctionTransitions(segment, segmentType.get(), argsMap, context, cu);
                 transitions.addAll(choiceTransitions);
             } else {
-                Transition t = parseStandardTransition(segment, argsMap);
+                Transition t = parseStandardTransition(segment, argsMap, context, cu);
                 // If it's a continuation's first segment and it has no source, it's just refining the receiver
                 if (isFirstSegmentOfContinuation && t.getSourceStates().isEmpty()) {
-                     // Don't add it as a new transition if it doesn't have source/target
+                    // Don't add it as a new transition if it doesn't have source/target
-                     if (!t.getTargetStates().isEmpty()) {
+                    if (!t.getTargetStates().isEmpty()) {
-                         transitions.add(t);
+                        transitions.add(t);
-                     }
+                    }
                 } else if (!t.getSourceStates().isEmpty() || !t.getTargetStates().isEmpty() || t.getType() != null) {
                     transitions.add(t);
                 }
@@ -121,16 +140,16 @@ public class AstTransitionParser {
         return transitions;
     }
 
-    private static List<Transition> parseChoiceOrJunctionTransitions(List<MethodInvocation> segment, TransitionType type, Map<String, Expression> argsMap) {
+    private static List<Transition> parseChoiceOrJunctionTransitions(List<MethodInvocation> segment, TransitionType type, Map<String, Expression> argsMap, CodebaseContext context, CompilationUnit cu) {
         List<Transition> transitions = new ArrayList<>();
-        String sourceState = null;
+        State sourceState = null;
         int orderCounter = 0;
 
         // Extract source state once
         for (MethodInvocation call : segment) {
             if ("source".equals(call.getName().getIdentifier()) && !call.arguments().isEmpty()) {
                 Expression arg = resolveArg((Expression) call.arguments().get(0), argsMap);
-                sourceState = QuotedExpression.of(arg).toStringWithoutQuotes();
+                sourceState = context.resolveState(arg, cu);
                 break; // only take the first occurrence
             }
         }
@@ -146,10 +165,10 @@ public class AstTransitionParser {
             if (sourceState != null) {
                 t.getSourceStates().add(sourceState);
             }
-            
+
             Expression targetArg = resolveArg((Expression) args.get(0), argsMap);
-            t.getTargetStates().add(QuotedExpression.of(targetArg).toStringWithoutQuotes());
+            t.getTargetStates().add(context.resolveState(targetArg, cu));
-            
+
             if (args.size() > 1) {
                 Expression secondArg = resolveArg((Expression) args.get(1), argsMap);
                 if ("last".equals(methodName)) {
@@ -170,8 +189,7 @@ public class AstTransitionParser {
         return transitions;
     }
 
-
+    private static Transition parseStandardTransition(List<MethodInvocation> segment, Map<String, Expression> argsMap, CodebaseContext context, CompilationUnit cu) {
-    private static Transition parseStandardTransition(List<MethodInvocation> segment, Map<String, Expression> argsMap) {
         Transition t = new Transition();
         for (MethodInvocation call : segment) {
             String methodName = call.getName().getIdentifier();
@@ -189,11 +207,11 @@ public class AstTransitionParser {
             switch (methodName) {
                 case "source" -> args.forEach(arg -> {
                     Expression resolved = resolveArg((Expression) arg, argsMap);
-                    t.getSourceStates().add(QuotedExpression.of(resolved).toStringWithoutQuotes());
+                    t.getSourceStates().add(context.resolveState(resolved, cu));
                 });
                 case "target" -> args.forEach(arg -> {
                     Expression resolved = resolveArg((Expression) arg, argsMap);
-                    t.getTargetStates().add(QuotedExpression.of(resolved).toStringWithoutQuotes());
+                    t.getTargetStates().add(context.resolveState(resolved, cu));
                 });
                 case "event" -> {
                     Expression resolved = resolveArg((Expression) args.get(0), argsMap);
@@ -229,7 +247,8 @@ public class AstTransitionParser {
     private static Expression resolveArg(Expression expr, Map<String, Expression> argsMap) {
         if (expr instanceof SimpleName sn) {
             Expression resolved = argsMap.get(sn.getIdentifier());
-            if (resolved != null) return resolved;
+            if (resolved != null)
+                return resolved;
         }
         return expr;
     }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/QuotedExpression.java

@@ -7,13 +7,14 @@ import org.eclipse.jdt.core.dom.StringLiteral;
 @Getter
 public class QuotedExpression {
     private final Expression expression;
-    
+
     private QuotedExpression(Expression expression) {
         this.expression = expression;
     }
 
     public String toStringWithoutQuotes() {
-        if (expression == null) return null;
+        if (expression == null)
+            return null;
         if (expression instanceof StringLiteral sl) {
             return sl.getLiteralValue();
         }
@@ -21,18 +22,19 @@ public class QuotedExpression {
     }
 
     public static String stripQuotes(String s) {
-        if (s == null) return null;
+        if (s == null)
+            return null;
         return s.replaceAll("^\"|\"$", "");
     }
-    
+
     public static QuotedExpression of(Expression expression) {
         return new QuotedExpression(expression);
     }
-    
+
     public static QuotedExpression of(Object obj) {
         if (obj instanceof Expression expr) {
             return new QuotedExpression(expr);
         }
         return null;
     }
-} 
+} 

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/StateMachineAggregator.java

@@ -4,15 +4,39 @@ import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
 import lombok.Getter;
-import org.eclipse.jdt.core.dom.*;
+import org.eclipse.jdt.core.dom.ArrayInitializer;
+import org.eclipse.jdt.core.dom.Block;
+import org.eclipse.jdt.core.dom.CompilationUnit;
+import org.eclipse.jdt.core.dom.EnhancedForStatement;
+import org.eclipse.jdt.core.dom.Expression;
+import org.eclipse.jdt.core.dom.ExpressionStatement;
+import org.eclipse.jdt.core.dom.LambdaExpression;
+import org.eclipse.jdt.core.dom.MethodDeclaration;
+import org.eclipse.jdt.core.dom.MethodInvocation;
+import org.eclipse.jdt.core.dom.NullLiteral;
+import org.eclipse.jdt.core.dom.ParameterizedType;
+import org.eclipse.jdt.core.dom.SimpleName;
+import org.eclipse.jdt.core.dom.SimpleType;
+import org.eclipse.jdt.core.dom.SingleVariableDeclaration;
+import org.eclipse.jdt.core.dom.TryStatement;
+import org.eclipse.jdt.core.dom.Type;
+import org.eclipse.jdt.core.dom.TypeDeclaration;
+import org.eclipse.jdt.core.dom.VariableDeclaration;
+import org.eclipse.jdt.core.dom.VariableDeclarationFragment;
+import org.eclipse.jdt.core.dom.VariableDeclarationStatement;
 
-import java.util.*;
+import java.util.ArrayList;
-import java.util.stream.Collectors;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 public class StateMachineAggregator {
     private final CodebaseContext context;
 
-    private enum AnalysisGoal { TRANSITIONS, STATES }
+    private enum AnalysisGoal {TRANSITIONS, STATES}
 
     public StateMachineAggregator(CodebaseContext context) {
         this.context = context;
@@ -20,17 +44,20 @@ public class StateMachineAggregator {
 
     public List<Transition> aggregateTransitions(TypeDeclaration startClass) {
         List<Transition> allTransitions = new ArrayList<>();
-        aggregateTransitionsRecursive(startClass, startClass, allTransitions, new HashSet<>(), new HashSet<>(), new HashMap<>());
+        Set<MethodVisit> visitedMethods = new HashSet<>();
+        aggregateTransitionsRecursive(startClass, startClass, allTransitions, new HashSet<>(), visitedMethods, new HashMap<>());
         return allTransitions;
     }
 
     private void aggregateTransitionsRecursive(TypeDeclaration currentTd, TypeDeclaration searchStartClass, List<Transition> allTransitions, Set<String> visitedClasses, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
-        if (currentTd == null || visitedClasses.contains(currentTd.getName().getFullyQualifiedName())) return;
+        if (currentTd == null || visitedClasses.contains(currentTd.getName().getFullyQualifiedName()))
+            return;
         visitedClasses.add(currentTd.getName().getFullyQualifiedName());
 
-        MethodDeclaration configureMethod = findConfigureTransitionsMethod(currentTd);
+        for (MethodDeclaration method : currentTd.getMethods()) {
-        if (configureMethod != null) {
+            if (isConfigureTransitionsMethod(method)) {
-            allTransitions.addAll(parseTransitionsWithMethodCalls(configureMethod, searchStartClass, visitedMethods, argsMap));
+                allTransitions.addAll(parseTransitionsWithMethodCalls(method, searchStartClass, visitedMethods, argsMap));
+            }
         }
 
         CompilationUnit cu = (CompilationUnit) currentTd.getRoot();
@@ -55,7 +82,8 @@ public class StateMachineAggregator {
         }
     }
 
-    private record MethodVisit(MethodDeclaration method, Map<String, Object> args) {}
+    private record MethodVisit(MethodDeclaration method, Map<String, Object> args) {
+    }
 
     private List<Transition> parseTransitionsWithMethodCalls(MethodDeclaration method, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
         MethodVisit visit = new MethodVisit(method, new HashMap<>(argsMap));
@@ -76,9 +104,9 @@ public class StateMachineAggregator {
         List<Transition> transitions = new ArrayList<>();
         if (expr instanceof MethodInvocation mi) {
             if (isTransitionChainEntry(mi)) {
-                transitions.addAll(AstTransitionParser.parseTransitionsFromExpression(mi, convertToExpressionMap(argsMap)));
+                transitions.addAll(AstTransitionParser.parseTransitionsFromExpression(mi, convertToExpressionMap(argsMap), context));
             } else if (isTransitionChainContinuation(mi, argsMap)) {
-                transitions.addAll(AstTransitionParser.parseTransitionsFromExpression(mi, convertToExpressionMap(argsMap), true));
+                transitions.addAll(AstTransitionParser.parseTransitionsFromExpression(mi, convertToExpressionMap(argsMap), true, context));
             } else if (isForEach(mi)) {
                 transitions.addAll(parseForEach(mi, searchStartClass, visitedMethods, argsMap, AnalysisGoal.TRANSITIONS, null, null));
             } else {
@@ -128,18 +156,23 @@ public class StateMachineAggregator {
         for (int i = 0; i < params.size() && i < arguments.size(); i++) {
             Object p = params.get(i);
             String paramName = "";
-            if (p instanceof VariableDeclaration vp) paramName = vp.getName().getIdentifier();
+            if (p instanceof VariableDeclaration vp)
+                paramName = vp.getName().getIdentifier();
             if (!paramName.isEmpty()) {
                 lambdaArgsMap.put(paramName, resolveExpression((Expression) arguments.get(i), argsMap));
             }
         }
 
         if (goal == AnalysisGoal.TRANSITIONS) {
-            if (lambda.getBody() instanceof Block block) return parseTransitionsInBlock(block, searchStartClass, visitedMethods, lambdaArgsMap);
+            if (lambda.getBody() instanceof Block block)
-            else if (lambda.getBody() instanceof Expression lambdaExpr) return parseTransitionsInExpression(lambdaExpr, searchStartClass, visitedMethods, lambdaArgsMap);
+                return parseTransitionsInBlock(block, searchStartClass, visitedMethods, lambdaArgsMap);
+            else if (lambda.getBody() instanceof Expression lambdaExpr)
+                return parseTransitionsInExpression(lambdaExpr, searchStartClass, visitedMethods, lambdaArgsMap);
         } else {
-            if (lambda.getBody() instanceof Block block) parseStatesInBlock(block, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
+            if (lambda.getBody() instanceof Block block)
-            else if (lambda.getBody() instanceof Expression lambdaExpr) parseStatesInExpression(lambdaExpr, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
+                parseStatesInBlock(block, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
+            else if (lambda.getBody() instanceof Expression lambdaExpr)
+                parseStatesInExpression(lambdaExpr, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
         }
         return Collections.emptyList();
     }
@@ -148,21 +181,27 @@ public class StateMachineAggregator {
         List<Transition> transitions = new ArrayList<>();
         Expression receiver = forEachMi.getExpression();
         List<Expression> elements = unrollCollection(receiver, argsMap);
-        if (elements.isEmpty() || forEachMi.arguments().isEmpty()) return transitions;
+        if (elements.isEmpty() || forEachMi.arguments().isEmpty())
+            return transitions;
 
         Object lambdaObj = forEachMi.arguments().get(0);
         if (lambdaObj instanceof LambdaExpression lambda) {
             String paramName = getLambdaParamName(lambda);
             for (Expression element : elements) {
                 Map<String, Object> lambdaArgsMap = new HashMap<>(argsMap);
-                if (!paramName.isEmpty()) lambdaArgsMap.put(paramName, element);
+                if (!paramName.isEmpty())
-                
+                    lambdaArgsMap.put(paramName, element);
+
                 if (goal == AnalysisGoal.TRANSITIONS) {
-                    if (lambda.getBody() instanceof Block block) transitions.addAll(parseTransitionsInBlock(block, searchStartClass, visitedMethods, lambdaArgsMap));
+                    if (lambda.getBody() instanceof Block block)
-                    else if (lambda.getBody() instanceof Expression lambdaExpr) transitions.addAll(parseTransitionsInExpression(lambdaExpr, searchStartClass, visitedMethods, lambdaArgsMap));
+                        transitions.addAll(parseTransitionsInBlock(block, searchStartClass, visitedMethods, lambdaArgsMap));
+                    else if (lambda.getBody() instanceof Expression lambdaExpr)
+                        transitions.addAll(parseTransitionsInExpression(lambdaExpr, searchStartClass, visitedMethods, lambdaArgsMap));
                 } else {
-                    if (lambda.getBody() instanceof Block block) parseStatesInBlock(block, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
+                    if (lambda.getBody() instanceof Block block)
-                    else if (lambda.getBody() instanceof Expression lambdaExpr) parseStatesInExpression(lambdaExpr, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
+                        parseStatesInBlock(block, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
+                    else if (lambda.getBody() instanceof Expression lambdaExpr)
+                        parseStatesInExpression(lambdaExpr, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
                 }
             }
         }
@@ -179,11 +218,15 @@ public class StateMachineAggregator {
             loopArgsMap.put(paramName, element);
 
             if (goal == AnalysisGoal.TRANSITIONS) {
-                if (efs.getBody() instanceof Block block) transitions.addAll(parseTransitionsInBlock(block, searchStartClass, visitedMethods, loopArgsMap));
+                if (efs.getBody() instanceof Block block)
-                else if (efs.getBody() instanceof ExpressionStatement es) transitions.addAll(parseTransitionsInExpression(es.getExpression(), searchStartClass, visitedMethods, loopArgsMap));
+                    transitions.addAll(parseTransitionsInBlock(block, searchStartClass, visitedMethods, loopArgsMap));
+                else if (efs.getBody() instanceof ExpressionStatement es)
+                    transitions.addAll(parseTransitionsInExpression(es.getExpression(), searchStartClass, visitedMethods, loopArgsMap));
             } else {
-                if (efs.getBody() instanceof Block block) parseStatesInBlock(block, searchStartClass, initialStates, endStates, visitedMethods, loopArgsMap);
+                if (efs.getBody() instanceof Block block)
-                else if (efs.getBody() instanceof ExpressionStatement es) parseStatesInExpression(es.getExpression(), searchStartClass, initialStates, endStates, visitedMethods, loopArgsMap);
+                    parseStatesInBlock(block, searchStartClass, initialStates, endStates, visitedMethods, loopArgsMap);
+                else if (efs.getBody() instanceof ExpressionStatement es)
+                    parseStatesInExpression(es.getExpression(), searchStartClass, initialStates, endStates, visitedMethods, loopArgsMap);
             }
         }
         return transitions;
@@ -194,12 +237,13 @@ public class StateMachineAggregator {
             String name = mi.getName().getIdentifier();
             Expression receiver = mi.getExpression();
             if ((name.equals("of") && receiver != null && (receiver.toString().equals("Stream") || receiver.toString().equals("List")))
-                || (name.equals("asList") && receiver != null && receiver.toString().equals("Arrays"))) {
+                    || (name.equals("asList") && receiver != null && receiver.toString().equals("Arrays"))) {
                 if (mi.arguments().size() == 1) {
                     Expression arg = (Expression) mi.arguments().get(0);
                     if (arg instanceof SimpleName sn) {
                         Object resolved = argsMap.get(sn.getIdentifier());
-                        if (resolved instanceof List list) return (List<Expression>) list;
+                        if (resolved instanceof List list)
+                            return (List<Expression>) list;
                     }
                 }
                 return (List<Expression>) mi.arguments();
@@ -207,17 +251,21 @@ public class StateMachineAggregator {
         }
         if (expr instanceof SimpleName sn) {
             Object resolved = argsMap.get(sn.getIdentifier());
-            if (resolved instanceof List list) return (List<Expression>) list;
+            if (resolved instanceof List list)
-            if (resolved instanceof Expression e) return unrollCollection(e, argsMap);
+                return (List<Expression>) list;
+            if (resolved instanceof Expression e)
+                return unrollCollection(e, argsMap);
         }
-        if (expr instanceof ArrayInitializer ai) return (List<Expression>) ai.expressions();
+        if (expr instanceof ArrayInitializer ai)
+            return (List<Expression>) ai.expressions();
         return Collections.emptyList();
     }
 
     private String getLambdaParamName(LambdaExpression lambda) {
         if (lambda.parameters().size() == 1) {
             Object p = lambda.parameters().get(0);
-            if (p instanceof VariableDeclaration vp) return vp.getName().getIdentifier();
+            if (p instanceof VariableDeclaration vp)
+                return vp.getName().getIdentifier();
         }
         return "";
     }
@@ -225,49 +273,57 @@ public class StateMachineAggregator {
     private List<Transition> parseTransitionsInBlock(Block block, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
         List<Transition> transitions = new ArrayList<>();
         for (Object stmt : block.statements()) {
-             if (stmt instanceof ExpressionStatement es) transitions.addAll(parseTransitionsInExpression(es.getExpression(), searchStartClass, visitedMethods, argsMap));
+            if (stmt instanceof ExpressionStatement es)
-             else if (stmt instanceof VariableDeclarationStatement vds) {
+                transitions.addAll(parseTransitionsInExpression(es.getExpression(), searchStartClass, visitedMethods, argsMap));
-                 for (Object fragmentObj : vds.fragments()) {
+            else if (stmt instanceof VariableDeclarationStatement vds) {
-                     if (fragmentObj instanceof VariableDeclarationFragment fragment) {
+                for (Object fragmentObj : vds.fragments()) {
-                         Expression initializer = fragment.getInitializer();
+                    if (fragmentObj instanceof VariableDeclarationFragment fragment) {
-                         if (initializer != null) {
+                        Expression initializer = fragment.getInitializer();
-                             Object resolved = resolve(initializer, argsMap);
+                        if (initializer != null) {
-                             if (resolved != null) argsMap.put(fragment.getName().getIdentifier(), resolved);
+                            Object resolved = resolve(initializer, argsMap);
-                             transitions.addAll(parseTransitionsInExpression(initializer, searchStartClass, visitedMethods, argsMap));
+                            if (resolved != null)
-                         }
+                                argsMap.put(fragment.getName().getIdentifier(), resolved);
-                     }
+                            transitions.addAll(parseTransitionsInExpression(initializer, searchStartClass, visitedMethods, argsMap));
-                 }
+                        }
-             }
+                    }
-             else if (stmt instanceof EnhancedForStatement efs) transitions.addAll(parseEnhancedForLoop(efs, searchStartClass, visitedMethods, argsMap, AnalysisGoal.TRANSITIONS, null, null));
+                }
-             else if (stmt instanceof TryStatement ts) transitions.addAll(parseTransitionsInBlock(ts.getBody(), searchStartClass, visitedMethods, argsMap));
+            } else if (stmt instanceof EnhancedForStatement efs)
+                transitions.addAll(parseEnhancedForLoop(efs, searchStartClass, visitedMethods, argsMap, AnalysisGoal.TRANSITIONS, null, null));
+            else if (stmt instanceof TryStatement ts)
+                transitions.addAll(parseTransitionsInBlock(ts.getBody(), searchStartClass, visitedMethods, argsMap));
         }
         return transitions;
     }
 
     private void parseStatesInBlock(Block block, TypeDeclaration searchStartClass, Set<String> initialStates, Set<String> endStates, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
         for (Object stmt : block.statements()) {
-             if (stmt instanceof ExpressionStatement es) parseStatesInExpression(es.getExpression(), searchStartClass, initialStates, endStates, visitedMethods, argsMap);
+            if (stmt instanceof ExpressionStatement es)
-             else if (stmt instanceof VariableDeclarationStatement vds) {
+                parseStatesInExpression(es.getExpression(), searchStartClass, initialStates, endStates, visitedMethods, argsMap);
-                 for (Object fragmentObj : vds.fragments()) {
+            else if (stmt instanceof VariableDeclarationStatement vds) {
-                     if (fragmentObj instanceof VariableDeclarationFragment fragment) {
+                for (Object fragmentObj : vds.fragments()) {
-                         Expression initializer = fragment.getInitializer();
+                    if (fragmentObj instanceof VariableDeclarationFragment fragment) {
-                         if (initializer != null) {
+                        Expression initializer = fragment.getInitializer();
-                             Object resolved = resolve(initializer, argsMap);
+                        if (initializer != null) {
-                             if (resolved != null) argsMap.put(fragment.getName().getIdentifier(), resolved);
+                            Object resolved = resolve(initializer, argsMap);
-                             parseStatesInExpression(initializer, searchStartClass, initialStates, endStates, visitedMethods, argsMap);
+                            if (resolved != null)
-                         }
+                                argsMap.put(fragment.getName().getIdentifier(), resolved);
-                     }
+                            parseStatesInExpression(initializer, searchStartClass, initialStates, endStates, visitedMethods, argsMap);
-                 }
+                        }
-             }
+                    }
-             else if (stmt instanceof EnhancedForStatement efs) parseEnhancedForLoop(efs, searchStartClass, visitedMethods, argsMap, AnalysisGoal.STATES, initialStates, endStates);
+                }
-             else if (stmt instanceof TryStatement ts) parseStatesInBlock(ts.getBody(), searchStartClass, initialStates, endStates, visitedMethods, argsMap);
+            } else if (stmt instanceof EnhancedForStatement efs)
+                parseEnhancedForLoop(efs, searchStartClass, visitedMethods, argsMap, AnalysisGoal.STATES, initialStates, endStates);
+            else if (stmt instanceof TryStatement ts)
+                parseStatesInBlock(ts.getBody(), searchStartClass, initialStates, endStates, visitedMethods, argsMap);
         }
     }
 
     private void parseStatesInExpression(Expression expr, TypeDeclaration searchStartClass, Set<String> initialStates, Set<String> endStates, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
         if (expr instanceof MethodInvocation mi) {
-            if (isWithStatesChain(mi, argsMap)) parseStateChain(mi, initialStates, endStates, argsMap);
+            if (isWithStatesChain(mi, argsMap))
-            else if (isForEach(mi)) parseForEach(mi, searchStartClass, visitedMethods, argsMap, AnalysisGoal.STATES, initialStates, endStates);
+                parseStateChain(mi, initialStates, endStates, argsMap);
+            else if (isForEach(mi))
+                parseForEach(mi, searchStartClass, visitedMethods, argsMap, AnalysisGoal.STATES, initialStates, endStates);
             else {
                 // Check if it's a call on a lambda parameter or a variable that resolved to a lambda
                 Expression receiver = mi.getExpression();
@@ -291,17 +347,18 @@ public class StateMachineAggregator {
         Map<String, Object> map = new HashMap<>(currentArgsMap);
         List<?> params = method.parameters();
         List<?> args = call.arguments();
-        
+
         for (int i = 0; i < params.size(); i++) {
             SingleVariableDeclaration param = (SingleVariableDeclaration) params.get(i);
             String paramName = param.getName().getIdentifier();
             if (param.isVarargs() && i < params.size()) {
                 List<Expression> varargList = new ArrayList<>();
                 for (int j = i; j < args.size(); j++) {
-                    Object resolved = resolve( (Expression) args.get(j), currentArgsMap);
+                    Object resolved = resolve((Expression) args.get(j), currentArgsMap);
                     if (resolved instanceof List<?> list) {
                         for (Object item : list) {
-                            if (item instanceof Expression e) varargList.add(e);
+                            if (item instanceof Expression e)
+                                varargList.add(e);
                         }
                     } else if (resolved instanceof Expression e) {
                         varargList.add(e);
@@ -316,36 +373,42 @@ public class StateMachineAggregator {
     }
 
     private Object resolve(Expression expr, Map<String, Object> argsMap) {
-        if (expr instanceof NullLiteral) return null;
+        if (expr instanceof NullLiteral)
+            return null;
         if (expr instanceof SimpleName sn) {
             Object resolved = argsMap.get(sn.getIdentifier());
             if (resolved instanceof Expression nextExpr && nextExpr != expr) {
                 return resolve(nextExpr, argsMap);
             }
-            if (resolved != null) return resolved;
+            if (resolved != null)
+                return resolved;
         }
         return expr;
     }
 
     private Expression resolveExpression(Expression expr, Map<String, Object> argsMap) {
         Object resolved = resolve(expr, argsMap);
-        if (resolved instanceof Expression e) return e;
+        if (resolved instanceof Expression e)
+            return e;
         return expr;
     }
 
     private Map<String, Expression> convertToExpressionMap(Map<String, Object> argsMap) {
         Map<String, Expression> map = new HashMap<>();
         for (Map.Entry<String, Object> entry : argsMap.entrySet()) {
-            if (entry.getValue() instanceof Expression e) map.put(entry.getKey(), e);
+            if (entry.getValue() instanceof Expression e)
+                map.put(entry.getKey(), e);
         }
         return map;
     }
 
     private boolean isTransitionChainEntry(MethodInvocation mi) {
         String name = mi.getName().getIdentifier();
-        if (TransitionType.fromMethodName(name).isPresent()) return true;
+        if (TransitionType.fromMethodName(name).isPresent())
+            return true;
         Expression receiver = mi.getExpression();
-        if (receiver instanceof MethodInvocation next) return isTransitionChainEntry(next);
+        if (receiver instanceof MethodInvocation next)
+            return isTransitionChainEntry(next);
         return false;
     }
 
@@ -353,7 +416,8 @@ public class StateMachineAggregator {
         String name = mi.getName().getIdentifier();
         if (Set.of("and", "source", "target", "event", "first", "then", "last").contains(name)) {
             Expression receiver = mi.getExpression();
-            if (receiver instanceof MethodInvocation next) return isTransitionChainContinuation(next, argsMap) || isTransitionChainEntry(next);
+            if (receiver instanceof MethodInvocation next)
+                return isTransitionChainContinuation(next, argsMap) || isTransitionChainEntry(next);
             Object resolved = resolve(receiver, argsMap);
             return resolved instanceof MethodInvocation || resolved == null; // null means it's a parameter we track
         }
@@ -361,8 +425,11 @@ public class StateMachineAggregator {
     }
 
     private MethodDeclaration findMethodInHierarchy(String methodName, TypeDeclaration td) {
-        if (td == null) return null;
+        if (td == null)
-        for (MethodDeclaration method : td.getMethods()) if (method.getName().getIdentifier().equals(methodName)) return method;
+            return null;
+        for (MethodDeclaration method : td.getMethods())
+            if (method.getName().getIdentifier().equals(methodName))
+                return method;
         Type superclassType = td.getSuperclassType();
         if (superclassType != null) {
             TypeDeclaration superclassTd = context.getTypeDeclaration(getSimpleName(superclassType));
@@ -371,16 +438,21 @@ public class StateMachineAggregator {
         return null;
     }
 
-    private boolean isRelevantMethod(MethodDeclaration method) { return true; }
+    private boolean isRelevantMethod(MethodDeclaration method) {
+        return true;
+    }
 
     private MethodDeclaration findConfigureTransitionsMethod(TypeDeclaration td) {
-        for (MethodDeclaration method : td.getMethods()) if (isConfigureTransitionsMethod(method)) return method;
+        for (MethodDeclaration method : td.getMethods())
+            if (isConfigureTransitionsMethod(method))
+                return method;
         return null;
     }
 
     private boolean isConfigureTransitionsMethod(MethodDeclaration method) {
-        if (!method.getName().getIdentifier().equals("configure") || method.parameters().size() != 1) return false;
+        if (!method.getName().getIdentifier().equals("configure") || method.parameters().isEmpty())
-        return ((SingleVariableDeclaration) method.parameters().get(0)).getType().toString().startsWith("StateMachineTransitionConfigurer");
+            return false;
+        return ((SingleVariableDeclaration) method.parameters().get(0)).getType().toString().contains("StateMachineTransitionConfigurer");
     }
 
     @Getter
@@ -389,15 +461,20 @@ public class StateMachineAggregator {
     private final Set<String> endStates = new HashSet<>();
 
     public void aggregateStates(TypeDeclaration startClass) {
-        aggregateStatesRecursive(startClass, startClass, initialStates, endStates, new HashSet<>(), new HashSet<>(), new HashMap<>());
+        Set<MethodVisit> visitedMethods = new HashSet<>();
+        aggregateStatesRecursive(startClass, startClass, initialStates, endStates, new HashSet<>(), visitedMethods, new HashMap<>());
     }
 
     private void aggregateStatesRecursive(TypeDeclaration currentTd, TypeDeclaration searchStartClass, Set<String> initialStates, Set<String> endStates, Set<String> visitedClasses, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
-        if (currentTd == null || visitedClasses.contains(currentTd.getName().getFullyQualifiedName())) return;
+        if (currentTd == null || visitedClasses.contains(currentTd.getName().getFullyQualifiedName()))
+            return;
         visitedClasses.add(currentTd.getName().getFullyQualifiedName());
 
-        MethodDeclaration configureMethod = findConfigureStatesMethod(currentTd);
+        for (MethodDeclaration method : currentTd.getMethods()) {
-        if (configureMethod != null) parseStatesWithMethodCalls(configureMethod, searchStartClass, initialStates, endStates, visitedMethods, argsMap);
+            if (isConfigureStatesMethod(method)) {
+                parseStatesWithMethodCalls(method, searchStartClass, initialStates, endStates, visitedMethods, argsMap);
+            }
+        }
 
         CompilationUnit cu = (CompilationUnit) currentTd.getRoot();
 
@@ -405,34 +482,41 @@ public class StateMachineAggregator {
         Type superclassType = currentTd.getSuperclassType();
         if (superclassType != null) {
             TypeDeclaration superclassTd = context.getTypeDeclaration(getSimpleName(superclassType), cu);
-            if (superclassTd != null) aggregateStatesRecursive(superclassTd, searchStartClass, initialStates, endStates, visitedClasses, visitedMethods, argsMap);
+            if (superclassTd != null)
+                aggregateStatesRecursive(superclassTd, searchStartClass, initialStates, endStates, visitedClasses, visitedMethods, argsMap);
         }
 
         // Hierarchy - Interfaces
         for (Object interfaceTypeObj : currentTd.superInterfaceTypes()) {
             if (interfaceTypeObj instanceof Type interfaceType) {
                 TypeDeclaration interfaceTd = context.getTypeDeclaration(getSimpleName(interfaceType), cu);
-                if (interfaceTd != null) aggregateStatesRecursive(interfaceTd, searchStartClass, initialStates, endStates, visitedClasses, visitedMethods, argsMap);
+                if (interfaceTd != null)
+                    aggregateStatesRecursive(interfaceTd, searchStartClass, initialStates, endStates, visitedClasses, visitedMethods, argsMap);
             }
         }
     }
 
     private void parseStatesWithMethodCalls(MethodDeclaration method, TypeDeclaration searchStartClass, Set<String> initialStates, Set<String> endStates, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
         MethodVisit visit = new MethodVisit(method, new HashMap<>(argsMap));
-        if (visitedMethods.contains(visit)) return;
+        if (visitedMethods.contains(visit))
+            return;
         visitedMethods.add(visit);
         Block body = method.getBody();
-        if (body == null) return;
+        if (body == null)
+            return;
 
         parseStatesInBlock(body, searchStartClass, initialStates, endStates, visitedMethods, argsMap);
     }
 
     private boolean isWithStatesChain(MethodInvocation mi, Map<String, Object> argsMap) {
-        if (mi.getName().getIdentifier().equals("withStates")) return true;
+        if (mi.getName().getIdentifier().equals("withStates"))
+            return true;
         Expression receiver = mi.getExpression();
-        if (receiver instanceof MethodInvocation next) return isWithStatesChain(next, argsMap);
+        if (receiver instanceof MethodInvocation next)
+            return isWithStatesChain(next, argsMap);
         Object resolved = resolve(receiver, argsMap);
-        if (resolved instanceof MethodInvocation next) return isWithStatesChain(next, argsMap);
+        if (resolved instanceof MethodInvocation next)
+            return isWithStatesChain(next, argsMap);
         return false;
     }
 
@@ -443,37 +527,50 @@ public class StateMachineAggregator {
             chain.addFirst(m);
             current = m.getExpression();
         }
+        CompilationUnit cu = (CompilationUnit) mi.getRoot();
         for (MethodInvocation call : chain) {
             String methodName = call.getName().getIdentifier();
-            if (call.arguments().isEmpty()) continue;
+            if (call.arguments().isEmpty())
+                continue;
             Expression arg = resolveExpression((Expression) call.arguments().get(0), argsMap);
-            if ("initial".equals(methodName)) { if (!isInsideRegionOrFork(call)) initialStates.add(QuotedExpression.of(arg).toStringWithoutQuotes()); }
+            if ("initial".equals(methodName)) {
-            else if ("end".equals(methodName)) endStates.add(QuotedExpression.of(arg).toStringWithoutQuotes());
+                if (!isInsideRegionOrFork(call)) {
+                    initialStates.add(context.resolveState(arg, cu).toString());
+                }
+            } else if ("end".equals(methodName)) {
+                endStates.add(context.resolveState(arg, cu).toString());
+            }
         }
     }
 
     private boolean isInsideRegionOrFork(MethodInvocation mi) {
         Expression expr = mi.getExpression();
         while (expr instanceof MethodInvocation parent) {
-            if (List.of("fork", "region").contains(parent.getName().getIdentifier())) return true;
+            if (List.of("fork", "region").contains(parent.getName().getIdentifier()))
+                return true;
             expr = parent.getExpression();
         }
         return false;
     }
 
     private MethodDeclaration findConfigureStatesMethod(TypeDeclaration td) {
-        for (MethodDeclaration method : td.getMethods()) if (isConfigureStatesMethod(method)) return method;
+        for (MethodDeclaration method : td.getMethods())
+            if (isConfigureStatesMethod(method))
+                return method;
         return null;
     }
 
     private boolean isConfigureStatesMethod(MethodDeclaration method) {
-        if (!method.getName().getIdentifier().equals("configure") || method.parameters().size() != 1) return false;
+        if (!method.getName().getIdentifier().equals("configure") || method.parameters().isEmpty())
-        return ((SingleVariableDeclaration) method.parameters().get(0)).getType().toString().startsWith("StateMachineStateConfigurer");
+            return false;
+        return ((SingleVariableDeclaration) method.parameters().get(0)).getType().toString().contains("StateMachineStateConfigurer");
     }
 
     private String getSimpleName(Type type) {
-        if (type.isSimpleType()) return ((SimpleType) type).getName().getFullyQualifiedName();
+        if (type.isSimpleType())
-        if (type.isParameterizedType()) return getSimpleName(((ParameterizedType) type).getType());
+            return ((SimpleType) type).getName().getFullyQualifiedName();
+        if (type.isParameterizedType())
+            return getSimpleName(((ParameterizedType) type).getType());
         return type.toString();
     }
 }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/TransitionStateUtils.java

@@ -1,5 +1,6 @@
 package click.kamil.springstatemachineexporter.ast.app;
 
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 
 import java.util.Collection;
@@ -43,8 +44,9 @@ public class TransitionStateUtils {
                 .collect(Collectors.toSet());
     }
 
-    private static Stream<String> normalizeStates(Collection<String> states) {
+    private static Stream<String> normalizeStates(Collection<State> states) {
         return states.stream()
+                .map(State::toString)
                 .map(QuotedExpression::stripQuotes)
                 .filter(Objects::nonNull);
     }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/domain/State.java

@@ -0,0 +1,16 @@
+package click.kamil.springstatemachineexporter.ast.app.domain;
+
+public record State(String rawName, String fullIdentifier) {
+    public static State of(String name) {
+        return new State(name, name);
+    }
+
+    public static State of(String raw, String full) {
+        return new State(raw, full);
+    }
+
+    @Override
+    public String toString() {
+        return fullIdentifier != null ? fullIdentifier : rawName;
+    }
+}

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/app/domain/Transition.java

@@ -12,8 +12,8 @@ import java.util.List;
 @Setter
 public class Transition {
     private TransitionType type;
-    private List<String> sourceStates = new ArrayList<>();
+    private List<State> sourceStates = new ArrayList<>();
-    private List<String> targetStates = new ArrayList<>();
+    private List<State> targetStates = new ArrayList<>();
     private String event;
 
     private Guard guard;

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/common/AstUtils.java

@@ -1,10 +1,19 @@
 package click.kamil.springstatemachineexporter.ast.common;
 
-import org.eclipse.jdt.core.dom.*;
+import org.eclipse.jdt.core.dom.ArrayType;
+import org.eclipse.jdt.core.dom.Name;
+import org.eclipse.jdt.core.dom.NameQualifiedType;
+import org.eclipse.jdt.core.dom.ParameterizedType;
+import org.eclipse.jdt.core.dom.PrimitiveType;
+import org.eclipse.jdt.core.dom.QualifiedName;
+import org.eclipse.jdt.core.dom.QualifiedType;
+import org.eclipse.jdt.core.dom.SimpleType;
+import org.eclipse.jdt.core.dom.Type;
 
 public final class AstUtils {
     private AstUtils() {
     }
+
     public static String extractSimpleTypeName(Type type) {
         if (type.isSimpleType()) {
             Name name = ((SimpleType) type).getName();

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/common/CodebaseContext.java

@@ -1,11 +1,28 @@
 package click.kamil.springstatemachineexporter.ast.common;
 
-import org.eclipse.jdt.core.dom.*;
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
+import org.eclipse.jdt.core.dom.AST;
+import org.eclipse.jdt.core.dom.ASTParser;
+import org.eclipse.jdt.core.dom.Annotation;
+import org.eclipse.jdt.core.dom.CompilationUnit;
+import org.eclipse.jdt.core.dom.Expression;
+import org.eclipse.jdt.core.dom.ImportDeclaration;
+import org.eclipse.jdt.core.dom.MethodDeclaration;
+import org.eclipse.jdt.core.dom.Modifier;
+import org.eclipse.jdt.core.dom.QualifiedName;
+import org.eclipse.jdt.core.dom.SimpleName;
+import org.eclipse.jdt.core.dom.StringLiteral;
+import org.eclipse.jdt.core.dom.Type;
+import org.eclipse.jdt.core.dom.TypeDeclaration;
 
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 public class CodebaseContext {
     private final Map<String, CompilationUnit> classes = new HashMap<>();
@@ -30,6 +47,46 @@ public class CodebaseContext {
         }
     }
 
+    public State resolveState(Expression expr, CompilationUnit cu) {
+        if (expr == null)
+            return null;
+        if (expr instanceof StringLiteral sl)
+            return State.of(sl.getLiteralValue());
+
+        String raw = expr.toString();
+        if (expr instanceof QualifiedName qn) {
+            String full = resolveQualifiedName(qn, cu);
+            return State.of(raw, full);
+        }
+        if (expr instanceof SimpleName sn) {
+            String full = resolveSimpleName(sn, cu);
+            return State.of(raw, full);
+        }
+        return State.of(raw);
+    }
+
+    private String resolveQualifiedName(QualifiedName qn, CompilationUnit cu) {
+        String qualifier = qn.getQualifier().toString();
+        String name = qn.getName().getIdentifier();
+        TypeDeclaration td = getTypeDeclaration(qualifier, cu);
+        if (td != null)
+            return getFqn(td) + "." + name;
+        return qn.getFullyQualifiedName();
+    }
+
+    private String resolveSimpleName(SimpleName sn, CompilationUnit cu) {
+        String name = sn.getIdentifier();
+        if (cu == null)
+            return name;
+        for (Object impObj : cu.imports()) {
+            ImportDeclaration imp = (ImportDeclaration) impObj;
+            String impName = imp.getName().getFullyQualifiedName();
+            if (impName.endsWith("." + name))
+                return impName;
+        }
+        return name;
+    }
+
     private CompilationUnit parse(String source) {
         ASTParser parser = ASTParser.newParser(AST.getJLSLatest());
         parser.setSource(source.toCharArray());
@@ -41,23 +98,27 @@ public class CodebaseContext {
     public TypeDeclaration getTypeDeclaration(String name) {
         // Try exact FQN match first
         CompilationUnit cu = classes.get(name);
-        if (cu != null) return findTypeInCu(cu, name);
+        if (cu != null)
-        
+            return findTypeInCu(cu, name);
+
         // If not found, it might be a simple name
         String fqn = simpleNameToFqn.get(name);
         if (fqn != null) {
             cu = classes.get(fqn);
-            if (cu != null) return findTypeInCu(cu, fqn);
+            if (cu != null)
+                return findTypeInCu(cu, fqn);
         }
-        
+
         return null;
     }
 
     public TypeDeclaration getTypeDeclaration(String name, CompilationUnit contextCu) {
-        if (name == null || name.isEmpty()) return null;
+        if (name == null || name.isEmpty())
-        
+            return null;
+
         // 1. Check if it's already an FQN
-        if (classes.containsKey(name)) return getTypeDeclaration(name);
+        if (classes.containsKey(name))
+            return getTypeDeclaration(name);
 
         // 2. Check imports in contextCu
         if (contextCu != null) {
@@ -65,27 +126,31 @@ public class CodebaseContext {
                 ImportDeclaration imp = (ImportDeclaration) impObj;
                 String impName = imp.getName().getFullyQualifiedName();
                 if (!imp.isStatic() && !imp.isOnDemand()) {
-                    if (impName.endsWith("." + name)) return getTypeDeclaration(impName);
+                    if (impName.endsWith("." + name))
+                        return getTypeDeclaration(impName);
                 }
             }
-            
+
             // 3. Check same package
             String packageName = contextCu.getPackage() != null ? contextCu.getPackage().getName().getFullyQualifiedName() : "";
             String localFqn = packageName.isEmpty() ? name : packageName + "." + name;
-            if (classes.containsKey(localFqn)) return getTypeDeclaration(localFqn);
+            if (classes.containsKey(localFqn))
-            
+                return getTypeDeclaration(localFqn);
+
             // 4. Check on-demand imports (star imports)
             for (Object impObj : contextCu.imports()) {
                 ImportDeclaration imp = (ImportDeclaration) impObj;
                 if (imp.isOnDemand()) {
                     String starFqn = imp.getName().getFullyQualifiedName() + "." + name;
-                    if (classes.containsKey(starFqn)) return getTypeDeclaration(starFqn);
+                    if (classes.containsKey(starFqn))
+                        return getTypeDeclaration(starFqn);
                 }
             }
-            
+
             // 5. Fallback to java.lang (common)
             String langFqn = "java.lang." + name;
-            if (classes.containsKey(langFqn)) return getTypeDeclaration(langFqn);
+            if (classes.containsKey(langFqn))
+                return getTypeDeclaration(langFqn);
         }
 
         // 6. Last resort: global simple name match (might be ambiguous but better than nothing)
@@ -105,7 +170,8 @@ public class CodebaseContext {
             if (type instanceof TypeDeclaration td) {
                 String simpleName = td.getName().getIdentifier();
                 String typeFqn = packageName.isEmpty() ? simpleName : packageName + "." + simpleName;
-                if (typeFqn.equals(fqn)) return td;
+                if (typeFqn.equals(fqn))
+                    return td;
             }
         }
         return null;
@@ -130,12 +196,21 @@ public class CodebaseContext {
     }
 
     private boolean isEntryPoint(TypeDeclaration td, List<String> targetAnnotations) {
+        if (td.isInterface())
+            return false;
+
         boolean isAbstract = false;
         // Has annotation
         for (Object modifier : td.modifiers()) {
-            if (modifier instanceof Annotation annotation &&
+            if (modifier instanceof Annotation annotation) {
-                    targetAnnotations.contains(annotation.getTypeName().getFullyQualifiedName())) {
+                String annotationName = annotation.getTypeName().getFullyQualifiedName();
-                return true;
+                String simpleAnnotationName = annotationName.contains(".")
+                        ? annotationName.substring(annotationName.lastIndexOf('.') + 1)
+                        : annotationName;
+
+                if (targetAnnotations.contains(annotationName) || targetAnnotations.contains(simpleAnnotationName)) {
+                    return true;
+                }
             }
             if (modifier instanceof Modifier m && m.isAbstract()) {
                 isAbstract = true;
@@ -158,21 +233,25 @@ public class CodebaseContext {
         Type superclass = td.getSuperclassType();
         if (superclass != null) {
             String superclassName = getSimpleName(superclass);
-            if (knownAdapters.contains(superclassName)) return true;
+            if (knownAdapters.contains(superclassName))
+                return true;
             TypeDeclaration superTd = getTypeDeclaration(superclassName, cu);
-            if (superTd != null && extendsStateMachineConfigurerAdapter(superTd)) return true;
+            if (superTd != null && extendsStateMachineConfigurerAdapter(superTd))
+                return true;
         }
 
         // Check interfaces
         for (Object interfaceTypeObj : td.superInterfaceTypes()) {
             if (interfaceTypeObj instanceof Type interfaceType) {
                 String interfaceName = getSimpleName(interfaceType);
-                if (knownAdapters.contains(interfaceName)) return true;
+                if (knownAdapters.contains(interfaceName))
+                    return true;
                 TypeDeclaration interfaceTd = getTypeDeclaration(interfaceName, cu);
-                if (interfaceTd != null && extendsStateMachineConfigurerAdapter(interfaceTd)) return true;
+                if (interfaceTd != null && extendsStateMachineConfigurerAdapter(interfaceTd))
+                    return true;
             }
         }
-        
+
         return false;
     }
 
@@ -195,9 +274,11 @@ public class CodebaseContext {
     private boolean isBeanMethodReturning(MethodDeclaration method, Set<String> targetReturnTypes) {
         // Check return type
         Type returnType = method.getReturnType2();
-        if (returnType == null) return false;
+        if (returnType == null)
+            return false;
         String typeName = AstUtils.extractSimpleTypeName(returnType);
-        if (!targetReturnTypes.contains(typeName)) return false;
+        if (!targetReturnTypes.contains(typeName))
+            return false;
 
         // Check for @Bean annotation
         for (Object modifier : method.modifiers()) {

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/common/StringUtils.java

@@ -47,5 +47,4 @@ public final class StringUtils {
 
     private StringUtils() {
     }
-
 }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/out/Dot.java

@@ -1,14 +1,19 @@
 package click.kamil.springstatemachineexporter.ast.out;
 
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 
-import java.util.*;
+import java.util.HashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 public class Dot implements StateMachineExporter {
 
     private static final String[] CHOICE_COLORS = {
-            "blue", "darkgreen", "darkred", "darkorange", "darkviolet", "teal", "brown", "crimson"
+            "blue", "green", "red", "purple", "orange", "brown", "darkgreen", "darkblue"
     };
 
     @Override
@@ -16,23 +21,21 @@ public class Dot implements StateMachineExporter {
                          Set<String> startStates,
                          Set<String> endStates,
                          boolean useLambdaGuards) {
-        boolean includeChoiceStates = true;
         StringBuilder sb = new StringBuilder();
-        sb.append("digraph stateMachine {\n");
+        sb.append("digraph statemachine {\n");
         sb.append("  rankdir=LR;\n");
-        sb.append("  node [shape=circle, style=filled, fillcolor=lightgray];\n\n");
+        sb.append("  node [shape=rounded, style=filled, fillcolor=white, fontname=\"Arial\"];\n");
+        sb.append("  edge [fontname=\"Arial\", fontsize=10];\n\n");
 
-        sb.append("  start [shape=point, label=\"\"];\n");
+        // Choices
-
+        Set<String> statesWithChoice = new LinkedHashSet<>();
-        // Identify choice states and assign colors
-        Set<String> statesWithChoice = new HashSet<>();
         Map<String, String> choiceColorMap = new HashMap<>();
         int colorIndex = 0;
 
         for (Transition t : transitions) {
             if (t.getType() == TransitionType.CHOICE && t.getSourceStates() != null) {
-                for (String source : t.getSourceStates()) {
+                for (State source : t.getSourceStates()) {
-                    String simplified = simplify(source);
+                    String simplified = simplify(source.toString());
                     statesWithChoice.add(simplified);
                     if (!choiceColorMap.containsKey(simplified)) {
                         choiceColorMap.put(simplified, CHOICE_COLORS[colorIndex % CHOICE_COLORS.length]);
@@ -42,57 +45,31 @@ public class Dot implements StateMachineExporter {
             }
         }
 
-        // Collect all states
+        boolean includeChoiceStates = !statesWithChoice.isEmpty();
-        Set<String> allStates = new HashSet<>();
-        for (Transition t : transitions) {
-            if (t.getSourceStates() != null) {
-                for (String s : t.getSourceStates()) allStates.add(simplify(s));
-            }
-            if (t.getTargetStates() != null) {
-                for (String tgt : t.getTargetStates()) allStates.add(simplify(tgt));
-            }
-        }
-
-        // Declare normal states
-        for (String state : allStates) {
-            if (!statesWithChoice.contains(includeChoiceStates ? state + "_choice" : state)) {
-                String shape = endStates.contains(state) ? "doublecircle" : "circle";
-                sb.append("  ").append(state)
-                        .append(" [shape=").append(shape).append(", fillcolor=lightgray];\n");
-            }
-        }
-
-        // Declare choice pseudostates
-        for (String state : statesWithChoice) {
-            String choiceState = includeChoiceStates ? state + "_choice" : state;
-            sb.append("  ").append(choiceState)
-                    .append(" [shape=diamond, label=\"\", fillcolor=lightyellow];\n");
-        }
-
-        sb.append("\n");
-
-        // Start state arrows
-        for (String start : startStates) {
-            sb.append("  start -> ").append(simplify(start)).append(";\n");
-        }
-
-        sb.append("\n");
-
         if (includeChoiceStates) {
-            // Connect normal states to their choice pseudostates
+            for (String choice : statesWithChoice) {
-            for (String state : statesWithChoice) {
+                String color = choiceColorMap.get(choice);
-                sb.append("  ").append(state).append(" -> ").append(state).append("_choice;\n");
+                sb.append("  ").append(choice).append("_choice [shape=diamond, label=\"\", fillcolor=\"").append(color).append("\", width=0.2, height=0.2];\n");
+                sb.append("  ").append(choice).append(" -> ").append(choice).append("_choice [style=dotted, color=\"").append(color).append("\"];\n");
             }
-            sb.append("\n");
         }
 
+        // Start/End nodes
+        sb.append("  _start [shape=circle, label=\"\", fillcolor=black, width=0.1];\n");
+        for (String start : startStates) {
+            sb.append("  _start -> ").append(simplify(start)).append(";\n");
+        }
+        for (String end : endStates) {
+            sb.append("  ").append(simplify(end)).append(" [fillcolor=lightgray];\n");
+        }
 
         // Transitions
         for (Transition t : transitions) {
-            if (t.getSourceStates() == null || t.getSourceStates().isEmpty()) continue;
+            if (t.getSourceStates() == null || t.getSourceStates().isEmpty())
+                continue;
 
             boolean allowNoTarget = t.getType() == TransitionType.JOIN || t.getType() == TransitionType.FORK;
-            List<String> targets = t.getTargetStates();
+            List<State> targets = t.getTargetStates();
 
             if (targets == null || targets.isEmpty()) {
                 if (!allowNoTarget) {
@@ -102,10 +79,12 @@ public class Dot implements StateMachineExporter {
                 }
             }
 
-            for (String rawSource : t.getSourceStates()) {
+            for (State rawSourceState : t.getSourceStates()) {
+                String rawSource = rawSourceState.toString();
                 String source = simplify(rawSource);
 
-                for (String rawTarget : targets) {
+                for (State rawTargetState : targets) {
+                    String rawTarget = rawTargetState.toString();
                     String target = simplify(rawTarget);
                     String edgeSource = source;
 
@@ -115,14 +94,14 @@ public class Dot implements StateMachineExporter {
 
                     // Label: event [guard] / actions
                     StringBuilder label = new StringBuilder();
-
                     if (t.getEvent() != null && !t.getEvent().isBlank()) {
-                        label.append(escapeDotString(simplify(t.getEvent())));
+                        label.append(escapeDotString(t.getEvent()));
                     }
 
                     // Guard
                     if (t.getGuard() != null && !t.getGuard().expression().isBlank()) {
-                        if (!label.isEmpty()) label.append(" ");
+                        if (!label.isEmpty())
+                            label.append(" ");
                         String guardText = useLambdaGuards && t.getGuard().isLambda() ? "λ"
                                 : escapeDotString(t.getGuard().expression().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim());
                         label.append("[").append(guardText).append("]");
@@ -132,17 +111,20 @@ public class Dot implements StateMachineExporter {
                     if (t.getActions() != null && !t.getActions().isEmpty()) {
                         StringBuilder actionBuilder = new StringBuilder();
                         for (int i = 0; i < t.getActions().size(); i++) {
-                            if (i > 0) actionBuilder.append(", ");
+                            if (i > 0)
+                                actionBuilder.append(", ");
                             var action = t.getActions().get(i);
                             actionBuilder.append(useLambdaGuards && action.isLambda() ? "λ" : escapeDotString(action.expression()));
                         }
-                        if (!label.isEmpty()) label.append(" / ");
+                        if (!label.isEmpty())
+                            label.append(" / ");
                         label.append(actionBuilder);
                     }
 
                     // Order
                     if ((t.getType() == TransitionType.CHOICE || t.getType() == TransitionType.JUNCTION) && t.getOrder() != null) {
-                        if (!label.isEmpty()) label.append(" ");
+                        if (!label.isEmpty())
+                            label.append(" ");
                         label.append("(order=").append(t.getOrder()).append(")");
                     }
 
@@ -208,13 +190,15 @@ public class Dot implements StateMachineExporter {
         return ".dot";
     }
 
-    private String escapeDotString(String input) {
+    private String escapeDotString(String s) {
-        if (input == null) return "";
+        if (s == null)
-        return input.replace("\\", "\\\\").replace("\"", "\\\"");
+            return "";
+        return s.replace("\"", "\\\"");
     }
 
     private void appendEdgeAttr(StringBuilder sb, String key, String value) {
-        if (!sb.isEmpty()) sb.append(", ");
+        if (!sb.isEmpty())
+            sb.append(", ");
         sb.append(key).append("=\"").append(value).append("\"");
     }
 }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/out/PlantUml.java

@@ -1,18 +1,25 @@
 package click.kamil.springstatemachineexporter.ast.out;
 
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 import click.kamil.springstatemachineexporter.ast.common.StringUtils;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
 import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 public class PlantUml implements StateMachineExporter {
-    // Color palette for choices
     private final List<String> choiceColors = List.of("FF6347", "4682B4", "32CD32", "FFD700", "6A5ACD", "FF69B4");
+
     private String getColor(Transition t, String source, Map<String, String> choiceStateColorMap) {
-        if (t.getType() == null) return "000000";
+        if (t.getType() == null)
+            return "000000";
         return switch (t.getType()) {
             case CHOICE -> choiceStateColorMap.getOrDefault(source, "000000");
             case EXTERNAL -> "1E90FF";
@@ -23,6 +30,7 @@ public class PlantUml implements StateMachineExporter {
             case FORK -> "20B2AA";
         };
     }
+
     private final String LAMBDA = "λ";
 
     @Override
@@ -30,97 +38,83 @@ public class PlantUml implements StateMachineExporter {
                          Set<String> startStates,
                          Set<String> endStates,
                          boolean useLambdaGuards) {
-        boolean includeChoiceStates = false;
         StringBuilder sb = new StringBuilder();
         sb.append("@startuml\n");
         sb.append("skinparam state {\n");
         sb.append("  BackgroundColor<<choice>> LightYellow\n");
         sb.append("}\n\n");
 
-        // 1. Print start states
         for (String start : startStates) {
             sb.append("[*] --> ").append(simplify(start)).append("\n");
         }
         sb.append("\n");
 
-        // 2. Find all states that have choice transitions originating from them
         Set<String> statesWithChoice = new HashSet<>();
         for (Transition t : transitions) {
             if (t.getType() == TransitionType.CHOICE) {
-                for (String source : t.getSourceStates()) {
+                for (State source : t.getSourceStates()) {
-                    statesWithChoice.add(simplify(source));
+                    statesWithChoice.add(simplify(source.toString()));
                 }
             }
         }
 
-        // 3. Declare choice pseudostates & connect from original state
-        for (String state : statesWithChoice) {
-            String choiceState = state;
-            if (includeChoiceStates) {
-                choiceState = state + "_choice";
-                sb.append("state ").append(choiceState).append(" <<choice>>\n");
-                sb.append(state).append(" --> ").append(choiceState).append("\n");
-            } else {
-                sb.append("state ").append(choiceState).append(" <<choice>>\n");
-            }
-        }
-        sb.append("\n");
-
-        // Assign color to each choice state
         Map<String, String> choiceStateColorMap = new HashMap<>();
         int colorIndex = 0;
+
         for (String state : statesWithChoice) {
+            sb.append("state ").append(state).append(" <<choice>>\n");
             choiceStateColorMap.put(state, choiceColors.get(colorIndex % choiceColors.size()));
             colorIndex++;
         }
 
-        Set<String> transitionLines = new HashSet<>();
+        Set<String> junctionStates = transitions.stream()
+                .filter(t -> t.getType() == TransitionType.JUNCTION)
+                .flatMap(t -> t.getSourceStates().stream())
+                .map(s -> simplify(s.toString()))
+                .collect(Collectors.toSet());
+        for (String junction : junctionStates) {
+            sb.append("state ").append(junction).append(" <<choice>>\n");
+        }
+
+        sb.append("\n");
 
-        // 4. Output transitions
         for (Transition t : transitions) {
-            if (t.getSourceStates() == null || t.getSourceStates().isEmpty()) continue;
+            if (t.getSourceStates() == null || t.getSourceStates().isEmpty())
+                continue;
 
             boolean allowNoTarget = t.getType() == TransitionType.JOIN || t.getType() == TransitionType.FORK;
 
             List<String> targets;
             if (t.getTargetStates() == null || t.getTargetStates().isEmpty()) {
                 if (!allowNoTarget) {
-                    targets = t.getSourceStates(); // Self-loop
+                    targets = t.getSourceStates().stream().map(s -> s.toString()).toList();
                 } else {
                     continue;
                 }
             } else {
-                targets = t.getTargetStates();
+                targets = t.getTargetStates().stream().map(s -> s.toString()).toList();
             }
 
-            for (String rawSource : t.getSourceStates()) {
+            for (State rawSourceState : t.getSourceStates()) {
-                String source = simplify(rawSource);
+                String source = simplify(rawSourceState.toString());
 
                 for (String rawTarget : targets) {
                     String target = simplify(rawTarget);
-                    String transitionSource;
+                    String transitionSource = source;
-                    if (includeChoiceStates) {
+
-                        transitionSource = t.getType() == TransitionType.CHOICE && statesWithChoice.contains(source)
+                    String color = getColor(t, source, choiceStateColorMap);
-                                ? source + "_choice"
+                    sb.append(transitionSource).append(" -[#").append(color).append(",bold]-> ").append(target);
-                                : source;
-                    } else  {
-                        transitionSource = source;
-                    }
 
                     String label = buildLabel(t, useLambdaGuards);
-                    String color = getColor(t, source, choiceStateColorMap);
+                    if (!label.isEmpty()) {
-                    String line = transitionSource + " -[#"+color+",bold]-> " + target + (label.isEmpty() ? "" : (" : " + label));
+                        sb.append(" : ").append(label);
-
-                    if (transitionLines.add(line)) {
-                        sb.append(line).append("\n");
                     }
+                    sb.append("\n");
                 }
             }
         }
 
         sb.append("\n");
-
-        // 5. Mark end states
         for (String end : endStates) {
             sb.append(simplify(end)).append(" --> [*]\n");
         }
@@ -135,23 +129,16 @@ public class PlantUml implements StateMachineExporter {
     }
 
     private String getGuardText(Transition t, boolean useLambdaGuards) {
-        if (t.getGuard() == null || StringUtils.isBlank(t.getGuard().expression())) {
+        if (t.getGuard() == null || StringUtils.isBlank(t.getGuard().expression()))
             return "";
-        }
+        if (useLambdaGuards && t.getGuard().isLambda())
-        if (useLambdaGuards && t.getGuard().isLambda()) {
             return LAMBDA;
-        }
+        return t.getGuard().expression().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim();
-        return t.getGuard().expression()
-                .replaceAll("[\\n\\r]+", " ")
-                .replaceAll("\\s{2,}", " ")
-                .trim();
     }
 
     private String getActionsText(Transition t, boolean useLambdaGuards) {
-        if (t.getActions() == null || t.getActions().isEmpty()) {
+        if (t.getActions() == null || t.getActions().isEmpty())
             return "";
-        }
-
         return t.getActions().stream()
                 .map(action -> (useLambdaGuards && action.isLambda()) ? LAMBDA : action.expression())
                 .collect(Collectors.joining(", "));
@@ -159,30 +146,22 @@ public class PlantUml implements StateMachineExporter {
 
     private String buildLabel(Transition t, boolean useLambdaGuards) {
         List<String> parts = new ArrayList<>();
-        Optional.ofNullable(t.getEvent())
+        if (t.getEvent() != null && !t.getEvent().isBlank())
-                .filter(e -> !e.isBlank())
+            parts.add(simplify(t.getEvent()));
-                .map(this::simplify)
+        String guard = getGuardText(t, useLambdaGuards);
-                .ifPresent(parts::add);
+        if (!guard.isEmpty())
-        Optional.of(getGuardText(t, useLambdaGuards))
+            parts.add("[" + guard + "]");
-                .filter(g -> !g.isBlank())
+        String actions = getActionsText(t, useLambdaGuards);
-                .map(g -> "[" + g + "]")
+        if (!actions.isEmpty()) {
-                .ifPresent(parts::add);
+            if (!parts.isEmpty())
-        Optional.of(getActionsText(t, useLambdaGuards))
+                parts.add("/ " + actions);
-                .filter(a -> !a.isBlank())
+            else
-                .ifPresent(actions -> {
+                parts.add(actions);
-                    if (!parts.isEmpty()) {
+        }
-                        int lastIndex = parts.size() - 1;
-                        String last = parts.get(lastIndex);
-                        parts.set(lastIndex, last + " / " + actions);
-                    } else {
-                        parts.add(actions);
-                    }
-                });
         Optional.ofNullable(t.getOrder())
                 .filter(order -> t.getType() == TransitionType.CHOICE || t.getType() == TransitionType.JUNCTION)
                 .map(order -> "(order=" + order + ")")
                 .ifPresent(parts::add);
-
         return String.join(" ", parts);
     }
 }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/out/Scxml.java

@@ -1,5 +1,6 @@
 package click.kamil.springstatemachineexporter.ast.out;
 
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
 
@@ -8,6 +9,7 @@ import java.util.List;
 import java.util.Set;
 
 public class Scxml implements StateMachineExporter {
+
     private static final String LAMBDA = "λ";
 
     @Override
@@ -15,38 +17,33 @@ public class Scxml implements StateMachineExporter {
                          Set<String> startStates,
                          Set<String> endStates,
                          boolean useLambdaGuards) {
-
         StringBuilder sb = new StringBuilder();
         sb.append("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
-        sb.append("<scxml xmlns=\"http://www.w3.org/2005/07/scxml\" version=\"1.0\">\n\n");
+        sb.append("<scxml xmlns=\"http://www.w3.org/2005/07/scxml\" version=\"1.0\" initial=\"");
+        sb.append(startStates.isEmpty() ? "" : simplify(startStates.iterator().next()));
+        sb.append("\">\n");
 
-        // 1. Collect all states
         Set<String> allStates = new HashSet<>();
         for (Transition t : transitions) {
-            if (t.getSourceStates() != null) allStates.addAll(t.getSourceStates());
+            t.getSourceStates().forEach(s -> allStates.add(s.toString()));
-            if (t.getTargetStates() != null) allStates.addAll(t.getTargetStates());
+            t.getTargetStates().forEach(s -> allStates.add(s.toString()));
         }
+        allStates.addAll(startStates);
+        allStates.addAll(endStates);
 
-        // 2. Emit <state> elements
         for (String state : allStates) {
-            String simpleState = simplify(state);
+            sb.append("  <state id=\"").append(simplify(state)).append("\">\n");
-            sb.append("  <state id=\"").append(simpleState).append("\">\n");
-
-            if (startStates.contains(state)) {
-                sb.append("    <initial>\n");
-                sb.append("      <transition target=\"").append(simpleState).append("\"/>\n");
-                sb.append("    </initial>\n");
-            }
 
             for (Transition t : transitions) {
-                if (t.getSourceStates() == null || !t.getSourceStates().contains(state)) continue;
+                if (t.getSourceStates() == null || t.getSourceStates().stream().noneMatch(s -> s.toString().equals(state)))
+                    continue;
 
                 boolean allowNoTarget = t.getType() == TransitionType.JOIN || t.getType() == TransitionType.FORK;
-                List<String> targets;
+                List<State> targets;
 
                 if (t.getTargetStates() == null || t.getTargetStates().isEmpty()) {
                     if (!allowNoTarget) {
-                        targets = List.of(state); // self-loop
+                        targets = List.of(State.of(state)); // self-loop
                     } else {
                         continue;
                     }
@@ -54,15 +51,11 @@ public class Scxml implements StateMachineExporter {
                     targets = t.getTargetStates();
                 }
 
-                for (String rawTarget : targets) {
+                for (State target : targets) {
-                    String target = simplify(rawTarget);
-                    if (target.isEmpty()) continue;
-
                     sb.append(renderTransition(t, target, useLambdaGuards));
                 }
             }
-
+            sb.append("  </state>\n");
-            sb.append("  </state>\n\n");
         }
 
         sb.append("</scxml>\n");
@@ -74,29 +67,14 @@ public class Scxml implements StateMachineExporter {
         return ".scxml.xml";
     }
 
-    private static String getGuardText(Transition t, boolean useLambdaGuards) {
+    private String renderTransition(Transition t, State target, boolean useLambdaGuards) {
-        if (t.getGuard() == null || t.getGuard().expression().isBlank()) return "";
+        StringBuilder sb = new StringBuilder();
-        if (useLambdaGuards && t.getGuard().isLambda()) {
+        sb.append("    <transition target=\"").append(simplify(target.toString())).append("\"");
-            return LAMBDA;
-        }
-        return t.getGuard().expression().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim();
-    }
-
-    private static String escapeXml(String s) {
-        if (s == null) return "";
-        return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
-                .replace("\"", "&quot;").replace("'", "&apos;");
-    }
-
-    private String renderTransition(Transition t, String target, boolean useLambdaGuards) {
-        StringBuilder sb = new StringBuilder("    <transition");
 
         if (t.getEvent() != null && !t.getEvent().isBlank()) {
-            sb.append(" event=\"").append(simplify(t.getEvent())).append("\"");
+            sb.append(" event=\"").append(escapeXml(t.getEvent())).append("\"");
         }
 
-        sb.append(" target=\"").append(target).append("\"");
-
         String guard = getGuardText(t, useLambdaGuards);
         if (!guard.isBlank()) {
             sb.append(" cond=\"").append(escapeXml(guard)).append("\"");
@@ -113,4 +91,20 @@ public class Scxml implements StateMachineExporter {
 
         return sb.toString();
     }
+
+    private static String getGuardText(Transition t, boolean useLambdaGuards) {
+        if (t.getGuard() == null || t.getGuard().expression().isBlank())
+            return "";
+        if (useLambdaGuards && t.getGuard().isLambda()) {
+            return LAMBDA;
+        }
+        return t.getGuard().expression().replaceAll("[\\n\\r]+", " ").replaceAll("\\s{2,}", " ").trim();
+    }
+
+    private static String escapeXml(String s) {
+        if (s == null)
+            return "";
+        return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
+                .replace("\"", "&quot;").replace("'", "&apos;");
+    }
 }

state_machine_exporter/src/main/java/click/kamil/springstatemachineexporter/ast/out/StateMachineExporter.java

@@ -13,9 +13,11 @@ public interface StateMachineExporter {
 
     // Helper: extract last enum part
     default String simplify(String full) {
-        if (full == null) return "";
+        if (full == null)
+            return "";
         int dot = full.lastIndexOf('.');
         return dot >= 0 ? full.substring(dot + 1) : full;
     }
+
     String getFileExtension();  // e.g. ".dot", ".plantuml.dot", ".scxml.xml"
 }

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/MainTest.java

@@ -1,11 +0,0 @@
-package click.kamil.springstatemachineexporter;
-
-import org.junit.jupiter.api.Test;
-
-class MainTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/RegressionTest.java

@@ -114,19 +114,19 @@ public class RegressionTest {
         String actualBaseName = actualDir.getFileName().toString();
 
         for (StateMachineExporter exporter : outputs) {
-        String fileName = actualBaseName + exporter.getFileExtension();
+            String fileName = actualBaseName + exporter.getFileExtension();
-        String goldenFileName = targetBaseName + exporter.getFileExtension();
+            String goldenFileName = targetBaseName + exporter.getFileExtension();
 
-        Path actualFile = actualDir.resolve(fileName);
+            Path actualFile = actualDir.resolve(fileName);
-        Path goldenFile = scenario.goldenPath().resolve(goldenFileName);
+            Path goldenFile = scenario.goldenPath().resolve(goldenFileName);
 
-        assertThat(actualFile)
+            assertThat(actualFile)
-                .as("File %s should exist in output", fileName)
+                    .as("File %s should exist in output", fileName)
-                .exists();
+                    .exists();
 
-        assertThat(actualFile)
+            assertThat(actualFile)
-                .as("Content mismatch for %s", fileName)
+                    .as("Content mismatch for %s", fileName)
-                .hasSameTextualContentAs(goldenFile);
+                    .content().isEqualTo(Files.readString(goldenFile));
-        }
         }
+    }
 }

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/app/AstTransitionParserTest.java

@@ -2,9 +2,16 @@ package click.kamil.springstatemachineexporter.ast.app;
 
 import click.kamil.springstatemachineexporter.ast.app.domain.Action;
 import click.kamil.springstatemachineexporter.ast.app.domain.Guard;
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import click.kamil.springstatemachineexporter.ast.app.domain.TransitionType;
-import org.eclipse.jdt.core.dom.*;
+import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
+import org.eclipse.jdt.core.dom.AST;
+import org.eclipse.jdt.core.dom.ASTParser;
+import org.eclipse.jdt.core.dom.CompilationUnit;
+import org.eclipse.jdt.core.dom.MethodDeclaration;
+import org.eclipse.jdt.core.dom.TypeDeclaration;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import java.util.List;
@@ -13,6 +20,13 @@ import static org.assertj.core.api.Assertions.assertThat;
 
 class AstTransitionParserTest {
 
+    private CodebaseContext context;
+
+    @BeforeEach
+    void setUp() {
+        context = new CodebaseContext();
+    }
+
     @Test
     void shouldParseSimpleWithExternalTransition() {
         String source = """
@@ -24,20 +38,59 @@ class AstTransitionParserTest {
                 """;
         MethodDeclaration method = createMethodDeclaration(source);
 
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
         assertThat(transition)
                 .extracting(Transition::getType)
                 .isEqualTo(TransitionType.EXTERNAL);
-        assertThat(transition.getSourceStates()).containsExactly("CREATED");
+        assertThat(transition.getSourceStates()).extracting(State::toString).containsExactly("CREATED");
-        assertThat(transition.getTargetStates()).containsExactly("PAID");
+        assertThat(transition.getTargetStates()).extracting(State::toString).containsExactly("PAID");
         assertThat(transition)
                 .extracting(Transition::getEvent)
                 .isEqualTo("PAY");
     }
 
+    @Test
+    void shouldResolveStaticImportsForStates() {
+        String source = """
+                import static com.example.OrderStates.CREATED;
+                import static com.example.OrderStates.PAID;
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
+                        transitions.withExternal().source(CREATED).target(PAID).event("PAY");
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
+
+        assertThat(transitions).hasSize(1);
+        System.out.println("Actual source state: " + transitions.getFirst().getSourceStates().get(0));
+        assertThat(transitions.getFirst().getSourceStates()).extracting(State::toString).containsExactly("com.example.OrderStates.CREATED");
+        assertThat(transitions.getFirst().getTargetStates()).extracting(State::toString).containsExactly("com.example.OrderStates.PAID");
+    }
+
+    @Test
+    void shouldResolveQualifiedNamesForStates() {
+        String source = """
+                public class TestClass {
+                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
+                        transitions.withExternal().source(OrderStates.CREATED).target(OrderStates.PAID);
+                    }
+                }
+                """;
+        MethodDeclaration method = createMethodDeclaration(source);
+
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
+
+        assertThat(transitions).hasSize(1);
+        // Without full scan, OrderStates is treated as a simple qualifier
+        assertThat(transitions.getFirst().getSourceStates()).extracting(State::toString).containsExactly("OrderStates.CREATED");
+    }
+
     @Test
     void shouldParseTransitionWithGuard() {
         String source = """
@@ -49,78 +102,36 @@ class AstTransitionParserTest {
                 """;
         MethodDeclaration method = createMethodDeclaration(source);
 
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
         assertThat(transition.getGuard())
                 .extracting(Guard::expression)
                 .isEqualTo("amount > 0");
-        assertThat(transition.getGuard())
-                .extracting(Guard::isLambda)
-                .isEqualTo(false);
     }
 
     @Test
-    void shouldParseTransitionWithLambdaGuard() {
+    void shouldParseTransitionWithLambdaGuardAndAction() {
         String source = """
                 public class TestClass {
                     public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").guard(context -> context.getExtendedState().getVariables().get("amount") > 0);
+                        transitions.withExternal()
+                            .source("CREATED")
+                            .target("PAID")
+                            .event("PAY")
+                            .guard(context -> true)
+                            .action(context -> System.out.println("Action"));
                     }
                 }
                 """;
         MethodDeclaration method = createMethodDeclaration(source);
 
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
 
         assertThat(transitions).hasSize(1);
         Transition transition = transitions.getFirst();
-        assertThat(transition.getGuard().expression()).contains("context -> context.getExtendedState().getVariables().get(\"amount\") > 0");
         assertThat(transition.getGuard().isLambda()).isTrue();
-    }
-
-    @Test
-    void shouldParseTransitionWithAction() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").action("processPayment()");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition.getActions())
-                .extracting(Action::expression)
-                .containsExactly("processPayment()");
-        assertThat(transition.getActions())
-                .extracting(Action::isLambda)
-                .containsExactly(false);
-    }
-
-    @Test
-    void shouldParseTransitionWithLambdaAction() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").action(context -> System.out.println("Payment processed"));
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition.getActions())
-                .extracting(Action::expression)
-                .anyMatch(s -> s.contains("context -> System.out.println(\"Payment processed\")"));
         assertThat(transition.getActions())
                 .extracting(Action::isLambda)
                 .containsExactly(true);
@@ -137,29 +148,9 @@ class AstTransitionParserTest {
                 """;
         MethodDeclaration method = createMethodDeclaration(source);
 
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
 
         assertThat(transitions).hasSize(2);
-        
-        Transition first = transitions.getFirst();
-        assertThat(first)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.EXTERNAL);
-        assertThat(first.getSourceStates()).containsExactly("CREATED");
-        assertThat(first.getTargetStates()).containsExactly("PAID");
-        assertThat(first)
-                .extracting(Transition::getEvent)
-                .isEqualTo("PAY");
-        
-        Transition second = transitions.get(1);
-        assertThat(second)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.EXTERNAL);
-        assertThat(second.getSourceStates()).containsExactly("PAID");
-        assertThat(second.getTargetStates()).containsExactly("SHIPPED");
-        assertThat(second)
-                .extracting(Transition::getEvent)
-                .isEqualTo("SHIP");
     }
 
     @Test
@@ -173,321 +164,34 @@ class AstTransitionParserTest {
                 """;
         MethodDeclaration method = createMethodDeclaration(source);
 
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
 
         assertThat(transitions).hasSize(3);
-        
+        assertThat(transitions.getFirst().getSourceStates()).extracting(State::toString).containsExactly("PAID");
-        Transition first = transitions.getFirst();
-        assertThat(first)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.CHOICE);
-        assertThat(first.getSourceStates()).containsExactly("PAID");
-        assertThat(first.getTargetStates()).containsExactly("SHIPPED");
-        assertThat(first.getGuard())
-                .extracting(Guard::expression)
-                .isEqualTo("isExpress()");
-        assertThat(first)
-                .extracting(Transition::getOrder)
-                .isEqualTo(0);
-        
-        Transition second = transitions.get(1);
-        assertThat(second)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.CHOICE);
-        assertThat(second.getSourceStates()).containsExactly("PAID");
-        assertThat(second.getTargetStates()).containsExactly("CANCELLED");
-        assertThat(second.getGuard())
-                .extracting(Guard::expression)
-                .isEqualTo("isCancelled()");
-        assertThat(second)
-                .extracting(Transition::getOrder)
-                .isEqualTo(1);
-        
-        Transition third = transitions.get(2);
-        assertThat(third)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.CHOICE);
-        assertThat(third.getSourceStates()).containsExactly("PAID");
-        assertThat(third.getTargetStates()).containsExactly("PENDING");
-        assertThat(third)
-                .extracting(Transition::getOrder)
-                .isEqualTo(2);
     }
 
     @Test
-    void shouldParseWithJunctionTransition() {
+    void shouldDetectAnonymousClassAsLambda() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withJunction().source("PAID").first("SHIPPED", "isExpress()").then("CANCELLED", "isCancelled()").last("PENDING");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(3);
-        assertThat(transitions.getFirst())
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.JUNCTION);
-        assertThat(transitions.get(1))
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.JUNCTION);
-        assertThat(transitions.get(2))
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.JUNCTION);
-    }
-
-    @Test
-    void shouldParseWithInternalTransition() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withInternal().source("PAID").event("UPDATE").action("updateStatus()");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.INTERNAL);
-        assertThat(transition.getSourceStates()).containsExactly("PAID");
-        assertThat(transition)
-                .extracting(Transition::getEvent)
-                .isEqualTo("UPDATE");
-        assertThat(transition.getActions())
-                .extracting(Action::expression)
-                .containsExactly("updateStatus()");
-    }
-
-    @Test
-    void shouldParseWithLocalTransition() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withLocal().source("PAID").target("PROCESSING").event("PROCESS");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.LOCAL);
-        assertThat(transition.getSourceStates()).containsExactly("PAID");
-        assertThat(transition.getTargetStates()).containsExactly("PROCESSING");
-        assertThat(transition)
-                .extracting(Transition::getEvent)
-                .isEqualTo("PROCESS");
-    }
-
-    @Test
-    void shouldParseWithForkTransition() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withFork().source("PAID").target("SHIPPING").target("BILLING");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.FORK);
-        assertThat(transition.getSourceStates()).containsExactly("PAID");
-        assertThat(transition.getTargetStates()).containsExactlyInAnyOrder("SHIPPING", "BILLING");
-    }
-
-    @Test
-    void shouldParseWithJoinTransition() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withJoin().source("SHIPPING").source("BILLING").target("COMPLETED");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition)
-                .extracting(Transition::getType)
-                .isEqualTo(TransitionType.JOIN);
-        assertThat(transition.getSourceStates()).containsExactlyInAnyOrder("SHIPPING", "BILLING");
-        assertThat(transition.getTargetStates()).containsExactly("COMPLETED");
-    }
-
-    @Test
-    void shouldReturnEmptyList_whenMethodIsNull() {
-        List<Transition> transitions = AstTransitionParser.parseTransitions(null);
-
-        assertThat(transitions).isEmpty();
-    }
-
-    @Test
-    void shouldReturnEmptyList_whenMethodHasNoBody() {
-        String source = """
-                public class TestClass {
-                    public void configure() {
-                        // Empty method
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).isEmpty();
-    }
-
-    @Test
-    void shouldReturnEmptyList_whenMethodHasNoTransitionStatements() {
-        String source = """
-                public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                    System.out.println("No transitions here");
-                    int x = 5;
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).isEmpty();
-    }
-
-    @Test
-    void shouldParseTransitionWithMultipleSources() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withExternal().source("CREATED").source("PENDING").target("PAID").event("PAY");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition.getSourceStates()).containsExactlyInAnyOrder("CREATED", "PENDING");
-        assertThat(transition.getTargetStates()).containsExactly("PAID");
-    }
-
-    @Test
-    void shouldParseTransitionWithMultipleTargets() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withExternal().source("PAID").target("SHIPPED").target("CANCELLED").event("PROCESS");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition.getSourceStates()).containsExactly("PAID");
-        assertThat(transition.getTargetStates()).containsExactlyInAnyOrder("SHIPPED", "CANCELLED");
-    }
-
-    @Test
-    void shouldParseChoiceTransitionWithAction() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions
-                            .withChoice()
-                            .source("PAID")
-                            .first("SHIPPED", "isExpress()", "shipExpress()")
-                            .then("CANCELLED", "isCancelled()", "cancelOrder()")
-                            .last("PENDING", "defaultAction()");
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(3);
-        assertThat(transitions.getFirst().getActions())
-                .extracting(Action::expression)
-                .containsExactly("shipExpress()");
-        assertThat(transitions.get(1).getActions())
-                .extracting(Action::expression)
-                .containsExactly("cancelOrder()");
-        assertThat(transitions.get(2).getActions())
-                .extracting(Action::expression)
-                .containsExactly("defaultAction()");
-    }
-
-    @Test
-    void shouldDetectLambdaExpressions() {
-        String source = """
-                public class TestClass {
-                    public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
-                        transitions.withExternal().source("CREATED").target("PAID").event("PAY").guard(context -> context.getExtendedState().getVariables().get("amount") > 0).action(context -> System.out.println("Payment processed"));
-                    }
-                }
-                """;
-        MethodDeclaration method = createMethodDeclaration(source);
-
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
-
-        assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
-        assertThat(transition.getGuard().isLambda()).isTrue();
-        assertThat(transition.getActions())
-                .extracting(Action::isLambda)
-                .containsExactly(true);
-    }
-
-    @Test
-    void shouldDetectAnonymousClassExpressions() {
         String source = """
                 public class TestClass {
                     public void configure(StateMachineTransitionConfigurer<String, String> transitions) throws Exception {
                         transitions
                             .withExternal()
-                            .source("CREATED")
+                            .source("S1")
-                            .target("PAID")
+                            .target("S2")
-                            .event("PAY")
                             .action(new Action<String, String>() {
                                 @Override
-                                public void execute(StateContext<String, String> context) {
+                                public void execute(StateContext<String, String> context) {}
-                                    System.out.println("Payment processed");
-                                }
                             });
                     }
                 }
                 """;
         MethodDeclaration method = createMethodDeclaration(source);
 
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
 
         assertThat(transitions).hasSize(1);
-        Transition transition = transitions.getFirst();
+        assertThat(transitions.getFirst().getActions())
-        assertThat(transition.getActions())
                 .extracting(Action::isLambda)
                 .containsExactly(true);
     }
@@ -504,7 +208,7 @@ class AstTransitionParserTest {
                 """;
         MethodDeclaration method = createMethodDeclaration(source);
 
-        List<Transition> transitions = AstTransitionParser.parseTransitions(method);
+        List<Transition> transitions = AstTransitionParser.parseTransitions(method, context);
 
         assertThat(transitions).hasSize(2);
         assertThat(transitions.get(0).getEvent()).isEqualTo("E1");
@@ -526,4 +230,4 @@ class AstTransitionParserTest {
         }
         return null;
     }
-} 
+}

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/app/StateMachineAggregatorTest.java

@@ -1,8 +1,9 @@
 package click.kamil.springstatemachineexporter.ast.app;
 
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
-import org.eclipse.jdt.core.dom.*;
+import org.eclipse.jdt.core.dom.TypeDeclaration;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -10,9 +11,7 @@ import org.junit.jupiter.api.io.TempDir;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -50,7 +49,7 @@ class StateMachineAggregatorTest {
                     }
                 }
                 """;
-        
+
         writeClass("BaseConfig", baseSource);
         writeClass("ChildConfig", childSource);
         context.scan(tempDir);
@@ -118,8 +117,8 @@ class StateMachineAggregatorTest {
         List<Transition> transitions = aggregator.aggregateTransitions(td);
 
         assertThat(transitions).hasSize(1);
-        assertThat(transitions.get(0).getSourceStates()).containsExactly("START");
+        assertThat(transitions.get(0).getSourceStates()).extracting(State::toString).containsExactly("START");
-        assertThat(transitions.get(0).getTargetStates()).containsExactly("END");
+        assertThat(transitions.get(0).getTargetStates()).extracting(State::toString).containsExactly("END");
         assertThat(transitions.get(0).getEvent()).isEqualTo("GO");
     }
 
@@ -438,6 +437,157 @@ class StateMachineAggregatorTest {
         assertThat(transitions.get(0).getEvent()).isEqualTo("TRY_E1");
     }
 
+    @Test
+    void shouldHandleComplexStateConfigurationWithMultipleBlocks() throws IOException {
+        String source = """
+                package com.example;
+                public class MultiBlockConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineStateConfigurer states) throws Exception {
+                        states
+                            .withStates()
+                            .initial("START")
+                            .state("S1")
+                            .end("END_S1")
+                            .and()
+                            .withStates()
+                            .parent("S1")
+                            .initial("SUB_START")
+                            .end("SUB_END");
+                    }
+                }
+                """;
+
+        writeClass("MultiBlockConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("MultiBlockConfig");
+        aggregator.aggregateStates(td);
+
+        assertThat(aggregator.getInitialStates()).containsExactlyInAnyOrder("START", "SUB_START");
+        assertThat(aggregator.getEndStates()).containsExactlyInAnyOrder("END_S1", "SUB_END");
+    }
+
+    @Test
+    void shouldIgnoreStatesInsideRegion() throws IOException {
+        String source = """
+                package com.example;
+                public class RegionStateConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineStateConfigurer states) throws Exception {
+                        states
+                            .withStates()
+                            .initial("CREATED")
+                            .region("PROCESSING")
+                            .and()
+                            .withStates()
+                            .parent("PROCESSING")
+                            .initial("SHIPPING")
+                            .end("COMPLETED");
+                    }
+                }
+                """;
+
+        writeClass("RegionStateConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("RegionStateConfig");
+        aggregator.aggregateStates(td);
+
+        assertThat(aggregator.getInitialStates()).containsExactly("CREATED");
+        assertThat(aggregator.getEndStates()).containsExactly("COMPLETED");
+    }
+
+    @Test
+    void shouldHandleEmptyConfigureMethod() throws IOException {
+        String source = """
+                package com.example;
+                public class EmptyConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineStateConfigurer states) throws Exception {}
+                }
+                """;
+
+        writeClass("EmptyConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("EmptyConfig");
+        aggregator.aggregateStates(td);
+
+        assertThat(aggregator.getInitialStates()).isEmpty();
+    }
+
+    @Test
+    void shouldHandleMultipleMethodsHandlingTransitions() throws IOException {
+        String source = """
+                package com.example;
+                public class MultiMethodConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
+                        transitions.withExternal().source("S1").target("S2").event("E1");
+                        otherTransitions(transitions);
+                    }
+                    private void otherTransitions(StateMachineTransitionConfigurer transitions) {
+                        transitions.withExternal().source("S2").target("S3").event("E2");
+                    }
+                }
+                """;
+
+        writeClass("MultiMethodConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("MultiMethodConfig");
+        List<Transition> transitions = aggregator.aggregateTransitions(td);
+
+        assertThat(transitions).hasSize(2);
+        assertThat(transitions).anyMatch(t -> t.getEvent().equals("E1"));
+        assertThat(transitions).anyMatch(t -> t.getEvent().equals("E2"));
+    }
+
+    @Test
+    void shouldHandleMethodCallsWithEnumArguments() throws IOException {
+        String source = """
+                package com.example;
+                public class EnumConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
+                        addTransition(transitions, States.S1, States.S2);
+                    }
+                    private void addTransition(StateMachineTransitionConfigurer t, States s, States e) {
+                        t.withExternal().source(s).target(e).event("ENUM_E1");
+                    }
+                }
+                """;
+
+        writeClass("EnumConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("EnumConfig");
+        List<Transition> transitions = aggregator.aggregateTransitions(td);
+
+        assertThat(transitions).hasSize(1);
+        assertThat(transitions.get(0).getSourceStates()).extracting(State::toString).containsExactly("States.S1");
+    }
+
+    @Test
+    void shouldIgnoreStaticMethodsThatAreNotCalled() throws IOException {
+        String source = """
+                package com.example;
+                public class UncalledStaticConfig extends StateMachineConfigurerAdapter {
+                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
+                        transitions.withExternal().source("S1").target("S2").event("E1");
+                    }
+                    public static void uncalled(StateMachineTransitionConfigurer transitions) {
+                        transitions.withExternal().source("S2").target("S3").event("E2");
+                    }
+                }
+                """;
+
+        writeClass("UncalledStaticConfig", source);
+        context.scan(tempDir);
+
+        TypeDeclaration td = context.getTypeDeclaration("UncalledStaticConfig");
+        List<Transition> transitions = aggregator.aggregateTransitions(td);
+
+        assertThat(transitions).hasSize(1);
+        assertThat(transitions.get(0).getEvent()).isEqualTo("E1");
+    }
+
     private void writeClass(String name, String source) throws IOException {
         Files.writeString(tempDir.resolve(name + ".java"), source);
     }

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/app/TransitionStateUtilsTest.java

@@ -1,5 +1,6 @@
 package click.kamil.springstatemachineexporter.ast.app;
 
+import click.kamil.springstatemachineexporter.ast.app.domain.State;
 import click.kamil.springstatemachineexporter.ast.app.domain.Transition;
 import org.junit.jupiter.api.Test;
 
@@ -186,59 +187,17 @@ class TransitionStateUtilsTest {
     }
 
     @Test
-    void shouldHandleTransitionsWithWhitespaceInStates() {
+    void shouldHandleVariousStateIdFormats() {
         List<Transition> transitions = List.of(
-                createTransition(" CREATED ", " PAID "),
+                createTransition(" OrderStates.START ", "STATE_1"),
-                createTransition(" PAID ", " SHIPPED ")
+                createTransition("STATE_1", "getEndState()")
         );
 
         Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
         Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
 
-        assertThat(startStates).containsExactly(" CREATED ");
+        assertThat(startStates).containsExactly(" OrderStates.START ");
-        assertThat(endStates).containsExactly(" SHIPPED ");
+        assertThat(endStates).containsExactly("getEndState()");
-    }
-
-    @Test
-    void shouldHandleTransitionsWithSpecialCharactersInStates() {
-        List<Transition> transitions = List.of(
-                createTransition("STATE_1", "STATE_2"),
-                createTransition("STATE_2", "STATE_3")
-        );
-
-        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
-        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
-
-        assertThat(startStates).containsExactly("STATE_1");
-        assertThat(endStates).containsExactly("STATE_3");
-    }
-
-    @Test
-    void shouldHandleTransitionsWithEnumValues() {
-        List<Transition> transitions = List.of(
-                createTransition("OrderStates.CREATED", "OrderStates.PAID"),
-                createTransition("OrderStates.PAID", "OrderStates.SHIPPED")
-        );
-
-        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
-        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
-
-        assertThat(startStates).containsExactly("OrderStates.CREATED");
-        assertThat(endStates).containsExactly("OrderStates.SHIPPED");
-    }
-
-    @Test
-    void shouldHandleTransitionsWithMethodCalls() {
-        List<Transition> transitions = List.of(
-                createTransition("getInitialState()", "getPaidState()"),
-                createTransition("getPaidState()", "getShippedState()")
-        );
-
-        Set<String> startStates = TransitionStateUtils.findStartStates(transitions);
-        Set<String> endStates = TransitionStateUtils.findEndStates(transitions);
-
-        assertThat(startStates).containsExactly("getInitialState()");
-        assertThat(endStates).containsExactly("getShippedState()");
     }
 
     @Test
@@ -261,10 +220,10 @@ class TransitionStateUtilsTest {
     private Transition createTransition(String source, String target) {
         Transition transition = new Transition();
         if (source != null) {
-            transition.getSourceStates().add(source);
+            transition.getSourceStates().add(State.of(source));
         }
         if (target != null) {
-            transition.getTargetStates().add(target);
+            transition.getTargetStates().add(State.of(target));
         }
         return transition;
     }
@@ -272,11 +231,13 @@ class TransitionStateUtilsTest {
     private Transition createTransition(List<String> sources, List<String> targets) {
         Transition transition = new Transition();
         if (sources != null) {
-            transition.getSourceStates().addAll(sources);
+            for (String s : sources)
+                transition.getSourceStates().add(State.of(s));
         }
         if (targets != null) {
-            transition.getTargetStates().addAll(targets);
+            for (String t : targets)
+                transition.getTargetStates().add(State.of(t));
         }
         return transition;
     }
-} 
+} 

state_machine_exporter/src/test/java/click/kamil/springstatemachineexporter/ast/common/CodebaseContextTest.java

@@ -109,7 +109,7 @@ class CodebaseContextTest {
                     public StateMachine<String, String> stateMachine() {
                         return null;
                     }
-                    
+                
                     public StateMachine<String, String> notABean() {
                         return null;
                     }
@@ -146,4 +146,103 @@ class CodebaseContextTest {
         assertThat(beanMethods).hasSize(1);
         assertThat(beanMethods.get(0).getName().getIdentifier()).isEqualTo("stateMachine");
     }
+
+    @Test
+    void shouldFindClassWithMultipleAnnotations() throws IOException {
+        String source = """
+                package com.example;
+                @org.springframework.context.annotation.Configuration
+                @click.kamil.MyAnnotation
+                public class MultiAnnotated {}
+                """;
+        Files.writeString(tempDir.resolve("MultiAnnotated.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        assertThat(context.findEntryPointClasses(List.of("Configuration"))).hasSize(1);
+        assertThat(context.findEntryPointClasses(List.of("MyAnnotation"))).hasSize(1);
+    }
+
+    @Test
+    void shouldIgnoreInterfacesWithAnnotations() throws IOException {
+        String source = """
+                package com.example;
+                import org.springframework.context.annotation.Configuration;
+                @Configuration
+                public interface MyInterface {}
+                """;
+        Files.writeString(tempDir.resolve("MyInterface.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        assertThat(context.findEntryPointClasses(List.of("Configuration"))).isEmpty();
+    }
+
+    @Test
+    void shouldFindBeanMethodsWithGenericParameters() throws IOException {
+        String source = """
+                package com.example;
+                import org.springframework.context.annotation.Bean;
+                import org.springframework.statemachine.StateMachine;
+                public class MyApp {
+                    @Bean
+                    public StateMachine<States, Events> stateMachine() {
+                        return null;
+                    }
+                }
+                """;
+        Files.writeString(tempDir.resolve("MyApp.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        var beanMethods = context.findBeanMethodsReturning(java.util.Set.of("StateMachine"));
+        assertThat(beanMethods).hasSize(1);
+    }
+
+    @Test
+    void shouldHandleFilesWithOnlyComments() throws IOException {
+        String source = """
+                // This is just a comment
+                /* And another one */
+                """;
+        Files.writeString(tempDir.resolve("Comments.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        assertThat(context.findEntryPointClasses(List.of("EnableStateMachine"))).isEmpty();
+    }
+
+    @Test
+    void shouldBeCaseSensitiveForAnnotations() throws IOException {
+        String source = """
+                package com.example;
+                @enablestatemachine
+                public class WrongCase {}
+                """;
+        Files.writeString(tempDir.resolve("WrongCase.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        assertThat(context.findEntryPointClasses(List.of("EnableStateMachine"))).isEmpty();
+    }
+
+    @Test
+    void shouldFindClassWithFullyQualifiedAnnotation() throws IOException {
+        String source = """
+                package com.example;
+                @org.springframework.statemachine.config.EnableStateMachine
+                public class FqnAnnotated {}
+                """;
+        Files.writeString(tempDir.resolve("FqnAnnotated.java"), source);
+
+        CodebaseContext context = new CodebaseContext();
+        context.scan(tempDir);
+
+        assertThat(context.findEntryPointClasses(List.of("EnableStateMachine"))).hasSize(1);
+    }
 }

state_machine_exporter_spring_based/src/main/java/click/kamil/SpringExporter.java

@@ -16,14 +16,14 @@ import org.springframework.stereotype.Component;
 import org.w3c.dom.Document;
 import org.w3c.dom.Element;
 
+import java.io.File;
+import java.io.PrintWriter;
 import javax.xml.parsers.DocumentBuilder;
 import javax.xml.parsers.DocumentBuilderFactory;
 import javax.xml.transform.Transformer;
 import javax.xml.transform.TransformerFactory;
 import javax.xml.transform.dom.DOMSource;
 import javax.xml.transform.stream.StreamResult;
-import java.io.File;
-import java.io.PrintWriter;
 
 public class SpringExporter {
 }
@@ -116,7 +116,8 @@ class PlantUmlExporter extends BaseExporter {
                     }
 
                     if (transition.getGuard() != null) {
-                        if (!label.isEmpty()) label += " ";
+                        if (!label.isEmpty())
+                            label += " ";
                         label += "[guard]";
                     }
 

state_machine_exporter_spring_based/src/main/java/click/kamil/StatemachinedemoApplication.java

@@ -7,7 +7,7 @@ import org.springframework.boot.autoconfigure.SpringBootApplication;
 @SpringBootApplication
 //@ComponentScan(basePackages = {"click.kamil."})
 class StatemachinedemoApplication {
-	public static void main(String[] args) {
+    public static void main(String[] args) {
-		SpringApplication.run(StatemachinedemoApplication.class, args);
+        SpringApplication.run(StatemachinedemoApplication.class, args);
-	}
+    }
 }

state_machine_exporter_spring_based/src/test/java/click/kamil/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/complex_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/dynamic_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/enumstate_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/forkjoin_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/inheritance_extra_functions2_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/inheritance_extra_functions3_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/inheritance_extra_functions_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/inheritance_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}

state_machines/simple_state_machine/src/test/java/click/kamil/examples/statemachine/MyTest.java

@@ -1,11 +0,0 @@
-package click.kamil.examples.statemachine;
-
-import org.junit.jupiter.api.Test;
-
-class MyTest {
-
-	@Test
-	void contextLoads() {
-	}
-
-}