package click.kamil.springstatemachineexporter.ast.common;

import click.kamil.springstatemachineexporter.analysis.index.AccessorIndex;
import click.kamil.springstatemachineexporter.analysis.index.AccessorIndexBuilder;
import click.kamil.springstatemachineexporter.analysis.index.FieldTypeIndex;
import click.kamil.springstatemachineexporter.analysis.index.FieldTypeIndexBuilder;
import click.kamil.springstatemachineexporter.analysis.index.GetterChainIndex;
import click.kamil.springstatemachineexporter.analysis.index.GetterChainIndexBuilder;
import click.kamil.springstatemachineexporter.analysis.index.IndexedFieldConstantsIndex;
import click.kamil.springstatemachineexporter.analysis.index.IndexedFieldConstantsIndexBuilder;
import click.kamil.springstatemachineexporter.analysis.index.PolymorphicAccessorIndex;
import click.kamil.springstatemachineexporter.analysis.index.PolymorphicAccessorIndexBuilder;
import click.kamil.springstatemachineexporter.analysis.index.PolymorphicMethodExpansionIndex;
import click.kamil.springstatemachineexporter.analysis.index.PolymorphicMethodExpansionIndexBuilder;
import click.kamil.springstatemachineexporter.analysis.pipeline.ClassCompatibilityCache;
import click.kamil.springstatemachineexporter.analysis.model.LibraryHint;
import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import click.kamil.springstatemachineexporter.analysis.resolver.PropertyResolver;
import click.kamil.springstatemachineexporter.analysis.util.PlaceholderResolver;
import click.kamil.springstatemachineexporter.model.State;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.eclipse.jdt.core.JavaCore;
import org.eclipse.jdt.core.dom.*;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Consumer;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class CodebaseContext {
    private final Map<String, CompilationUnit> classes = new HashMap<>();
    private final Map<String, TypeDeclaration> typeDeclarations = new HashMap<>();
    private final Map<String, org.eclipse.jdt.core.dom.RecordDeclaration> recordDeclarations = new HashMap<>();
    private final Map<String, Path> classPaths = new HashMap<>();
    private final Map<String, String> simpleNameToFqn = new HashMap<>();
    private final Set<String> ambiguousSimpleNames = new HashSet<>();
    private final ConstantResolver constantResolver = new ConstantResolver();
    private final StateResolver stateResolver = new StateResolver(this);
    private final PropertyResolver propertyResolver = new PropertyResolver();

    public ConstantResolver getConstantResolver() {
        return constantResolver;
    }
    private final List<String> activeProfiles = new ArrayList<>();
    private Map<String, Map<String, String>> allProperties = new HashMap<>();
    private List<LibraryHint> libraryHints = new ArrayList<>();
    private final ObjectMapper objectMapper = new ObjectMapper();
    private Path projectRoot;
    private final Map<String, Object> cache = new HashMap<>();
    private AccessorIndex accessorIndex = AccessorIndex.empty();
    private PolymorphicAccessorIndex polymorphicAccessorIndex = PolymorphicAccessorIndex.empty();
    private GetterChainIndex getterChainIndex = GetterChainIndex.empty();
    private IndexedFieldConstantsIndex indexedFieldConstantsIndex = IndexedFieldConstantsIndex.empty();
    private FieldTypeIndex fieldTypeIndex = FieldTypeIndex.empty();
    private PolymorphicMethodExpansionIndex polymorphicMethodExpansionIndex =
            PolymorphicMethodExpansionIndex.empty();
    private final ClassCompatibilityCache classCompatibilityCache = new ClassCompatibilityCache();
    private final Map<String, MethodDeclaration> methodDeclarationCache = new HashMap<>();
    private final Map<String, String> superclassFqnCache = new HashMap<>();
    private final Map<String, List<String>> implementationsCache = new HashMap<>();
    private final Map<String, Boolean> subtypeOrSameCache = new HashMap<>();
    private boolean inlineAccessors = true;

    public Map<String, Object> getCache() {
        return cache;
    }

    public AccessorIndex getAccessorIndex() {
        return accessorIndex;
    }

    public PolymorphicAccessorIndex getPolymorphicAccessorIndex() {
        return polymorphicAccessorIndex;
    }

    public GetterChainIndex getGetterChainIndex() {
        return getterChainIndex;
    }

    public IndexedFieldConstantsIndex getIndexedFieldConstantsIndex() {
        return indexedFieldConstantsIndex;
    }

    public FieldTypeIndex getFieldTypeIndex() {
        return fieldTypeIndex;
    }

    public PolymorphicMethodExpansionIndex getPolymorphicMethodExpansionIndex() {
        return polymorphicMethodExpansionIndex;
    }

    public ClassCompatibilityCache getClassCompatibilityCache() {
        return classCompatibilityCache;
    }

    public void clearAnalysisCaches() {
        methodDeclarationCache.clear();
        superclassFqnCache.clear();
        classCompatibilityCache.clear();
        subtypeOrSameCache.clear();
    }

    public Set<String> getIndexedTypeFqns() {
        return Collections.unmodifiableSet(classes.keySet());
    }

    public boolean isInlineAccessors() {
        return inlineAccessors;
    }

    public void setInlineAccessors(boolean inlineAccessors) {
        this.inlineAccessors = inlineAccessors;
    }

    public void setProjectRoot(Path root) {
        this.projectRoot = root;
    }

    public String getRelativePath(Path fullPath) {
        if (projectRoot != null && fullPath.isAbsolute() && fullPath.startsWith(projectRoot)) {
            return projectRoot.relativize(fullPath).toString();
        }
        return fullPath.getFileName().toString();
    }

    public String getRelativePath(String fqn) {
        Path path = classPaths.get(fqn);
        return path != null ? getRelativePath(path) : null;
    }

    public void loadLibraryHints(Path hintsFile) throws IOException {
        if (Files.exists(hintsFile)) {
            log.debug("Loading optional hints.json override from {}", hintsFile.toAbsolutePath());
            this.libraryHints = objectMapper.readValue(hintsFile.toFile(), new TypeReference<List<LibraryHint>>() {});
            log.debug("Loaded {} optional library hints from hints.json", libraryHints.size());
        }
    }

    public List<LibraryHint> getLibraryHints() {
        return Collections.unmodifiableList(libraryHints);
    }

    private final Set<String> ignorePatterns = new HashSet<>(Set.of(
            "**/src/test/**",
            "**/src/test/java/**",
            "**/src/testFixtures/**",
            "**/src/test-fixtures/**",
            "**/target/classes/**",
            "**/target/test-classes/**",
            "**/target/generated-test-sources/**",
            "**/target/surefire/**",
            "**/target/failsafe/**",
            "**/target/surefire-reports/**",
            "**/target/failsafe-reports/**",
            "**/build/classes/**",
            "**/build/test-classes/**",
            "**/build/generated-test-sources/**",
            "**/build/generated/sources/**/test/**",
            "**/build/generated/sources/**/testFixtures/**",
            "**/build/libs/**",
            "**/build/tmp/**",
            "**/build/reports/**",
            "**/build/test-results/**",
            "**/build/kotlin/**",
            "**/build/intermediates/**",
            "**/node_modules/**",
            "**/out/**"));
    private String[] classpath = new String[0];
    private String[] sourcepath = new String[0];
    private boolean resolveBindings = false;
    private transient List<String> analysisCallPath;
    private transient Map<String, String> analysisNamedBindings;

    public void setAnalysisCallPath(List<String> path) {
        this.analysisCallPath = path;
    }

    public List<String> getAnalysisCallPath() {
        return analysisCallPath;
    }

    public void clearAnalysisCallPath() {
        this.analysisCallPath = null;
    }

    public void setAnalysisNamedBindings(Map<String, String> bindings) {
        this.analysisNamedBindings = bindings;
    }

    public Map<String, String> getAnalysisNamedBindings() {
        return analysisNamedBindings;
    }

    public void clearAnalysisNamedBindings() {
        this.analysisNamedBindings = null;
    }

    public Map<String, Map<String, String>> getProperties() {
        return Collections.unmodifiableMap(allProperties);
    }

    public void setClasspath(List<String> classpath) {
        this.classpath = classpath.toArray(new String[0]);
    }

    public void setSourcepath(List<String> sourcepath) {
        this.sourcepath = sourcepath.toArray(new String[0]);
    }

    public void setResolveBindings(boolean resolveBindings) {
        this.resolveBindings = resolveBindings;
    }

    public boolean isResolveBindings() {
        return resolveBindings;
    }

    public void setActiveProfiles(List<String> profiles) {
        this.activeProfiles.clear();
        if (profiles != null) {
            this.activeProfiles.addAll(profiles);
        }
        this.propertyResolver.setActiveProfiles(profiles);
    }

    public String resolveString(String input) {
        if (input == null) return null;

        Map<String, String> merged = new HashMap<>();
        // 1. Default properties
        Map<String, String> defaultProps = allProperties.get("default");
        if (defaultProps != null) {
            merged.putAll(defaultProps);
        }

        // 2. Active profiles (later ones override earlier ones)
        for (String profile : activeProfiles) {
            Map<String, String> profileProps = allProperties.get(profile);
            if (profileProps != null) {
                merged.putAll(profileProps);
            }
        }

        return PlaceholderResolver.resolve(input, merged);
    }

    public String resolveExpression(Expression expr) {
        if (expr == null) return null;

        String astResolved = constantResolver.resolve(expr, this);
        if (astResolved == null) {
            astResolved = expr.toString().replaceAll("^\"|\"$", "");
        }

        return resolveString(astResolved);
    }
    public TypeDeclaration resolveStaticImport(String memberName, CompilationUnit contextCu) {
        if (contextCu == null) return null;

        for (Object impObj : contextCu.imports()) {
            ImportDeclaration imp = (ImportDeclaration) impObj;
            if (imp.isStatic()) {
                String impName = imp.getName().getFullyQualifiedName();
                if (imp.isOnDemand()) {
                    // Star import: import static com.example.C.*;
                    TypeDeclaration td = getTypeDeclaration(impName, contextCu);
                    if (td != null && (hasField(td, memberName) || hasMethod(td, memberName))) {
                        return td;
                    }
                } else {
                    // Single import: import static com.example.C.A;
                    if (impName.endsWith("." + memberName)) {
                        String typeFqn = impName.substring(0, impName.lastIndexOf('.'));
                        return getTypeDeclaration(typeFqn, contextCu);
                    }
                }
            }
        }
        return null;
    }

    public boolean hasMethod(TypeDeclaration td, String name) {
        for (MethodDeclaration md : td.getMethods()) {
            if (md.getName().getIdentifier().equals(name)) {
                return true;
            }
        }
        return false;
    }

    public boolean hasField(TypeDeclaration td, String name) {
        for (FieldDeclaration fd : td.getFields()) {
            for (Object fragObj : fd.fragments()) {
                if (fragObj instanceof VariableDeclarationFragment fragment) {
                    if (fragment.getName().getIdentifier().equals(name)) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

    public void scan(Path rootDir) throws IOException {
        scan(Set.of(rootDir), Collections.emptySet());
    }

    private final List<CompilationUnit> allCus = new ArrayList<>();

    private final Map<String, List<String>> interfaceToImpls = new HashMap<>();
    private final Map<String, List<String>> enumValues = new HashMap<>();

    public void scan(Set<Path> rootDirs, Set<String> customIgnorePatterns) throws IOException {
        this.allProperties = propertyResolver.resolveAllProperties(rootDirs);

        Set<String> activeIgnore = new HashSet<>(ignorePatterns);
        activeIgnore.addAll(customIgnorePatterns);

        List<Path> javaFiles = FileUtils.findFilesWithExtension(rootDirs, ".java", activeIgnore);
        log.info("CodebaseContext found {} java files across {} root directories", javaFiles.size(), rootDirs.size());
        typeDeclarations.clear();
        for (Path javaFile : javaFiles) {
            String source = Files.readString(javaFile);
            CompilationUnit cu = parse(source, javaFile.toAbsolutePath().toString());
            allCus.add(cu);
            
            String packageName = cu.getPackage() != null ? cu.getPackage().getName().getFullyQualifiedName() : "";
            for (Object type : cu.types()) {
                if (type instanceof TypeDeclaration td) {
                    indexType(td, packageName, javaFile);
                } else if (type instanceof EnumDeclaration ed) {
                    indexEnum(ed, packageName, javaFile);
                } else if (type instanceof org.eclipse.jdt.core.dom.RecordDeclaration rd) {
                    indexRecord(rd, packageName, javaFile);
                }
            }
        }

        // Pass 2: register inheritance under canonical FQN keys once all types are indexed.
        // Single-pass writes in indexType are unsafe when a subclass is scanned before its superclass.
        indexInheritanceEdges();

        this.implementationsCache.clear();
        this.accessorIndex = inlineAccessors
                ? new AccessorIndexBuilder().build(this)
                : AccessorIndex.empty();
        this.polymorphicAccessorIndex = inlineAccessors
                ? new PolymorphicAccessorIndexBuilder().build(this, accessorIndex)
                : PolymorphicAccessorIndex.empty();
        this.getterChainIndex = inlineAccessors
                ? new GetterChainIndexBuilder().build(this, accessorIndex)
                : GetterChainIndex.empty();
        this.indexedFieldConstantsIndex = inlineAccessors
                ? new IndexedFieldConstantsIndexBuilder().build(this, accessorIndex)
                : IndexedFieldConstantsIndex.empty();
        this.fieldTypeIndex = new FieldTypeIndexBuilder().build(this);
        this.polymorphicMethodExpansionIndex = new PolymorphicMethodExpansionIndexBuilder().build(this);
        clearAnalysisCaches();
        if (inlineAccessors) {
            log.info("Built accessor index: {} trivial accessors across {} types",
                    accessorIndex.trivialCount(), accessorIndex.typeCount());
        } else {
            log.info("Accessor inlining disabled; skipping accessor index build.");
        }
    }

    private void indexType(TypeDeclaration td, String parentFqn, Path javaFile) {
        String simpleName = td.getName().getIdentifier();
        String fqn = parentFqn.isEmpty() ? simpleName : parentFqn + "." + simpleName;

        if (shouldReplaceTypeIndex(fqn, javaFile)) {
            classes.put(fqn, (CompilationUnit) td.getRoot());
            classPaths.put(fqn, javaFile);
            typeDeclarations.put(fqn, td);
            registerSimpleNameForFqn(simpleName, fqn);
        }

        // Recursively index nested types
        for (Object decl : td.bodyDeclarations()) {
            if (decl instanceof TypeDeclaration nestedTd) {
                indexType(nestedTd, fqn, javaFile);
            } else if (decl instanceof EnumDeclaration nestedEd) {
                indexEnum(nestedEd, fqn, javaFile);
            } else if (decl instanceof org.eclipse.jdt.core.dom.RecordDeclaration nestedRd) {
                indexRecord(nestedRd, fqn, javaFile);
            }
        }
    }

    /**
     * After all types are indexed, register extends/implements under a canonical key:
     * resolved FQN when the supertype is in the codebase, otherwise the literal AST name.
     */
    private void indexInheritanceEdges() {
        interfaceToImpls.clear();
        // Snapshot keys: resolveSupertypeIndexKey -> getTypeDeclaration may cache into typeDeclarations.
        for (String implFqn : new ArrayList<>(typeDeclarations.keySet())) {
            TypeDeclaration td = typeDeclarations.get(implFqn);
            if (td == null) {
                continue;
            }
            CompilationUnit cu = td.getRoot() instanceof CompilationUnit root ? root : null;

            if (td.superInterfaceTypes() != null) {
                for (Object itfObj : td.superInterfaceTypes()) {
                    if (itfObj instanceof Type itf) {
                        String key = resolveSupertypeIndexKey(itf, cu);
                        interfaceToImpls.computeIfAbsent(key, k -> new ArrayList<>()).add(implFqn);
                    }
                }
            }
            Type superclass = td.getSuperclassType();
            if (superclass != null) {
                String key = resolveSupertypeIndexKey(superclass, cu);
                interfaceToImpls.computeIfAbsent(key, k -> new ArrayList<>()).add(implFqn);
            }
        }
    }

    /**
     * Canonical index key for a supertype reference: FQN when resolvable in-context, else raw AST name.
     */
    private String resolveSupertypeIndexKey(Type type, CompilationUnit cu) {
        String raw = extractTypeName(type);
        if (raw == null || raw.isBlank()) {
            return raw;
        }
        TypeDeclaration superTd = getTypeDeclaration(raw, cu);
        if (superTd != null) {
            String fqn = getFqn(superTd);
            if (fqn != null && !fqn.isBlank()) {
                return fqn;
            }
        }
        return raw;
    }

    /**
     * Higher rank wins on duplicate FQN indexing: main sources beat build artifacts.
     */
    static int sourcePathRank(Path path) {
        if (path == null) {
            return 0;
        }
        String normalized = path.toString().replace('\\', '/');
        if (normalized.contains("/src/main/java/")) {
            return 2;
        }
        if (normalized.contains("/target/") || normalized.contains("/build/")) {
            return 0;
        }
        return 1;
    }

    private boolean shouldReplaceTypeIndex(String fqn, Path javaFile) {
        Path existingPath = classPaths.get(fqn);
        if (existingPath == null) {
            return true;
        }
        int existingRank = sourcePathRank(existingPath);
        int newRank = sourcePathRank(javaFile);
        if (newRank > existingRank) {
            return true;
        }
        if (newRank < existingRank) {
            if (log.isDebugEnabled()) {
                log.debug(
                        "Skipping lower-rank duplicate type index fqn={} path={} rank={} kept={} rank={}",
                        fqn,
                        javaFile,
                        newRank,
                        existingPath,
                        existingRank);
            }
            return false;
        }
        return false;
    }

    private void registerSimpleNameForFqn(String simpleName, String fqn) {
        if (simpleNameToFqn.containsKey(simpleName)) {
            String existingFqn = simpleNameToFqn.get(simpleName);
            if (existingFqn != null && !existingFqn.equals(fqn)) {
                ambiguousSimpleNames.add(simpleName);
            }
        } else {
            simpleNameToFqn.put(simpleName, fqn);
        }
    }

    /**
     * Walk {@code start} then each superclass; stop on null, missing type, {@code java.lang.Object}, or cycle.
     */
    public void forEachSuperclass(TypeDeclaration start, Consumer<TypeDeclaration> visitor) {
        if (start == null || visitor == null) {
            return;
        }
        Set<String> visited = new LinkedHashSet<>();
        TypeDeclaration current = start;
        while (current != null) {
            String fqn = getFqn(current);
            if (fqn != null && !visited.add(fqn)) {
                break;
            }
            visitor.accept(current);
            String superFqn = getSuperclassFqn(current);
            if (superFqn == null || "java.lang.Object".equals(superFqn)) {
                break;
            }
            current = getTypeDeclaration(superFqn);
        }
    }

    private void indexRecord(org.eclipse.jdt.core.dom.RecordDeclaration rd, String parentFqn, Path javaFile) {
        String simpleName = rd.getName().getIdentifier();
        String fqn = parentFqn.isEmpty() ? simpleName : parentFqn + "." + simpleName;

        if (shouldReplaceTypeIndex(fqn, javaFile)) {
            // Treat as a class since it is a type
            classes.put(fqn, (CompilationUnit) rd.getRoot());
            classPaths.put(fqn, javaFile);
            recordDeclarations.put(fqn, rd);
            registerSimpleNameForFqn(simpleName, fqn);
        }

        // Recursively index nested types
        for (Object decl : rd.bodyDeclarations()) {
            if (decl instanceof TypeDeclaration nestedTd) {
                indexType(nestedTd, fqn, javaFile);
            } else if (decl instanceof EnumDeclaration nestedEd) {
                indexEnum(nestedEd, fqn, javaFile);
            } else if (decl instanceof org.eclipse.jdt.core.dom.RecordDeclaration nestedRd) {
                indexRecord(nestedRd, fqn, javaFile);
            }
        }
    }

    public List<String> getImplementations(String interfaceName) {
        if (interfaceName == null) {
            return List.of();
        }
        String cleanName = interfaceName.contains("<")
                ? interfaceName.substring(0, interfaceName.indexOf('<'))
                : interfaceName;
        return implementationsCache.computeIfAbsent(cleanName, this::computeImplementations);
    }

    private List<String> computeImplementations(String cleanName) {
        Set<String> allImpls = new LinkedHashSet<>();
        collectImplementations(cleanName, allImpls, new HashSet<>());
        List<String> sorted = new ArrayList<>(allImpls);
        Collections.sort(sorted);
        return Collections.unmodifiableList(sorted);
    }

    private void collectImplementations(String typeName, Set<String> results, Set<String> visited) {
        String cleanName = typeName;
        if (typeName != null && typeName.contains("<")) {
            cleanName = typeName.substring(0, typeName.indexOf('<'));
        }
        if (!visited.add(cleanName)) return;

        // Fail closed on ambiguous simple type names: we must not pick an arbitrary FQN and widen
        // to unrelated implementations across packages.
        if (cleanName != null && !cleanName.contains(".") && ambiguousSimpleNames.contains(cleanName)) {
            return;
        }

        List<String> directImpls = resolveDirectImplementations(typeName, cleanName);
        String directImplsSource = describeDirectImplsSource(typeName, cleanName);

        if (log.isDebugEnabled()) {
            log.debug(
                    "collectImplementations typeName={} cleanName={} directImplsSource={} directImpls={}",
                    typeName,
                    cleanName,
                    directImplsSource,
                    directImpls);
        }

        if (directImpls != null) {
            for (String impl : directImpls) {
                results.add(impl);
                // Recursively find implementations of the implementation (e.g., subclasses of an abstract class)
                collectImplementations(impl, results, visited);
            }
        }
    }

    /**
     * Union direct subclasses/implementors from every index bucket that refers to the same supertype.
     * Generated stubs often index under a literal FQN ({@code extends abcd.ccc.AML}) while hand-written
     * sources index under the simple name ({@code extends AML}) — merge instead of stopping at the first hit.
     */
    private List<String> resolveDirectImplementations(String typeName, String cleanName) {
        LinkedHashSet<String> merged = new LinkedHashSet<>();
        if (cleanName == null || cleanName.isBlank()) {
            return null;
        }

        addImplsForKey(merged, cleanName);

        if (!cleanName.contains(".")) {
            String fqn = simpleNameToFqn.get(typeName != null ? typeName : cleanName);
            if (fqn != null && !fqn.equals(cleanName)) {
                addImplsForKey(merged, fqn);
            }
        } else {
            String simpleName = cleanName.substring(cleanName.lastIndexOf('.') + 1);
            if (!ambiguousSimpleNames.contains(simpleName)) {
                addImplsForKey(merged, simpleName);
            } else {
                String pkg = cleanName.substring(0, cleanName.lastIndexOf('.'));
                List<String> simpleImpls = interfaceToImpls.get(simpleName);
                if (simpleImpls != null) {
                    for (String impl : simpleImpls) {
                        if (impl.startsWith(pkg + ".")) {
                            merged.add(impl);
                        }
                    }
                }
            }
        }

        return merged.isEmpty() ? null : new ArrayList<>(merged);
    }

    private void addImplsForKey(Set<String> merged, String key) {
        List<String> impls = interfaceToImpls.get(key);
        if (impls != null) {
            merged.addAll(impls);
        }
    }

    private String describeDirectImplsSource(String typeName, String cleanName) {
        if (cleanName == null || cleanName.isBlank()) {
            return "none";
        }
        boolean hasDirect = interfaceToImpls.containsKey(cleanName);
        boolean hasFqnAlias = false;
        boolean hasSimpleAlias = false;
        if (!cleanName.contains(".")) {
            String fqn = simpleNameToFqn.get(typeName != null ? typeName : cleanName);
            hasFqnAlias = fqn != null && interfaceToImpls.containsKey(fqn);
        } else {
            String simpleName = cleanName.substring(cleanName.lastIndexOf('.') + 1);
            hasSimpleAlias = interfaceToImpls.containsKey(simpleName);
        }
        if (hasDirect && (hasFqnAlias || hasSimpleAlias)) {
            return "merged";
        }
        if (hasDirect) {
            return "direct";
        }
        if (hasFqnAlias) {
            return "simple→fqn";
        }
        if (hasSimpleAlias) {
            return "fqn→simple";
        }
        return "none";
    }

    private void indexEnum(EnumDeclaration ed, String parentFqn, Path javaFile) {
        String simpleName = ed.getName().getIdentifier();
        String fqn = parentFqn.isEmpty() ? simpleName : parentFqn + "." + simpleName;

        if (shouldReplaceTypeIndex(fqn, javaFile)) {
            classes.put(fqn, (CompilationUnit) ed.getRoot());
            classPaths.put(fqn, javaFile);

            List<String> values = new ArrayList<>();
            for (Object o : ed.enumConstants()) {
                if (o instanceof EnumConstantDeclaration ecd) {
                    values.add(fqn + "." + ecd.getName().getIdentifier());
                }
            }
            enumValues.put(fqn, values);
            registerSimpleNameForFqn(simpleName, fqn);
        }

        // Recursively index nested types inside enum body
        for (Object decl : ed.bodyDeclarations()) {
            if (decl instanceof TypeDeclaration nestedTd) {
                indexType(nestedTd, fqn, javaFile);
            } else if (decl instanceof EnumDeclaration nestedEd) {
                indexEnum(nestedEd, fqn, javaFile);
            } else if (decl instanceof org.eclipse.jdt.core.dom.RecordDeclaration nestedRd) {
                indexRecord(nestedRd, fqn, javaFile);
            }
        }
    }

    public Map<String, List<String>> getEnumValuesMap() {
        return enumValues;
    }

    /**
     * Returns true when multiple types share the same simple name across packages, so simple-name
     * lookup must fail closed.
     */
    public boolean isAmbiguousSimpleName(String name) {
        if (name == null || name.isBlank() || name.contains(".")) {
            return false;
        }
        return ambiguousSimpleNames.contains(name);
    }

    /**
     * True when two type references denote the same type, or share an unambiguous simple name.
     * Fails closed when the shared simple name exists in multiple packages.
     */
    public boolean areSameTypeOrUnambiguousSimpleMatch(String leftFqn, String rightFqn) {
        if (leftFqn == null || rightFqn == null) {
            return false;
        }
        if (leftFqn.equals(rightFqn)) {
            return true;
        }
        String leftSimple = simpleTypeName(leftFqn);
        String rightSimple = simpleTypeName(rightFqn);
        if (!leftSimple.equals(rightSimple)) {
            return false;
        }
        return !isAmbiguousSimpleName(leftSimple);
    }

    public List<String> getEnumValues(String fqnOrSimpleName) {
        if (fqnOrSimpleName == null) return null;
        
        String cleanName = fqnOrSimpleName;
        if (cleanName.contains("<")) {
            cleanName = cleanName.substring(0, cleanName.indexOf('<'));
        }
        if (cleanName.contains("[")) {
            cleanName = cleanName.substring(0, cleanName.indexOf('['));
        }
        cleanName = cleanName.trim();

        List<String> values = enumValues.get(cleanName);
        if (values == null) {
            if (!cleanName.contains(".") && ambiguousSimpleNames.contains(cleanName)) {
                return null;
            }
            String fqn = simpleNameToFqn.get(cleanName);
            if (fqn != null) values = enumValues.get(fqn);
        }
        log.debug("getEnumValues called for: {} resolved to {}. Returns: {}", fqnOrSimpleName, cleanName, values);
        return values;
    }

    public State resolveState(Expression expr, CompilationUnit cu) {
        return stateResolver.resolveState(expr, cu);
    }

    private CompilationUnit parse(String source, String unitName) {
        ASTParser parser = ASTParser.newParser(AST.getJLSLatest());
        parser.setSource(source.toCharArray());
        parser.setKind(ASTParser.K_COMPILATION_UNIT);
        parser.setResolveBindings(resolveBindings);
        parser.setBindingsRecovery(true);

        Map<String, String> options = JavaCore.getOptions();
        JavaCore.setComplianceOptions(JavaCore.VERSION_21, options);
        parser.setCompilerOptions(options);

        if (resolveBindings) {
            parser.setUnitName(unitName);
            parser.setEnvironment(classpath, sourcepath, null, true);
        }
        return (CompilationUnit) parser.createAST(null);
    }

    public TypeDeclaration getTypeDeclaration(String name) {
        TypeDeclaration cached = typeDeclarations.get(name);
        if (cached != null) {
            return cached;
        }

        // Try exact FQN match first
        CompilationUnit cu = classes.get(name);
        if (cu != null) {
            TypeDeclaration resolved = findTypeInCu(cu, name);
            if (resolved != null) {
                typeDeclarations.put(name, resolved);
            }
            return resolved;
        }

        // If it's a simple name, check if it's ambiguous
        if (ambiguousSimpleNames.contains(name)) {
            // System.err.println("Warning: Ambiguous simple name '" + name + "'. Provide FQN to resolve correctly.");
            return null; // Don't return a random class if it's ambiguous
        }

        // If not ambiguous, return the one we found during scan
        String fqn = simpleNameToFqn.get(name);
        if (fqn != null) {
            cached = typeDeclarations.get(fqn);
            if (cached != null) {
                return cached;
            }
            cu = classes.get(fqn);
            if (cu != null) {
                TypeDeclaration resolved = findTypeInCu(cu, fqn);
                if (resolved != null) {
                    typeDeclarations.put(fqn, resolved);
                }
                return resolved;
            }
        }

        return null;
    }

    public TypeDeclaration getTypeDeclaration(String name, CompilationUnit contextCu) {
        if (name == null || name.isEmpty())
            return null;

        // 1. Check if it's already an FQN
        if (classes.containsKey(name))
            return getTypeDeclaration(name);

        // 2. Check imports in contextCu
        if (contextCu != null) {
            for (Object typeObj : contextCu.types()) {
                if (typeObj instanceof TypeDeclaration tdDecl) {
                    if (tdDecl.getName().getIdentifier().equals(name)) {
                        return tdDecl;
                    }
                }
            }

            for (Object impObj : contextCu.imports()) {
                ImportDeclaration imp = (ImportDeclaration) impObj;
                String impName = imp.getName().getFullyQualifiedName();
                if (!imp.isStatic() && !imp.isOnDemand()) {
                    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);

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

            // 5. Fallback to java.lang (common)
            String langFqn = "java.lang." + name;
            if (classes.containsKey(langFqn))
                return getTypeDeclaration(langFqn);
        }

        // 6. Last resort: global simple name match (might be ambiguous but better than nothing)
        return getTypeDeclaration(name);
    }

    public String getSuperclassFqn(TypeDeclaration td) {
        if (td == null) {
            return null;
        }
        String typeFqn = getFqn(td);
        if (typeFqn != null && superclassFqnCache.containsKey(typeFqn)) {
            return superclassFqnCache.get(typeFqn);
        }
        String resolved = resolveSuperclassFqn(td);
        if (typeFqn != null) {
            superclassFqnCache.put(typeFqn, resolved);
        }
        return resolved;
    }

    /**
     * AST-backed subtype check: {@code true} when {@code childFqn} is the same type as
     * {@code parentFqn}, or extends/implements it (transitively). Does not require JDT bindings.
     */
    public boolean isSubtypeOrSame(String childFqn, String parentFqn) {
        if (childFqn == null || parentFqn == null || childFqn.isBlank() || parentFqn.isBlank()) {
            return false;
        }
        if (childFqn.equals(parentFqn) || areSameTypeOrUnambiguousSimpleMatch(childFqn, parentFqn)) {
            return true;
        }
        String cacheKey = childFqn + "<:" + parentFqn;
        Boolean cached = subtypeOrSameCache.get(cacheKey);
        if (cached != null) {
            return cached;
        }
        boolean result = computeIsSubtypeOrSame(childFqn, parentFqn, new HashSet<>());
        subtypeOrSameCache.put(cacheKey, result);
        return result;
    }

    /**
     * {@code true} when the type is a known non-abstract class (not an interface).
     * Unknown types are treated as non-concrete so pathfinding fails open.
     */
    public boolean isConcreteClassType(String typeFqn) {
        if (typeFqn == null || typeFqn.isBlank()) {
            return false;
        }
        TypeDeclaration td = getTypeDeclaration(typeFqn);
        if (td == null) {
            return false;
        }
        return !td.isInterface() && !Modifier.isAbstract(td.getModifiers());
    }

    private boolean computeIsSubtypeOrSame(String childFqn, String parentFqn, Set<String> visited) {
        if (childFqn == null || !visited.add(childFqn)) {
            return false;
        }
        if (childFqn.equals(parentFqn) || areSameTypeOrUnambiguousSimpleMatch(childFqn, parentFqn)) {
            return true;
        }
        TypeDeclaration childTd = getTypeDeclaration(childFqn);
        if (childTd == null) {
            return false;
        }
        String superclass = getSuperclassFqn(childTd);
        if (superclass != null && computeIsSubtypeOrSame(superclass, parentFqn, visited)) {
            return true;
        }
        CompilationUnit cu = childTd.getRoot() instanceof CompilationUnit root ? root : null;
        for (Object interfaceTypeObj : childTd.superInterfaceTypes()) {
            if (!(interfaceTypeObj instanceof Type interfaceType)) {
                continue;
            }
            String interfaceName = extractTypeName(interfaceType);
            TypeDeclaration interfaceTd = getTypeDeclaration(interfaceName, cu);
            String interfaceFqn = interfaceTd != null ? getFqn(interfaceTd) : interfaceName;
            if (interfaceFqn != null && computeIsSubtypeOrSame(interfaceFqn, parentFqn, visited)) {
                return true;
            }
        }
        return false;
    }

    private String resolveSuperclassFqn(TypeDeclaration td) {
        String bindingSuper = null;
        ITypeBinding binding = td.resolveBinding();
        if (binding != null) {
            ITypeBinding superBinding = binding.getSuperclass();
            if (superBinding != null) {
                bindingSuper = superBinding.getErasure().getQualifiedName();
            }
        }

        Type superType = td.getSuperclassType();
        if (superType == null) {
            return bindingSuper;
        }

        String superName = extractTypeName(superType);
        CompilationUnit cu = (td.getRoot() instanceof CompilationUnit) ? (CompilationUnit) td.getRoot() : null;
        TypeDeclaration superTd = getTypeDeclaration(superName, cu);
        String astSuper = superTd != null ? getFqn(superTd) : null;

        if (astSuper != null) {
            if (bindingSuper == null || !astSuper.equals(bindingSuper)) {
                if (bindingSuper == null
                        || getTypeDeclaration(bindingSuper) == null
                        || classes.containsKey(astSuper)) {
                    return astSuper;
                }
            }
            return astSuper;
        }
        return bindingSuper != null ? bindingSuper : superName;
    }

    private String extractTypeName(Type type) {
        if (type.isParameterizedType()) {
            return extractTypeName(((ParameterizedType) type).getType());
        } else if (type.isSimpleType()) {
            return ((SimpleType) type).getName().getFullyQualifiedName();
        } else if (type.isQualifiedType()) {
            return ((QualifiedType) type).getName().getFullyQualifiedName();
        }
        return type.toString();
    }

    public MethodDeclaration findMethodDeclaration(TypeDeclaration td, String methodName, boolean includeSuper) {
        if (td == null) {
            return null;
        }
        String typeFqn = getFqn(td);
        if (typeFqn != null) {
            String cacheKey = typeFqn + "#" + methodName + "#" + includeSuper;
            if (methodDeclarationCache.containsKey(cacheKey)) {
                return methodDeclarationCache.get(cacheKey);
            }
            MethodDeclaration resolved = findMethodDeclarationUncached(td, methodName, includeSuper);
            methodDeclarationCache.put(cacheKey, resolved);
            return resolved;
        }
        return findMethodDeclarationUncached(td, methodName, includeSuper);
    }

    private MethodDeclaration findMethodDeclarationUncached(TypeDeclaration td, String methodName, boolean includeSuper) {
        ITypeBinding binding = td.resolveBinding();
        if (binding != null && includeSuper) {
            IMethodBinding methodBinding = findMethodBinding(binding, methodName);
            if (methodBinding != null && methodBinding.getDeclaringClass() != null) {
                CompilationUnit declaringCu = classes.get(methodBinding.getDeclaringClass().getErasure().getQualifiedName());
                if (declaringCu != null) {
                    ASTNode declNode = declaringCu.findDeclaringNode(methodBinding.getKey());
                    if (declNode instanceof MethodDeclaration) {
                        return (MethodDeclaration) declNode;
                    }
                }
            }
        }
        return findMethodDeclarationLegacy(td, methodName, includeSuper);
    }

    private IMethodBinding findMethodBinding(ITypeBinding typeBinding, String methodName) {
        return findMethodBinding(typeBinding, methodName, new java.util.HashSet<>());
    }

    private IMethodBinding findMethodBinding(ITypeBinding typeBinding, String methodName, Set<String> visited) {
        if (typeBinding == null) return null;
        String key = typeBinding.getKey();
        if (key != null && !visited.add(key)) {
            return null;
        }

        for (IMethodBinding mb : typeBinding.getDeclaredMethods()) {
            if (mb.getName().equals(methodName)) {
                return mb;
            }
        }
        
        ITypeBinding superclass = typeBinding.getSuperclass();
        if (superclass != null) {
            IMethodBinding mb = findMethodBinding(superclass, methodName, visited);
            if (mb != null) return mb;
        }

        for (ITypeBinding intf : typeBinding.getInterfaces()) {
            IMethodBinding mb = findMethodBinding(intf, methodName, visited);
            if (mb != null) return mb;
        }

        return null;
    }

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

    private MethodDeclaration findMethodDeclarationLegacy(
            TypeDeclaration td,
            String methodName,
            boolean includeSuper,
            Set<String> visited) {
        if (td == null) {
            return null;
        }
        String typeFqn = getFqn(td);
        if (typeFqn != null && !visited.add(typeFqn)) {
            return null;
        }
        for (MethodDeclaration method : td.getMethods()) {
            if (method.getName().getIdentifier().equals(methodName)) {
                return method;
            }
        }

        if (includeSuper) {
            String superFqn = getSuperclassFqn(td);
            if (superFqn != null) {
                TypeDeclaration superTd = getTypeDeclaration(superFqn);
                if (superTd != null) {
                    MethodDeclaration found = findMethodDeclarationLegacy(superTd, methodName, true, visited);
                    if (found != null) {
                        return found;
                    }
                }
            }

            // Check interfaces
            for (Object interfaceTypeObj : td.superInterfaceTypes()) {
                Type interfaceType = (Type) interfaceTypeObj;
                String interfaceName = extractTypeName(interfaceType);
                TypeDeclaration interfaceTd = getTypeDeclaration(
                        interfaceName,
                        (td.getRoot() instanceof CompilationUnit) ? (CompilationUnit) td.getRoot() : null);
                if (interfaceTd != null) {
                    MethodDeclaration found = findMethodDeclarationLegacy(interfaceTd, methodName, true, visited);
                    if (found != null) {
                        return found;
                    }
                }
            }
        }
        return null;
    }

    public boolean areClassesPolymorphicallyCompatible(String classNeighbor, String classTarget) {
        if (classNeighbor == null || classTarget == null) {
            return false;
        }
        if (classNeighbor.equals(classTarget)) {
            return true;
        }
        if (areSameTypeOrUnambiguousSimpleMatch(classNeighbor, classTarget)) {
            return true;
        }
        return classCompatibilityCache.areClassesCompatible(
                classNeighbor,
                classTarget,
                this::computePolymorphicClassCompatibility);
    }

    private boolean computePolymorphicClassCompatibility(String classNeighbor, String classTarget) {
        List<String> targetImplementations = getImplementations(classTarget);
        for (String implementation : targetImplementations) {
            if (implementation.equals(classNeighbor) || implementation.endsWith("." + classNeighbor)) {
                return true;
            }
        }
        List<String> neighborImplementations = getImplementations(classNeighbor);
        for (String implementation : neighborImplementations) {
            if (implementation.equals(classTarget) || implementation.endsWith("." + classTarget)) {
                return true;
            }
        }
        return false;
    }

    private static String simpleTypeName(String typeFqn) {
        return typeFqn.contains(".") ? typeFqn.substring(typeFqn.lastIndexOf('.') + 1) : typeFqn;
    }

    public org.eclipse.jdt.core.dom.RecordDeclaration getRecordDeclaration(String fqn) {
        if (fqn == null || fqn.isEmpty()) {
            return null;
        }
        org.eclipse.jdt.core.dom.RecordDeclaration indexed = recordDeclarations.get(fqn);
        if (indexed != null) {
            return indexed;
        }
        AbstractTypeDeclaration abstractType = getAbstractTypeDeclaration(fqn);
        if (abstractType instanceof org.eclipse.jdt.core.dom.RecordDeclaration recordDeclaration) {
            return recordDeclaration;
        }
        return null;
    }

    public AbstractTypeDeclaration getAbstractTypeDeclaration(String name) {
        // Try exact FQN match first
        CompilationUnit cu = classes.get(name);
        if (cu != null)
            return findAbstractTypeInCu(cu, name);

        // If it's a simple name, check if it's ambiguous
        if (ambiguousSimpleNames.contains(name)) {
            return null;
        }

        // If not ambiguous, return the one we found during scan
        String fqn = simpleNameToFqn.get(name);
        if (fqn != null) {
            cu = classes.get(fqn);
            if (cu != null)
                return findAbstractTypeInCu(cu, fqn);
        }

        return null;
    }

    public AbstractTypeDeclaration getAbstractTypeDeclaration(String name, CompilationUnit contextCu) {
        if (name == null || name.isEmpty())
            return null;

        // 1. Check if it's already an FQN
        if (classes.containsKey(name))
            return getAbstractTypeDeclaration(name);

        // 2. Check imports in contextCu
        if (contextCu != null) {
            for (Object impObj : contextCu.imports()) {
                ImportDeclaration imp = (ImportDeclaration) impObj;
                String impName = imp.getName().getFullyQualifiedName();
                if (!imp.isStatic() && !imp.isOnDemand()) {
                    if (impName.endsWith("." + name))
                        return getAbstractTypeDeclaration(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 getAbstractTypeDeclaration(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 getAbstractTypeDeclaration(starFqn);
                }
            }

            // 5. Fallback to java.lang (common)
            String langFqn = "java.lang." + name;
            if (classes.containsKey(langFqn))
                return getAbstractTypeDeclaration(langFqn);
        }

        // 6. Last resort: global simple name match
        return getAbstractTypeDeclaration(name);
    }

    private AbstractTypeDeclaration findAbstractTypeInCu(CompilationUnit cu, String fqn) {
        for (Object type : cu.types()) {
            if (type instanceof AbstractTypeDeclaration atd) {
                AbstractTypeDeclaration found = findNestedAbstractType(atd, fqn);
                if (found != null) return found;
            }
        }
        return null;
    }

    private AbstractTypeDeclaration findNestedAbstractType(AbstractTypeDeclaration atd, String targetFqn) {
        if (getFqn(atd).equals(targetFqn)) return atd;
        if (atd instanceof TypeDeclaration td) {
            for (TypeDeclaration nested : td.getTypes()) {
                AbstractTypeDeclaration found = findNestedAbstractType(nested, targetFqn);
                if (found != null) return found;
            }
        }
        return null;
    }

    public String getFqn(AbstractTypeDeclaration atd) {
        if (atd == null) return null;
        StringBuilder sb = new StringBuilder(atd.getName().getIdentifier());
        ASTNode current = atd.getParent();
        while (current instanceof AbstractTypeDeclaration parent) {
            sb.insert(0, parent.getName().getIdentifier() + ".");
            current = parent.getParent();
        }

        CompilationUnit cu = (atd.getRoot() instanceof CompilationUnit) ? (CompilationUnit) atd.getRoot() : null;
        if (cu == null) {
            return sb.toString();
        }
        String packageName = cu.getPackage() != null ? cu.getPackage().getName().getFullyQualifiedName() : "";
        String fqn = sb.toString();
        return packageName.isEmpty() ? fqn : packageName + "." + fqn;
    }

    private TypeDeclaration findTypeInCu(CompilationUnit cu, String fqn) {
        for (Object type : cu.types()) {
            if (type instanceof TypeDeclaration td) {
                TypeDeclaration found = findNestedType(td, fqn);
                if (found != null) return found;
            }
        }
        return null;
    }

    private TypeDeclaration findNestedType(TypeDeclaration td, String targetFqn) {
        if (getFqn(td).equals(targetFqn)) return td;
        for (TypeDeclaration nested : td.getTypes()) {
            TypeDeclaration found = findNestedType(nested, targetFqn);
            if (found != null) return found;
        }
        return null;
    }

    public Path getPath(String className) {
        return classPaths.get(className);
    }

    public List<TypeDeclaration> findEntryPointClasses(List<String> targetAnnotations) {
        List<TypeDeclaration> entryPoints = new ArrayList<>();
        for (CompilationUnit cu : classes.values()) {
            for (Object type : cu.types()) {
                if (type instanceof TypeDeclaration td) {
                    if (isEntryPoint(td, targetAnnotations)) {
                        entryPoints.add(td);
                    }
                }
            }
        }
        return entryPoints;
    }

    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) {
                String annotationName = annotation.getTypeName().getFullyQualifiedName();
                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;
            }
        }

        // If it's a leaf class (not abstract) and extends StateMachineConfigurerAdapter indirectly
        if (!isAbstract && extendsStateMachineConfigurerAdapter(td)) {
            return true;
        }

        return false;
    }

    public boolean extendsStateMachineConfigurerAdapter(TypeDeclaration td) {
        ITypeBinding binding = td.resolveBinding();
        if (binding != null) {
            boolean isSubtype = BindingResolver.isSubtypeOf(binding, "org.springframework.statemachine.config.builders.StateMachineConfigurer")
                || BindingResolver.isSubtypeOf(binding, "org.springframework.statemachine.config.StateMachineConfigurerAdapter")
                || BindingResolver.isSubtypeOf(binding, "org.springframework.statemachine.config.EnumStateMachineConfigurerAdapter");
            if (isSubtype) {
                return true;
            }
        }
        return extendsStateMachineConfigurerAdapterLegacy(td, new HashSet<>());
    }

    private boolean extendsStateMachineConfigurerAdapterLegacy(TypeDeclaration td, Set<String> visited) {
        String fqn = getFqn(td);
        if (visited.contains(fqn)) {
            return false;
        }
        visited.add(fqn);

        Set<String> knownAdapters = Set.of("EnumStateMachineConfigurerAdapter", "StateMachineConfigurerAdapter", "StateMachineConfigurer");
        CompilationUnit cu = (CompilationUnit) td.getRoot();

        // Check superclass
        Type superclass = td.getSuperclassType();
        if (superclass != null) {
            String superclassName = getSimpleName(superclass);
            if (knownAdapters.contains(superclassName)
                    || superclass.toString().contains("StateMachineConfigurerAdapter")
                    || superclass.toString().contains("EnumStateMachineConfigurerAdapter")) {
                return true;
            }
            TypeDeclaration superTd = getTypeDeclaration(superclassName, cu);
            if (superTd != null && extendsStateMachineConfigurerAdapterLegacy(superTd, visited))
                return true;
        }

        // Check interfaces
        for (Object interfaceTypeObj : td.superInterfaceTypes()) {
            if (interfaceTypeObj instanceof Type interfaceType) {
                String interfaceName = getSimpleName(interfaceType);
                if (knownAdapters.contains(interfaceName))
                    return true;
                TypeDeclaration interfaceTd = getTypeDeclaration(interfaceName, cu);
                if (interfaceTd != null && extendsStateMachineConfigurerAdapterLegacy(interfaceTd, visited))
                    return true;
            }
        }

        return false;
    }

    public List<MethodDeclaration> findBeanMethodsReturning(Set<String> targetReturnTypes) {
        List<MethodDeclaration> beanMethods = new ArrayList<>();
        for (CompilationUnit cu : classes.values()) {
            for (Object type : cu.types()) {
                if (type instanceof TypeDeclaration td) {
                    for (MethodDeclaration method : td.getMethods()) {
                        if (isBeanMethodReturning(method, targetReturnTypes)) {
                            beanMethods.add(method);
                        }
                    }
                }
            }
        }
        return beanMethods;
    }

    private boolean isBeanMethodReturning(MethodDeclaration method, Set<String> targetReturnTypes) {
        // Check return type
        Type returnType = method.getReturnType2();
        if (returnType == null)
            return false;
        String typeName = AstUtils.extractSimpleTypeName(returnType);
        if (!targetReturnTypes.contains(typeName))
            return false;

        // Check for @Bean annotation
        for (Object modifier : method.modifiers()) {
            if (modifier instanceof Annotation annotation) {
                String annotationName = annotation.getTypeName().getFullyQualifiedName();
                if ("Bean".equals(annotationName) || "org.springframework.context.annotation.Bean".equals(annotationName)) {
                    return true;
                }
            }
        }
        return false;
    }

    private String getSimpleName(Type type) {
        return AstUtils.extractSimpleTypeName(type);
    }

    public Collection<CompilationUnit> getCompilationUnits() {
        return Collections.unmodifiableCollection(allCus);
    }

    public Collection<TypeDeclaration> getTypeDeclarations() {
        List<TypeDeclaration> types = new ArrayList<>();
        for (CompilationUnit cu : allCus) {
            for (Object typeObj : cu.types()) {
                if (typeObj instanceof TypeDeclaration td) {
                    types.add(td);
                    // Handle nested types
                    types.addAll(getNestedTypes(td));
                }
            }
        }
        return types;
    }

    private List<TypeDeclaration> getNestedTypes(TypeDeclaration td) {
        List<TypeDeclaration> nested = new ArrayList<>();
        for (TypeDeclaration inner : td.getTypes()) {
            nested.add(inner);
            nested.addAll(getNestedTypes(inner));
        }
        return nested;
    }
}
