package click.kamil.springstatemachineexporter.analysis.service;

import click.kamil.springstatemachineexporter.analysis.model.CallEdge;
import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.*;
import java.util.*;

public class VariableTracer {
    private final CodebaseContext context;
    private final ConstantResolver constantResolver;
    private ConstantExtractor constantExtractor;

    private final click.kamil.springstatemachineexporter.ast.common.DataFlowModel dataFlowModel;

    public VariableTracer(CodebaseContext context, ConstantResolver constantResolver) {
        this.context = context;
        this.constantResolver = constantResolver;
        this.dataFlowModel = new click.kamil.springstatemachineexporter.ast.common.JdtDataFlowModel(context);
    }

    public void setConstantExtractor(ConstantExtractor constantExtractor) {
        this.constantExtractor = constantExtractor;
    }

    public Expression traceVariable(Expression expr) {
        List<Expression> all = traceVariableAll(expr);
        return all.isEmpty() ? expr : all.get(all.size() - 1);
    }

    private Expression peelExpression(Expression expr) {
        while (expr instanceof ParenthesizedExpression pe) {
            expr = pe.getExpression();
        }
        if (expr instanceof CastExpression ce) {
            return peelExpression(ce.getExpression());
        }
        return expr;
    }

    public List<Expression> traceVariableAll(Expression expr) {
        return dataFlowModel.getReachingDefinitions(expr);
    }

    public List<Expression> traceVariableAll(Expression expr, Set<String> visitedVariables) {
        return dataFlowModel.getReachingDefinitions(expr);
    }

    public Expression traceLocalSetter(String methodFqn, String varName, String getterName) {
        if (methodFqn == null || !methodFqn.contains(".")) return null;
        String className = methodFqn.substring(0, methodFqn.lastIndexOf('.'));
        String methodName = methodFqn.substring(methodFqn.lastIndexOf('.') + 1);
        TypeDeclaration td = context.getTypeDeclaration(className);
        if (td != null) {
            MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
            if (md != null && md.getBody() != null) {
                final Expression[] setterArg = new Expression[1];
                String propName = getterName.startsWith("get") ? getterName.substring(3) : getterName;
                md.getBody().accept(new ASTVisitor() {
                    @Override
                    public boolean visit(MethodInvocation node) {
                        if (node.getName().getIdentifier().equalsIgnoreCase("set" + propName) || node.getName().getIdentifier().equalsIgnoreCase(propName)) {
                            if (node.getExpression() instanceof SimpleName sn && sn.getIdentifier().equals(varName)) {
                                if (!node.arguments().isEmpty()) {
                                    setterArg[0] = (Expression) node.arguments().get(0);
                                }
                            }
                        }
                        return super.visit(node);
                    }
                    @Override
                    public boolean visit(Assignment node) {
                        if (node.getLeftHandSide() instanceof FieldAccess fa) {
                            if (fa.getExpression() instanceof SimpleName faSn && faSn.getIdentifier().equals(varName)) {
                                if (fa.getName().getIdentifier().equalsIgnoreCase(propName)) {
                                    setterArg[0] = node.getRightHandSide();
                                }
                            }
                        } else if (node.getLeftHandSide() instanceof QualifiedName qqn) {
                            if (qqn.getQualifier().getFullyQualifiedName().equals(varName)) {
                                if (qqn.getName().getIdentifier().equalsIgnoreCase(propName)) {
                                    setterArg[0] = node.getRightHandSide();
                                }
                            }
                        }
                        return super.visit(node);
                    }
                });
                return setterArg[0];
            }
        }
        return null;
    }

    private String getFieldType(TypeDeclaration td, String fieldName, CodebaseContext context, java.util.Set<String> visited) {
        if (td == null) return null;
        String typeFqn = context.getFqn(td);
        if (!visited.add(typeFqn)) return null;

        for (FieldDeclaration fd : td.getFields()) {
            for (Object fragObj : fd.fragments()) {
                VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
                if (frag.getName().getIdentifier().equals(fieldName)) {
                    return fd.getType().toString();
                }
            }
        }
        
        String superFqn = context.getSuperclassFqn(td);
        if (superFqn != null) {
            TypeDeclaration superTd = context.getTypeDeclaration(superFqn);
            if (superTd != null) {
                String result = getFieldType(superTd, fieldName, context, visited);
                if (result != null) return result;
            }
        }
        
        for (Object intfObj : td.superInterfaceTypes()) {
            Type intfType = (Type) intfObj;
            String intfName = click.kamil.springstatemachineexporter.ast.common.AstUtils.extractSimpleTypeName(intfType);
            String resolvedIntf = resolveTypeFqn(intfName, context, td);
            if (resolvedIntf != null) {
                TypeDeclaration intfTd = context.getTypeDeclaration(resolvedIntf);
                if (intfTd != null) {
                    String result = getFieldType(intfTd, fieldName, context, visited);
                    if (result != null) return result;
                }
            }
        }
        
        return null;
    }
    
    private String resolveTypeFqn(String simpleName, CodebaseContext context, TypeDeclaration td) {
        CompilationUnit cu = td.getRoot() instanceof CompilationUnit ? (CompilationUnit) td.getRoot() : null;
        if (cu != null) {
            for (Object impObj : cu.imports()) {
                ImportDeclaration imp = (ImportDeclaration) impObj;
                String name = imp.getName().getFullyQualifiedName();
                if (name.endsWith("." + simpleName)) {
                    return name;
                }
            }
            if (cu.getPackage() != null) {
                return cu.getPackage().getName().getFullyQualifiedName() + "." + simpleName;
            }
        }
        return simpleName;
    }

    public String getVariableDeclaredType(String methodFqn, String cleanFieldName) {
        if (methodFqn == null) return null;
        int lastDot = methodFqn.lastIndexOf('.');
        if (lastDot > 0) {
            String fqn = methodFqn.substring(0, lastDot);
            String methodName = methodFqn.substring(lastDot + 1);
            TypeDeclaration td = context.getTypeDeclaration(fqn);
            if (td != null) {
                MethodDeclaration md = context.findMethodDeclaration(td, methodName, false);
                if (md != null) {
                    for (Object pObj : md.parameters()) {
                        SingleVariableDeclaration svd = (SingleVariableDeclaration) pObj;
                        if (svd.getName().getIdentifier().equals(cleanFieldName)) {
                            return svd.getType().toString();
                        }
                    }
                    final String[] foundType = new String[1];
                    if (md.getBody() != null) {
                        md.getBody().accept(new ASTVisitor() {
                        @Override
                        public boolean visit(VariableDeclarationStatement node) {
                            for (Object fragObj : node.fragments()) {
                                VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
                                if (frag.getName().getIdentifier().equals(cleanFieldName)) {
                                    foundType[0] = node.getType().toString();
                                }
                            }
                            return super.visit(node);
                        }
                        
                        @Override
                        public boolean visit(LambdaExpression node) {
                            for (Object paramObj : node.parameters()) {
                                VariableDeclaration param = (VariableDeclaration) paramObj;
                                if (param.getName().getIdentifier().equals(cleanFieldName)) {
                                    if (param instanceof SingleVariableDeclaration svd && svd.getType() != null) {
                                        foundType[0] = svd.getType().toString();
                                        return false;
                                    }
                                    
                                    ASTNode parent = node.getParent();
                                    if (parent instanceof MethodInvocation mi) {
                                        Expression caller = mi.getExpression();
                                        
                                        boolean hasMap = false;
                                        while (caller instanceof MethodInvocation chainMi) {
                                            String mName = chainMi.getName().getIdentifier();
                                            if (mName.equals("map") || mName.equals("flatMap")) {
                                                hasMap = true;
                                                if (chainMi.arguments().size() == 1 && chainMi.arguments().get(0) instanceof LambdaExpression lambda) {
                                                    if (lambda.getBody() instanceof MethodInvocation mapMi) {
                                                        Expression mapCaller = mapMi.getExpression();
                                                        if (mapCaller instanceof SimpleName mapSn) {
                                                            String mapCallerType = getVariableDeclaredType(methodFqn, mapSn.getIdentifier());
                                                            if (mapCallerType != null) {
                                                                TypeDeclaration currentTd = context.getTypeDeclaration(methodFqn.substring(0, methodFqn.lastIndexOf('.')));
                                                                CompilationUnit cu = currentTd != null && currentTd.getRoot() instanceof CompilationUnit ? (CompilationUnit) currentTd.getRoot() : null;
                                                                TypeDeclaration td = context.getTypeDeclaration(mapCallerType, cu);
                                                                if (td != null) {
                                                                    MethodDeclaration mapMd = context.findMethodDeclaration(td, mapMi.getName().getIdentifier(), true);
                                                                    if (mapMd != null && mapMd.getReturnType2() != null) {
                                                                        foundType[0] = mapMd.getReturnType2().toString();
                                                                        return false;
                                                                    }
                                                                }
                                                            }
                                                        }
                                                    }
                                                }
                                            }
                                            if (mName.equals("stream") || mName.equals("filter") || mName.equals("map") || 
                                                mName.equals("flatMap") || mName.equals("peek") || mName.equals("distinct") || 
                                                mName.equals("sorted") || mName.equals("limit") || mName.equals("skip")) {
                                                caller = chainMi.getExpression();
                                            } else {
                                                break;
                                            }
                                        }
                                        
                                        if (!hasMap) {
                                            if (caller instanceof SimpleName callerSn) {
                                                String callerName = callerSn.getIdentifier();
                                                if (!callerName.equals(cleanFieldName)) {
                                                    String callerType = getVariableDeclaredType(methodFqn, callerName);
                                                    if (callerType != null && callerType.contains("<")) {
                                                        int start = callerType.indexOf('<');
                                                        int end = callerType.lastIndexOf('>');
                                                        if (start >= 0 && end > start) {
                                                            foundType[0] = callerType.substring(start + 1, end);
                                                            return false;
                                                        }
                                                    }
                                                }
                                            } else if (caller instanceof FieldAccess fa) {
                                                String callerName = fa.getName().getIdentifier();
                                                if (!callerName.equals(cleanFieldName)) {
                                                    String callerType = getVariableDeclaredType(methodFqn, callerName);
                                                    if (callerType != null && callerType.contains("<")) {
                                                        int start = callerType.indexOf('<');
                                                        int end = callerType.lastIndexOf('>');
                                                        if (start >= 0 && end > start) {
                                                            foundType[0] = callerType.substring(start + 1, end);
                                                            return false;
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                            return super.visit(node);
                        }
                    });
                }
                if (foundType[0] != null) return foundType[0];
            }
            
            // Fallback to fields including superclasses
            String fieldType = getFieldType(td, cleanFieldName, context, new java.util.HashSet<>());
            if (fieldType != null) return fieldType;
        }
    }
        
        // Final fallback: ask constantExtractor to resolve it via method return constant
        if (constantExtractor != null && methodFqn.contains(".")) {
            String className = methodFqn.substring(0, methodFqn.lastIndexOf('.'));
            String methodName = methodFqn.substring(methodFqn.lastIndexOf('.') + 1);
            TypeDeclaration td = context.getTypeDeclaration(className);
            CompilationUnit cu = td != null && td.getRoot() instanceof CompilationUnit ? (CompilationUnit) td.getRoot() : null;
            List<String> values = constantExtractor.resolveMethodReturnConstant(cleanFieldName, methodName, 0, new HashSet<>(), cu);
            if (values != null && !values.isEmpty()) {
                // If it resolves to some class constant, we could infer type or return a fallback
                return null; 
            }
        }
        return null;
    }

    public String traceLocalVariable(String methodFqn, String varName) {
        if (methodFqn == null || !methodFqn.contains(".")) return null;
        String className = methodFqn.substring(0, methodFqn.lastIndexOf('.'));
        String methodName = methodFqn.substring(methodFqn.lastIndexOf('.') + 1);
        TypeDeclaration td = context.getTypeDeclaration(className);
        if (td != null) {
            MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
            if (md != null && md.getBody() != null) {
                List<Expression> initializers = new ArrayList<>();
                md.getBody().accept(new ASTVisitor() {
                    @Override
                    public boolean visit(VariableDeclarationFragment node) {
                        if (node.getName().getIdentifier().equals(varName) && node.getInitializer() != null) {
                            initializers.add(peelExpression(node.getInitializer()));
                        }
                        return super.visit(node);
                    }
                    @Override
                    public boolean visit(Assignment node) {
                        if (node.getLeftHandSide() instanceof SimpleName asn && asn.getIdentifier().equals(varName)) {
                            initializers.add(peelExpression(node.getRightHandSide()));
                        }
                        return super.visit(node);
                    }
                });
                
                if (!initializers.isEmpty()) {
                    List<String> stringified = new ArrayList<>();
                    for (Expression expr : initializers) {
                        Expression traced = traceVariable(expr);
                        if (traced instanceof MethodInvocation mi) {
                            Expression innerMost = unwrapMethodInvocation(mi, 0, methodFqn);
                            stringified.add(innerMost.toString());
                        } else if (traced instanceof ArrayInitializer) {
                            stringified.add("new Object[]" + traced.toString());
                        } else {
                            stringified.add(traced.toString());
                        }
                    }
                    if (stringified.size() == 1) {
                        return stringified.get(0).replaceAll("->\\s*yield\\s+", "-> ");
                    }
                    
                    StringBuilder sb = new StringBuilder();
                    for (int i = 0; i < stringified.size() - 1; i++) {
                        sb.append("true ? ").append(stringified.get(i)).append(" : ");
                    }
                    sb.append(stringified.get(stringified.size() - 1));
                    return sb.toString().replaceAll("->\\s*yield\\s+", "-> ");
                }
            }
            
            // If local variable not found, check field assignments
            String cleanFieldName = varName.startsWith("this.") ? varName.substring(5) : varName;
            final Expression[] fieldInitializer = new Expression[1];
            ASTVisitor fieldVisitor = new ASTVisitor() {
                @Override
                public boolean visit(VariableDeclarationFragment node) {
                    if (node.getName().getIdentifier().equals(cleanFieldName) && node.getInitializer() != null) {
                        fieldInitializer[0] = node.getInitializer();
                    }
                    return super.visit(node);
                }
                @Override
                public boolean visit(Assignment node) {
                    Expression left = node.getLeftHandSide();
                    if ((left instanceof SimpleName && ((SimpleName) left).getIdentifier().equals(cleanFieldName)) ||
                        (left instanceof FieldAccess && ((FieldAccess) left).getName().getIdentifier().equals(cleanFieldName))) {
                        fieldInitializer[0] = node.getRightHandSide();
                    }
                    return super.visit(node);
                }
            };
            
            for (FieldDeclaration fd : td.getFields()) {
                for (Object fragObj : fd.fragments()) {
                    VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
                    if (frag.getName().getIdentifier().equals(cleanFieldName) && frag.getInitializer() != null) {
                        fieldInitializer[0] = frag.getInitializer();
                    }
                }
            }
            if (fieldInitializer[0] != null) return fieldInitializer[0].toString();

            for (MethodDeclaration classMd : td.getMethods()) {
                if (classMd.isConstructor() && classMd.getBody() != null) {
                    classMd.getBody().accept(fieldVisitor);
                }
            }
            if (fieldInitializer[0] != null) return fieldInitializer[0].toString();
            
            for (MethodDeclaration classMd : td.getMethods()) {
                if (!classMd.isConstructor() && classMd.getBody() != null) {
                    classMd.getBody().accept(fieldVisitor);
                }
            }
            if (fieldInitializer[0] != null) return fieldInitializer[0].toString();
        }
        return null;
    }

    public String traceLocalVariable(String methodFqn, String varName, Map<String, String> parameterValues) {
        String result = traceLocalVariable(methodFqn, varName);
        if (result != null && parameterValues != null && !parameterValues.isEmpty()) {
            TypeDeclaration td = context.getTypeDeclaration(methodFqn.substring(0, methodFqn.lastIndexOf('.')));
            if (td != null) {
                String methodName = methodFqn.substring(methodFqn.lastIndexOf('.') + 1);
                MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
                if (md != null && md.getBody() != null) {
                    final ASTNode[] switchNode = new ASTNode[1];
                    md.getBody().accept(new ASTVisitor() {
                        @Override
                        public boolean visit(VariableDeclarationFragment node) {
                            if (node.getName().getIdentifier().equals(varName) && node.getInitializer() != null) {
                                Expression init = node.getInitializer();
                                if (init.getNodeType() == ASTNode.SWITCH_EXPRESSION) {
                                    switchNode[0] = init;
                                    return false;
                                }
                                if (init.getNodeType() == ASTNode.SWITCH_STATEMENT) {
                                    switchNode[0] = init;
                                    return false;
                                }
                            }
                            return super.visit(node);
                        }
                        @Override
                        public boolean visit(Assignment node) {
                            if (node.getLeftHandSide() instanceof SimpleName asn && asn.getIdentifier().equals(varName)) {
                                Expression rhs = node.getRightHandSide();
                                if (rhs.getNodeType() == ASTNode.SWITCH_EXPRESSION) {
                                    switchNode[0] = rhs;
                                    return false;
                                }
                                if (rhs.getNodeType() == ASTNode.SWITCH_STATEMENT) {
                                    switchNode[0] = rhs;
                                    return false;
                                }
                            }
                            return super.visit(node);
                        }
                    });
                    if (switchNode[0] != null) {
                        if (switchNode[0] instanceof SwitchExpression se) {
                            String evaluated = constantResolver.evaluateSwitchWithParams(se, parameterValues, context);
                            if (evaluated != null) return evaluated;
                        } else if (switchNode[0] instanceof SwitchStatement ss) {
                            String evaluated = constantResolver.evaluateSwitchWithParams(ss, parameterValues, context);
                            if (evaluated != null) return evaluated;
                        }
                    }
                }
            }
        }
        return result;
    }

    public Map<String, String> buildParameterValuesMap(String caller, String target, Map<String, List<CallEdge>> callGraph, List<String> path, int pathIndex) {
        Map<String, String> paramValues = new HashMap<>();
        List<CallEdge> edges = callGraph.get(caller);
        
        boolean hasEdge = false;
        if (edges != null) {
            for (CallEdge edge : edges) {
                if (edge.getTargetMethod().equals(target) || isHeuristicMatch(edge.getTargetMethod(), target)) {
                    hasEdge = true;
                    break;
                }
            }
        }

        // Bridge polymorphic interfaces to implementations without graph edges
        if (!hasEdge) {
            String callerSimple = caller.contains(".") ? caller.substring(caller.lastIndexOf('.') + 1) : caller;
            String targetSimple = target.contains(".") ? target.substring(target.lastIndexOf('.') + 1) : target;
            if (callerSimple.equals(targetSimple)) {
                List<String> callerParams = getParameterNames(caller);
                List<String> targetParams = getParameterNames(target);
                if (callerParams != null && targetParams != null && callerParams.size() == targetParams.size()) {
                    for (int i = 0; i < callerParams.size(); i++) {
                        paramValues.put(targetParams.get(i), callerParams.get(i));
                    }
                    return paramValues;
                }
            }
        }

        if (edges == null) {
            return paramValues;
        }

        // We also want to support heuristic matches using a helper if isHeuristicMatch is not directly available, but let's check
        // We can check if name equals or if it resolves to target
        for (CallEdge edge : edges) {
            if (edge.getTargetMethod().equals(target) || isHeuristicMatch(edge.getTargetMethod(), target)) {
                List<String> args = edge.getArguments();
                if (args == null || args.isEmpty()) break;

                int lastDot = target.lastIndexOf('.');
                if (lastDot < 0) break;
                String className = target.substring(0, lastDot);
                String methodName = target.substring(lastDot + 1);

                TypeDeclaration td = context.getTypeDeclaration(className);
                if (td == null) break;

                MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
                if (md == null) break;

                List<String> paramNames = new ArrayList<>();
                for (Object paramObj : md.parameters()) {
                    SingleVariableDeclaration param = (SingleVariableDeclaration) paramObj;
                    paramNames.add(param.getName().getIdentifier());
                }

                int minSize = Math.min(paramNames.size(), args.size());
                for (int j = 0; j < minSize; j++) {
                    String argValue = args.get(j);
                    String resolvedArgValue = resolveArgumentValue(argValue, caller, path, pathIndex, callGraph);
                    paramValues.put(paramNames.get(j), resolvedArgValue);
                }
                break;
            }
        }
        return paramValues;
    }

    private List<String> getParameterNames(String methodFqn) {
        int lastDot = methodFqn.lastIndexOf('.');
        if (lastDot < 0) return null;
        String className = methodFqn.substring(0, lastDot);
        String methodName = methodFqn.substring(lastDot + 1);

        TypeDeclaration td = context.getTypeDeclaration(className);
        if (td == null) return null;

        MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
        if (md == null) return null;

        List<String> paramNames = new ArrayList<>();
        for (Object paramObj : md.parameters()) {
            SingleVariableDeclaration param = (SingleVariableDeclaration) paramObj;
            paramNames.add(param.getName().getIdentifier());
        }
        return paramNames;
    }

    public String resolveArgumentValue(String argValue, String callerMethod, List<String> path, int pathIndex, Map<String, List<CallEdge>> callGraph) {
        if (argValue == null) return null;
        if (argValue.contains(".") || argValue.startsWith("new ") || argValue.matches(".*[0-9].*") || (argValue.startsWith("\"") && argValue.endsWith("\""))) {
            return argValue;
        }

        String tracedLocal = traceLocalVariable(callerMethod, argValue);
        if (tracedLocal != null && !tracedLocal.equals(argValue)) {
            if (tracedLocal.contains(".") || tracedLocal.startsWith("new ") || tracedLocal.matches(".*[0-9].*") || tracedLocal.contains("(") || (tracedLocal.startsWith("\"") && tracedLocal.endsWith("\""))) {
                return tracedLocal;
            }
            argValue = tracedLocal;
        }

        int callerIndex = path.indexOf(callerMethod);
        if (callerIndex <= 0) return argValue;

        String prevCaller = path.get(callerIndex - 1);
        List<CallEdge> prevEdges = callGraph.get(prevCaller);
        if (prevEdges == null) return argValue;

        for (CallEdge edge : prevEdges) {
            if (edge.getTargetMethod().equals(callerMethod) || isHeuristicMatch(edge.getTargetMethod(), callerMethod)) {
                List<String> prevArgs = edge.getArguments();
                if (prevArgs == null || prevArgs.isEmpty()) break;

                int lastDot = callerMethod.lastIndexOf('.');
                if (lastDot < 0) break;
                String className = callerMethod.substring(0, lastDot);
                String methodName = callerMethod.substring(lastDot + 1);

                TypeDeclaration td = context.getTypeDeclaration(className);
                if (td == null) break;

                MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
                if (md == null) break;

                List<String> paramNames = new ArrayList<>();
                for (Object paramObj : md.parameters()) {
                    SingleVariableDeclaration param = (SingleVariableDeclaration) paramObj;
                    paramNames.add(param.getName().getIdentifier());
                }

                int paramIdx = paramNames.indexOf(argValue);
                if (paramIdx >= 0 && paramIdx < prevArgs.size()) {
                    String prevArg = prevArgs.get(paramIdx);
                    return resolveArgumentValue(prevArg, prevCaller, path, pathIndex, callGraph);
                }
                break;
            }
        }
        return argValue;
    }

    private boolean isHeuristicMatch(String neighbor, String target) {
        if (neighbor == null || target == null) return false;
        if (neighbor.equals(target)) return true;

        if (neighbor.contains(".") && target.contains(".")) {
            String methodNeighbor = neighbor.substring(neighbor.lastIndexOf('.') + 1);
            String methodTarget = target.substring(target.lastIndexOf('.') + 1);
            if (!methodNeighbor.equals(methodTarget)) return false;

            String classNeighbor = neighbor.substring(0, neighbor.lastIndexOf('.'));
            String classTarget = target.substring(0, target.lastIndexOf('.'));
            if (classNeighbor.equals(classTarget)) return true;

            String simpleClassNeighbor = classNeighbor.contains(".") ? classNeighbor.substring(classNeighbor.lastIndexOf('.') + 1) : classNeighbor;
            String simpleClassTarget = classTarget.contains(".") ? classTarget.substring(classTarget.lastIndexOf('.') + 1) : classTarget;
            if (simpleClassNeighbor.equals(simpleClassTarget)) return true;

            List<String> impls = context.getImplementations(classTarget);
            if (impls != null) {
                for (String impl : impls) {
                    if (impl.equals(classNeighbor) || impl.endsWith("." + classNeighbor)) return true;
                }
            }

            List<String> implsN = context.getImplementations(classNeighbor);
            if (implsN != null) {
                for (String impl : implsN) {
                    if (impl.equals(classTarget) || impl.endsWith("." + classTarget)) return true;
                }
            }
            return false;
        }

        String simpleTarget = target.contains(".") ? target.substring(target.lastIndexOf('.') + 1) : target;
        String simpleNeighbor = neighbor.contains(".") ? neighbor.substring(neighbor.lastIndexOf('.') + 1) : neighbor;
        if (!simpleNeighbor.equals(simpleTarget)) return false;
        
        String classNeighbor = neighbor.contains(".") ? neighbor.substring(0, neighbor.lastIndexOf('.')) : null;
        String classTarget = target.contains(".") ? target.substring(0, target.lastIndexOf('.')) : null;
        String simpleClassTarget = classTarget != null && classTarget.contains(".") ? classTarget.substring(classTarget.lastIndexOf('.') + 1) : classTarget;
        String simpleClassNeighbor = classNeighbor != null && classNeighbor.contains(".") ? classNeighbor.substring(classNeighbor.lastIndexOf('.') + 1) : classNeighbor;

        if (simpleClassNeighbor != null && simpleClassTarget != null) {
            if (simpleClassNeighbor.equalsIgnoreCase(simpleClassTarget)) return true;
            if (simpleClassNeighbor.equals("this") || simpleClassNeighbor.equals("super")) return true;
            return false;
        }
        
        return true;
    }

    private MethodDeclaration findEnclosingMethod(ASTNode node) {
        ASTNode parent = node.getParent();
        while (parent != null && !(parent instanceof MethodDeclaration)) {
            parent = parent.getParent();
        }
        return (MethodDeclaration) parent;
    }

    private TypeDeclaration findEnclosingType(ASTNode node) {
        ASTNode parent = node.getParent();
        while (parent != null && !(parent instanceof TypeDeclaration)) {
            parent = parent.getParent();
        }
        return (TypeDeclaration) parent;
    }

    private Expression unwrapMethodInvocation(MethodInvocation mi, int depth, String methodFqn) {
        if (depth > 5) return mi;
        
        String targetMethodFqn = resolveTargetMethodFqn(mi, methodFqn);
        if (targetMethodFqn != null && targetMethodFqn.contains(".")) {
            String className = targetMethodFqn.substring(0, targetMethodFqn.lastIndexOf('.'));
            String methodName = targetMethodFqn.substring(targetMethodFqn.lastIndexOf('.') + 1);
            
            TypeDeclaration currentTd = methodFqn != null && methodFqn.contains(".") ? context.getTypeDeclaration(methodFqn.substring(0, methodFqn.lastIndexOf('.'))) : null;
            CompilationUnit cu = currentTd != null && currentTd.getRoot() instanceof CompilationUnit ? (CompilationUnit) currentTd.getRoot() : null;
            
            TypeDeclaration td = context.getTypeDeclaration(className, cu);
            if (td != null) {
                MethodDeclaration localMd = context.findMethodDeclaration(td, methodName, true);
                if (localMd != null) {
                    int paramIdx = getReturnedParameterIndex(localMd);
                    if (paramIdx >= 0 && paramIdx < mi.arguments().size()) {
                        Expression arg = peelExpression((Expression) mi.arguments().get(paramIdx));
                        if (arg instanceof MethodInvocation innerMi) {
                            return unwrapMethodInvocation(innerMi, depth + 1, methodFqn);
                        }
                        return arg;
                    }
                }
            }
        }
        
        if (mi.getExpression() != null) {
            String exprStr = mi.getExpression().toString();
            if (!exprStr.equals("Optional") && !exprStr.equals("Stream") && !exprStr.equals("List") && !exprStr.equals("Set") && !exprStr.equals("Arrays") && !exprStr.equals("Objects") && !exprStr.equals("Mono") && !exprStr.equals("Flux")) {
                return mi;
            }
        }
        
        if (!mi.arguments().isEmpty()) {
            String mName = mi.getName().getIdentifier();
            String lowerName = mName.toLowerCase();
            
            if (lowerName.startsWith("map") || lowerName.startsWith("parse") || 
                lowerName.startsWith("convert") || lowerName.startsWith("to") || 
                lowerName.startsWith("resolve") || lowerName.startsWith("build") || 
                lowerName.startsWith("create") || lowerName.startsWith("from") ||
                lowerName.startsWith("transform") || lowerName.startsWith("translate") ||
                lowerName.startsWith("derive") || lowerName.startsWith("determine") ||
                lowerName.startsWith("calculate") || lowerName.startsWith("decode") ||
                lowerName.startsWith("extract") || lowerName.startsWith("get") ||
                lowerName.startsWith("find") || lowerName.startsWith("validate")) {
                return mi;
            }
            
            Expression arg = peelExpression((Expression) mi.arguments().get(0));
            if (arg instanceof MethodInvocation innerMi) {
                return unwrapMethodInvocation(innerMi, depth + 1, methodFqn);
            }
            return arg;
        }
        return mi;
    }
    
    private String resolveTargetMethodFqn(MethodInvocation mi, String currentMethodFqn) {
        if (currentMethodFqn == null || !currentMethodFqn.contains(".")) return null;
        String currentClass = currentMethodFqn.substring(0, currentMethodFqn.lastIndexOf('.'));
        if (mi.getExpression() == null || mi.getExpression().toString().equals("this")) {
            return currentClass + "." + mi.getName().getIdentifier();
        }
        if (mi.getExpression() instanceof SimpleName sn) {
            String declaredType = getVariableDeclaredType(currentMethodFqn, sn.getIdentifier());
            if (declaredType != null) {
                return declaredType + "." + mi.getName().getIdentifier();
            }
        } else if (mi.getExpression() instanceof FieldAccess fa) {
            String declaredType = getVariableDeclaredType(currentMethodFqn, fa.getName().getIdentifier());
            if (declaredType != null) {
                return declaredType + "." + mi.getName().getIdentifier();
            }
        }
        return null;
    }

    private int getReturnedParameterIndex(MethodDeclaration md) {
        if (md.getBody() == null) return -1;
        List<?> parameters = md.parameters();
        if (parameters.isEmpty()) return -1;
        
        final List<String> returnedExprs = new ArrayList<>();
        md.getBody().accept(new ASTVisitor() {
            @Override
            public boolean visit(ReturnStatement node) {
                if (node.getExpression() != null) {
                    returnedExprs.add(node.getExpression().toString());
                }
                return super.visit(node);
            }
        });
        
        if (returnedExprs.size() == 1) {
            String retStr = returnedExprs.get(0);
            for (int i = 0; i < parameters.size(); i++) {
                if (parameters.get(i) instanceof SingleVariableDeclaration svd) {
                    if (svd.getName().getIdentifier().equals(retStr)) {
                        return i;
                    }
                }
            }
        }
        return -1;
    }
}
