package click.kamil.springstatemachineexporter.ast.app;

import click.kamil.springstatemachineexporter.model.Transition;
import click.kamil.springstatemachineexporter.model.TransitionType;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.*;

import java.util.*;

public class StateMachineAdapterParser {
    private final CodebaseContext context;

    private enum AnalysisGoal {TRANSITIONS, STATES}

    public record AdapterStates(Set<String> initialStates, Set<String> endStates) {}

    public record AdapterConfig(
            List<Transition> transitions,
            Set<String> initialStates,
            Set<String> endStates
    ) {}

    public StateMachineAdapterParser(CodebaseContext context) {
        this.context = context;
    }

    public AdapterConfig parse(TypeDeclaration startClass) {
        List<Transition> transitions = parseTransitions(startClass);
        AdapterStates states = parseStates(startClass);
        return new AdapterConfig(transitions, states.initialStates(), states.endStates());
    }

    public List<Transition> parseTransitions(TypeDeclaration startClass) {
        List<Transition> allTransitions = new ArrayList<>();
        Set<MethodVisit> visitedMethods = new HashSet<>();
        aggregateTransitionsRecursive(startClass, startClass, allTransitions, new HashSet<>(), visitedMethods, new HashMap<>());
        return allTransitions;
    }

    public AdapterStates parseStates(TypeDeclaration startClass) {
        Set<String> initialStates = new HashSet<>();
        Set<String> endStates = new HashSet<>();
        Set<MethodVisit> visitedMethods = new HashSet<>();
        aggregateStatesRecursive(startClass, startClass, initialStates, endStates, new HashSet<>(), visitedMethods, new HashMap<>());
        return new AdapterStates(initialStates, endStates);
    }

    private void aggregateTransitionsRecursive(TypeDeclaration currentTd, TypeDeclaration searchStartClass, List<Transition> allTransitions, Set<String> visitedClasses, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
        if (currentTd == null)
            return;
        String fqn = context.getFqn(currentTd);
        if (visitedClasses.contains(fqn))
            return;
        visitedClasses.add(fqn);

        for (MethodDeclaration method : currentTd.getMethods()) {
            if (isConfigureTransitionsMethod(method)) {
                allTransitions.addAll(parseTransitionsWithMethodCalls(method, searchStartClass, visitedMethods, argsMap));
            }
        }

        CompilationUnit cu = (CompilationUnit) currentTd.getRoot();

        // Hierarchy - Superclass
        Type superclassType = currentTd.getSuperclassType();
        if (superclassType != null) {
            TypeDeclaration superclassTd = context.getTypeDeclaration(getSimpleName(superclassType), cu);
            if (superclassTd != null) {
                aggregateTransitionsRecursive(superclassTd, searchStartClass, allTransitions, 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) {
                    aggregateTransitionsRecursive(interfaceTd, searchStartClass, allTransitions, visitedClasses, visitedMethods, argsMap);
                }
            }
        }
    }

    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));
        if (visitedMethods.contains(visit)) {
            return new ArrayList<>();
        }
        visitedMethods.add(visit);

        Block body = method.getBody();
        if (body == null) {
            return new ArrayList<>();
        }

        return parseTransitionsInBlock(body, searchStartClass, visitedMethods, argsMap);
    }

    private List<Transition> parseTransitionsInExpression(Expression expr, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap) {
        List<Transition> transitions = new ArrayList<>();
        if (expr instanceof MethodInvocation mi) {
            if (isTransitionChainEntry(mi)) {
                transitions.addAll(AstTransitionParser.parseTransitionsFromExpression(mi, convertToExpressionMap(argsMap), context));
            } else if (isTransitionChainContinuation(mi, argsMap)) {
                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 {
                // Check if it's a call on a lambda parameter or a variable that resolved to a lambda
                Expression lambdaReceiver = mi.getExpression();
                if (lambdaReceiver instanceof SimpleName sn) {
                    Object resolved = argsMap.get(sn.getIdentifier());
                    if (resolved instanceof LambdaExpression lambda) {
                        transitions.addAll(parseLambdaInvocation(lambda, mi.arguments(), searchStartClass, visitedMethods, argsMap, AnalysisGoal.TRANSITIONS, null, null));
                    }
                }

                MethodDeclaration calledMethod = findMethodInHierarchy(mi.getName().getIdentifier(), searchStartClass);
                if (calledMethod == null) {
                    // Try to resolve method in another class if it's a static call or on a known type
                    Expression miReceiver = mi.getExpression();
                    if (miReceiver instanceof SimpleName sn) {
                        TypeDeclaration otherTd = context.getTypeDeclaration(sn.getIdentifier(), (CompilationUnit) searchStartClass.getRoot());
                        if (otherTd != null) {
                            calledMethod = findMethodInHierarchy(mi.getName().getIdentifier(), otherTd);
                            if (calledMethod != null) {
                                // For calls to other classes, we switch the "searchStartClass" to that class
                                Map<String, Object> nextArgsMap = buildArgsMap(calledMethod, mi, argsMap);
                                transitions.addAll(parseTransitionsWithMethodCalls(calledMethod, otherTd, visitedMethods, nextArgsMap));
                                return transitions; // Handled
                            }
                        }
                    }
                }

                if (calledMethod != null) {
                    Map<String, Object> nextArgsMap = buildArgsMap(calledMethod, mi, argsMap);
                    transitions.addAll(parseTransitionsWithMethodCalls(calledMethod, searchStartClass, visitedMethods, nextArgsMap));
                }
            }
        }
        return transitions;
    }

    private boolean isForEach(MethodInvocation mi) {
        return mi.getName().getIdentifier().equals("forEach");
    }

    private List<Transition> parseLambdaInvocation(LambdaExpression lambda, List<?> arguments, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, AnalysisGoal goal, Set<String> initialStates, Set<String> endStates) {
        Map<String, Object> lambdaArgsMap = new HashMap<>(argsMap);
        List<?> params = lambda.parameters();
        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 (!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);
            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);
            else if (lambda.getBody() instanceof Expression lambdaExpr)
                parseStatesInExpression(lambdaExpr, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
        }
        return Collections.emptyList();
    }

    private List<Transition> parseForEach(MethodInvocation forEachMi, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, AnalysisGoal goal, Set<String> initialStates, Set<String> endStates) {
        List<Transition> transitions = new ArrayList<>();
        Expression receiver = forEachMi.getExpression();
        List<Expression> elements = unrollCollection(receiver, argsMap);
        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 (goal == AnalysisGoal.TRANSITIONS) {
                    if (lambda.getBody() instanceof Block block)
                        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);
                    else if (lambda.getBody() instanceof Expression lambdaExpr)
                        parseStatesInExpression(lambdaExpr, searchStartClass, initialStates, endStates, visitedMethods, lambdaArgsMap);
                }
            }
        }
        return transitions;
    }

    @SuppressWarnings("unchecked")
    private List<Transition> parseEnhancedForLoop(EnhancedForStatement efs, TypeDeclaration searchStartClass, Set<MethodVisit> visitedMethods, Map<String, Object> argsMap, AnalysisGoal goal, Set<String> initialStates, Set<String> endStates) {
        List<Transition> transitions = new ArrayList<>();
        List<Expression> elements = unrollCollection(efs.getExpression(), argsMap);
        String paramName = efs.getParameter().getName().getIdentifier();

        for (Expression element : elements) {
            Map<String, Object> loopArgsMap = new HashMap<>(argsMap);
            loopArgsMap.put(paramName, element);

            if (goal == AnalysisGoal.TRANSITIONS) {
                if (efs.getBody() instanceof Block block)
                    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);
                else if (efs.getBody() instanceof ExpressionStatement es)
                    parseStatesInExpression(es.getExpression(), searchStartClass, initialStates, endStates, visitedMethods, loopArgsMap);
            }
        }
        return transitions;
    }

    @SuppressWarnings("unchecked")
    private List<Expression> unrollCollection(Expression expr, Map<String, Object> argsMap) {
        Set<Expression> visited = new HashSet<>();
        Expression current = expr;
        while (current != null) {
            if (!visited.add(current))
                break;

            if (current instanceof MethodInvocation mi) {
                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"))) {
                    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;
                        }
                    }
                    return (List<Expression>) mi.arguments();
                }
            }
            if (current instanceof SimpleName sn) {
                Object resolved = argsMap.get(sn.getIdentifier());
                if (resolved instanceof List list)
                    return (List<Expression>) list;
                if (resolved instanceof Expression e) {
                    current = e;
                    continue;
                }
            }
            if (current instanceof ArrayInitializer ai)
                return (List<Expression>) ai.expressions();
            break;
        }
        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();
        }
        return "";
    }

    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));
            else if (stmt instanceof VariableDeclarationStatement vds) {
                for (Object fragmentObj : vds.fragments()) {
                    if (fragmentObj instanceof VariableDeclarationFragment fragment) {
                        Expression initializer = fragment.getInitializer();
                        if (initializer != null) {
                            Object resolved = resolve(initializer, 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));
        }
        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);
            else if (stmt instanceof VariableDeclarationStatement vds) {
                for (Object fragmentObj : vds.fragments()) {
                    if (fragmentObj instanceof VariableDeclarationFragment fragment) {
                        Expression initializer = fragment.getInitializer();
                        if (initializer != null) {
                            Object resolved = resolve(initializer, 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);
        }
    }

    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);
            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();
                if (receiver instanceof SimpleName sn) {
                    Object resolved = argsMap.get(sn.getIdentifier());
                    if (resolved instanceof LambdaExpression lambda) {
                        parseLambdaInvocation(lambda, mi.arguments(), searchStartClass, visitedMethods, argsMap, AnalysisGoal.STATES, initialStates, endStates);
                    }
                }

                MethodDeclaration calledMethod = findMethodInHierarchy(mi.getName().getIdentifier(), searchStartClass);
                if (calledMethod != null) {
                    Map<String, Object> nextArgsMap = buildArgsMap(calledMethod, mi, argsMap);
                    parseStatesWithMethodCalls(calledMethod, searchStartClass, initialStates, endStates, visitedMethods, nextArgsMap);
                }
            }
        }
    }

    private Map<String, Object> buildArgsMap(MethodDeclaration method, MethodInvocation call, Map<String, Object> currentArgsMap) {
        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);
                    if (resolved instanceof List<?> list) {
                        for (Object item : list) {
                            if (item instanceof Expression e)
                                varargList.add(e);
                        }
                    } else if (resolved instanceof Expression e) {
                        varargList.add(e);
                    }
                }
                map.put(paramName, varargList);
            } else if (i < args.size()) {
                map.put(paramName, resolve((Expression) args.get(i), currentArgsMap));
            }
        }
        return map;
    }

    private Object resolve(Expression expr, Map<String, Object> argsMap) {
        if (expr instanceof NullLiteral)
            return null;
        Set<Expression> visited = new HashSet<>();
        Object current = expr;
        while (current instanceof SimpleName sn) {
            if (!visited.add((SimpleName) current))
                break;
            Object resolved = argsMap.get(sn.getIdentifier());
            if (resolved == null || resolved == current)
                break;
            current = resolved;
        }
        return current;
    }

    private Expression resolveExpression(Expression expr, Map<String, Object> argsMap) {
        Object resolved = resolve(expr, argsMap);
        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);
        }
        return map;
    }

    private boolean isTransitionChainEntry(MethodInvocation mi) {
        Expression current = mi;
        while (current instanceof MethodInvocation call) {
            String name = call.getName().getIdentifier();
            if (TransitionType.fromMethodName(name).isPresent())
                return true;
            current = call.getExpression();
        }
        return false;
    }

    private boolean isTransitionChainContinuation(MethodInvocation mi, Map<String, Object> argsMap) {
        Expression current = mi;
        Set<Expression> visited = new HashSet<>();
        while (current instanceof MethodInvocation call) {
            if (!visited.add(current))
                break;
            String name = call.getName().getIdentifier();
            if (Set.of("and", "source", "target", "event", "first", "then", "last").contains(name)) {
                Expression receiver = call.getExpression();
                if (receiver instanceof MethodInvocation next) {
                    current = next;
                    continue;
                }
                Object resolved = resolve(receiver, argsMap);
                if (resolved instanceof MethodInvocation next) {
                    current = next;
                    continue;
                }
                return resolved == null; // null means it's a parameter we track
            }
            if (TransitionType.fromMethodName(name).isPresent())
                return true;
            break;
        }
        return false;
    }

    private MethodDeclaration findMethodInHierarchy(String methodName, TypeDeclaration td) {
        return findMethodInHierarchy(methodName, td, new HashSet<>());
    }

    private MethodDeclaration findMethodInHierarchy(String methodName, TypeDeclaration td, Set<String> visited) {
        if (td == null)
            return null;

        String fqn = context.getFqn(td);
        if (visited.contains(fqn))
            return null;
        visited.add(fqn);

        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), (CompilationUnit) td.getRoot());
            return findMethodInHierarchy(methodName, superclassTd, visited);
        }
        return null;
    }

    private boolean isRelevantMethod(MethodDeclaration method) {
        return true;
    }

    private MethodDeclaration findConfigureTransitionsMethod(TypeDeclaration td) {
        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().isEmpty())
            return false;
        return ((SingleVariableDeclaration) method.parameters().get(0)).getType().toString().contains("StateMachineTransitionConfigurer");
    }

    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)
            return;
        String fqn = context.getFqn(currentTd);
        if (visitedClasses.contains(fqn))
            return;
        visitedClasses.add(fqn);

        for (MethodDeclaration method : currentTd.getMethods()) {
            if (isConfigureStatesMethod(method)) {
                parseStatesWithMethodCalls(method, searchStartClass, initialStates, endStates, visitedMethods, argsMap);
            }
        }

        CompilationUnit cu = (CompilationUnit) currentTd.getRoot();

        // Hierarchy - Superclass
        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);
        }

        // 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);
            }
        }
    }

    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;
        visitedMethods.add(visit);
        Block body = method.getBody();
        if (body == null)
            return;

        parseStatesInBlock(body, searchStartClass, initialStates, endStates, visitedMethods, argsMap);
    }

    private boolean isWithStatesChain(MethodInvocation mi, Map<String, Object> argsMap) {
        Expression current = mi;
        Set<Expression> visited = new HashSet<>();
        while (current instanceof MethodInvocation call) {
            if (!visited.add(current))
                break;
            if (call.getName().getIdentifier().equals("withStates"))
                return true;
            Expression receiver = call.getExpression();
            if (receiver instanceof MethodInvocation next) {
                current = next;
                continue;
            }
            Object resolved = resolve(receiver, argsMap);
            if (resolved instanceof MethodInvocation next) {
                current = next;
                continue;
            }
            break;
        }
        return false;
    }

    private void parseStateChain(MethodInvocation mi, Set<String> initialStates, Set<String> endStates, Map<String, Object> argsMap) {
        List<MethodInvocation> chain = new ArrayList<>();
        Expression current = mi;
        while (current instanceof MethodInvocation m) {
            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;
            Expression arg = resolveExpression((Expression) call.arguments().get(0), argsMap);
            if ("initial".equals(methodName)) {
                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;
            expr = parent.getExpression();
        }
        return false;
    }

    private MethodDeclaration findConfigureStatesMethod(TypeDeclaration td) {
        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().isEmpty())
            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.isParameterizedType())
            return getSimpleName(((ParameterizedType) type).getType());
        return type.toString();
    }
}
