package click.kamil.springstatemachineexporter.analysis.service;

import click.kamil.springstatemachineexporter.analysis.model.CallChain;
import click.kamil.springstatemachineexporter.analysis.model.CallEdge;
import click.kamil.springstatemachineexporter.analysis.model.EntryPoint;
import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import lombok.extern.slf4j.Slf4j;
import org.eclipse.jdt.core.dom.*;

import java.util.*;

import click.kamil.springstatemachineexporter.analysis.spring.InjectionPointAnalyzer;

@Slf4j
public class JdtCallGraphEngine extends AbstractCallGraphEngine {
    private final InjectionPointAnalyzer injectionAnalyzer;
    private String currentMethodFqn;
    private Map<String, List<CallEdge>> graph;

    public JdtCallGraphEngine(CodebaseContext context, InjectionPointAnalyzer injectionAnalyzer) {
        super(context);
        this.injectionAnalyzer = injectionAnalyzer;
    }

    @SuppressWarnings("unchecked")
    protected Map<String, List<CallEdge>> buildCallGraph() {
        Map<String, List<CallEdge>> cached = (Map<String, List<CallEdge>>) context.getCache().get("jdtCallGraph");
        if (cached != null) {
            this.graph = cached;
            return cached;
        }
        graph = new HashMap<>();
        for (CompilationUnit cu : context.getCompilationUnits()) {
            cu.accept(new ASTVisitor() {

                @Override
                public boolean visit(MethodInvocation node) {
                    MethodDeclaration md = findEnclosingMethod(node);
                    if (md != null) {
                        TypeDeclaration td = findEnclosingType(md);
                        if (td != null) {
                            String currentMethodFqn = context.getFqn(td) + "." + md.getName().getIdentifier();
                            List<String> calledMethods = resolveCalledMethodsPolymorphic(node);
                            List<String> args = resolveArguments(node.arguments());
                            String constraint = click.kamil.springstatemachineexporter.ast.common.AstUtils.findConditionConstraint(node);
                            for (String calledMethod : calledMethods) {
                                CallEdge edge = new CallEdge(calledMethod, args);
                                edge.setConstraint(resolveConstraint(node, calledMethod, constraint));
                                graph.computeIfAbsent(currentMethodFqn, k -> new ArrayList<>()).add(edge);
                            }
                            for (Object argObj : node.arguments()) {
                                if (argObj instanceof ExpressionMethodReference emr) {
                                    String typeName = emr.getExpression().toString();
                                    if ("this".equals(typeName) || "super".equals(typeName)) {
                                        TypeDeclaration td2 = findEnclosingType(node);
                                        if (td2 != null) {
                                            String refMethod = resolveMethodInType(td2, emr.getName().getIdentifier());
                                            if (refMethod != null) {
                                                List<String> implicitArgs = new ArrayList<>();
                                                if (click.kamil.springstatemachineexporter.ast.common.AstUtils.isFunctionalCollectionMethod(node)) {
                                                    implicitArgs.add(node.getExpression().toString());
                                                } else {
                                                    implicitArgs.addAll(args);
                                                }
                                                CallEdge edge = new CallEdge(refMethod, implicitArgs);
                                                edge.setConstraint(constraint);
                                                graph.computeIfAbsent(currentMethodFqn, k -> new ArrayList<>()).add(edge);
                                            }
                                        }
                                    } else {
                                        String fallbackTypeFqn = null;
                                        ITypeBinding binding = emr.getExpression().resolveTypeBinding();
                                        if (binding != null) {
                                            fallbackTypeFqn = binding.getQualifiedName();
                                        } else if (emr.getExpression() instanceof SimpleName sn) {
                                            fallbackTypeFqn = resolveReceiverTypeFallback(sn);
                                        }
                                        if (fallbackTypeFqn != null) {
                                            String calledMethod = fallbackTypeFqn + "." + emr.getName().getIdentifier();
                                            List<String> implicitArgs = new ArrayList<>();
                                            if (click.kamil.springstatemachineexporter.ast.common.AstUtils.isFunctionalCollectionMethod(node)) {
                                                implicitArgs.add(node.getExpression().toString());
                                            } else {
                                                implicitArgs.addAll(args);
                                            }
                                            CallEdge edge = new CallEdge(calledMethod, implicitArgs);
                                            edge.setConstraint(constraint);
                                            graph.computeIfAbsent(currentMethodFqn, k -> new ArrayList<>()).add(edge);
                                            
                                            List<String> impls = context.getImplementations(fallbackTypeFqn);
                                            for (String impl : impls) {
                                                CallEdge implEdge = new CallEdge(impl + "." + emr.getName().getIdentifier(), args);
                                                implEdge.setConstraint(constraint);
                                                graph.computeIfAbsent(currentMethodFqn, k -> new ArrayList<>()).add(implEdge);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                    return super.visit(node);
                }

                @Override
                public boolean visit(SuperMethodInvocation node) {
                    MethodDeclaration md = findEnclosingMethod(node);
                    if (md != null) {
                        TypeDeclaration tdOuter = findEnclosingType(md);
                        if (tdOuter != null) {
                            String currentMethodFqn = context.getFqn(tdOuter) + "." + md.getName().getIdentifier();
                            String methodName = node.getName().getIdentifier();
                            TypeDeclaration td = findEnclosingType(node);
                            if (td != null) {
                                String superFqn = context.getSuperclassFqn(td);
                                if (superFqn != null) {
                                    TypeDeclaration superTd = context.getTypeDeclaration(superFqn);
                                    String calledMethod = null;
                                    if (superTd != null) {
                                        calledMethod = resolveMethodInType(superTd, methodName);
                                    }
                                    if (calledMethod == null) {
                                        calledMethod = superFqn + "." + methodName;
                                    }
                                    List<String> args = resolveArguments(node.arguments());
                                    String constraint = click.kamil.springstatemachineexporter.ast.common.AstUtils.findConditionConstraint(node);
                                    CallEdge edge = new CallEdge(calledMethod, args);
                                    edge.setConstraint(constraint);
                                    graph.computeIfAbsent(currentMethodFqn, k -> new ArrayList<>()).add(edge);
                                }
                            }
                        }
                    }
                    return super.visit(node);
                }
            });
        }
        context.getCache().put("jdtCallGraph", graph);
        return graph;
    }

    @Override
    protected String resolveArgument(org.eclipse.jdt.core.dom.Expression expr) {
        if (expr == null) {
            return "null";
        }
        // Extract from constructor args (e.g., new CustomMessage(OrderEvent.PROCESS, ...))
        org.eclipse.jdt.core.dom.Expression unwrappedBuilder = unwrapMessageBuilder(expr);
        if (unwrappedBuilder != expr) {
            expr = unwrappedBuilder;
        }

        // Trace variable to resolve local variable references
        org.eclipse.jdt.core.dom.Expression tracedExpr = traceVariable(expr);
        if (tracedExpr instanceof org.eclipse.jdt.core.dom.QualifiedName 
                || tracedExpr instanceof org.eclipse.jdt.core.dom.ClassInstanceCreation 
                || tracedExpr instanceof org.eclipse.jdt.core.dom.StringLiteral 
                || tracedExpr instanceof org.eclipse.jdt.core.dom.NumberLiteral) {
            expr = tracedExpr; // Only accept trace if it resolves to a pure constant/primitive/constructor
        }

        // Delegate to base class for lambda, method reference, and constructor unwrapping
        String result = super.resolveArgument(expr);
        if (result != null) {
            result = result.replaceAll("->\\s*yield\\s+", "-> ");
        }
        return result;
    }

    protected String resolveCalledMethod(MethodInvocation node) {
        Expression receiver = node.getExpression();
        String methodName = node.getName().getIdentifier();

        if (receiver instanceof MethodInvocation miReceiver) {
            String returnType = resolveMethodInvocationReturnType(miReceiver);
            if (returnType != null) {
                return returnType + "." + methodName;
            }
        }

        if (receiver == null) {
            TypeDeclaration td = findEnclosingType(node);
            if (td != null) {
                return resolveMethodInType(td, methodName);
            }
            return null;
        }

        if (injectionAnalyzer != null) {
            IBinding nameBinding = null;
            if (receiver instanceof SimpleName sn) {
                nameBinding = sn.resolveBinding();
            } else if (receiver instanceof FieldAccess fa) {
                nameBinding = fa.resolveFieldBinding();
            }
            if (nameBinding instanceof IVariableBinding varBinding) {
                log.debug("CALLGRAPH RESOLVING BEAN FOR BINDING: {} calling {}", varBinding.getName(), methodName);
                String concreteFqn = injectionAnalyzer.resolveInjectedBeanFqn(varBinding);
                if (concreteFqn != null) {
                    log.debug("  -> RESOLVED CONCRETE FQN: {}", concreteFqn);
                    return concreteFqn + "." + methodName;
                } else {
                    log.debug("  -> RESOLVE FAILED, RETURNED NULL");
                }
            } else {
                log.debug("CALLGRAPH NAME BINDING IS NOT VARIABLE: {} calling {}", nameBinding, methodName);
            }
        }

        ITypeBinding binding = receiver.resolveTypeBinding();
        if (binding != null) {
            String typeName = null;
            try {
                if (binding.getErasure() != null) {
                    typeName = binding.getErasure().getQualifiedName();
                    if (typeName == null || typeName.isEmpty()) {
                        typeName = binding.getErasure().getName();
                    }
                }
            } catch (Exception e) {
                // Ignore JDT internal exceptions
            }

            if (typeName == null || typeName.isEmpty()) {
                typeName = binding.getQualifiedName();
                if (typeName == null || typeName.isEmpty()) {
                    typeName = binding.getName();
                }
                if (typeName != null && typeName.contains("<")) {
                    typeName = typeName.replaceAll("<.*>", "");
                }
            }

            if (typeName != null && !typeName.isEmpty()) {
                org.eclipse.jdt.core.dom.TypeDeclaration resolvedTd = context.getTypeDeclaration(typeName);
                if (resolvedTd != null && context.findMethodDeclaration(resolvedTd, methodName, true) != null) {
                    return typeName + "." + methodName;
                }
            }
        }

        if (receiver instanceof ThisExpression) {
            TypeDeclaration td = findEnclosingType(node);
            if (td != null) {
                return context.getFqn(td) + "." + methodName;
            }
        }

        if (receiver instanceof SimpleName sn) {
            String fallbackTypeFqn = resolveReceiverTypeFallback(sn);
            if (fallbackTypeFqn != null) {
                return fallbackTypeFqn + "." + methodName;
            }
            String receiverName = sn.getIdentifier();
            return receiverName + "." + methodName; 
        }

        return null;
    }

    private String resolveReceiverTypeFallback(SimpleName receiverNameNode) {
        String varName = receiverNameNode.getIdentifier();
        
        // 1. Check local variables in enclosing method
        MethodDeclaration enclosingMethod = findEnclosingMethod(receiverNameNode);
        if (enclosingMethod != null) {
            // Check parameters
            for (Object paramObj : enclosingMethod.parameters()) {
                SingleVariableDeclaration svd = (SingleVariableDeclaration) paramObj;
                if (svd.getName().getIdentifier().equals(varName)) {
                    return resolveTypeToFqn(svd.getType(), receiverNameNode);
                }
            }
            // Check method body (local variables)
            if (enclosingMethod.getBody() != null) {
                Type[] foundType = new Type[1];
                enclosingMethod.getBody().accept(new ASTVisitor() {
                    @Override
                    public boolean visit(VariableDeclarationStatement node) {
                        for (Object fragObj : node.fragments()) {
                            VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
                            if (frag.getName().getIdentifier().equals(varName)) {
                                foundType[0] = node.getType();
                            }
                        }
                        return super.visit(node);
                    }
                });
                if (foundType[0] != null) {
                    return resolveTypeToFqn(foundType[0], receiverNameNode);
                }
            }
        }

        // 2. Check fields in enclosing class and its superclasses
        TypeDeclaration enclosingType = findEnclosingType(receiverNameNode);
        TypeDeclaration currentType = enclosingType;
        while (currentType != null) {
            for (FieldDeclaration field : currentType.getFields()) {
                for (Object fragObj : field.fragments()) {
                    VariableDeclarationFragment frag = (VariableDeclarationFragment) fragObj;
                    if (frag.getName().getIdentifier().equals(varName)) {
                        return resolveTypeToFqn(field.getType(), receiverNameNode);
                    }
                }
            }
            String superclass = context.getSuperclassFqn(currentType);
            currentType = superclass != null ? context.getTypeDeclaration(superclass) : null;
        }

        return null;
    }

    private String resolveTypeToFqn(Type type, ASTNode contextNode) {
        if (type == null) return null;
        String simpleName;
        if (type.isSimpleType()) {
            simpleName = ((SimpleType) type).getName().getFullyQualifiedName();
        } else if (type.isParameterizedType()) {
            return resolveTypeToFqn(((ParameterizedType) type).getType(), contextNode);
        } else {
            simpleName = type.toString();
        }

        CompilationUnit cu = (CompilationUnit) contextNode.getRoot();
        TypeDeclaration td = context.getTypeDeclaration(simpleName, cu);
        if (td != null) {
            return context.getFqn(td);
        }
        
        // Fallback to import matching if CodebaseContext doesn't know it (e.g., external library)
        for (Object impObj : cu.imports()) {
            ImportDeclaration imp = (ImportDeclaration) impObj;
            String impName = imp.getName().getFullyQualifiedName();
            if (impName.endsWith("." + simpleName)) {
                return impName;
            }
        }

        return simpleName;
    }

    private Expression unwrapMessageBuilder(Expression expr) {
        if (expr instanceof MethodInvocation mi) {
            MethodInvocation current = mi;
            while (current != null) {
                String name = current.getName().getIdentifier();
                if (("withPayload".equals(name) || "just".equals(name)) && !current.arguments().isEmpty()) {
                    return (Expression) current.arguments().get(0);
                }
                Expression receiver = current.getExpression();
                if (receiver instanceof MethodInvocation nextMi) {
                    current = nextMi;
                } else {
                    current = null;
                }
            }
        }
        return expr;
    }

    protected String resolveMethodInvocationReturnType(MethodInvocation mi) {
        Expression receiver = mi.getExpression();
        String receiverType = null;
        if (receiver == null) {
            TypeDeclaration td = findEnclosingType(mi);
            if (td != null) {
                receiverType = context.getFqn(td);
            }
        } else if (receiver instanceof SimpleName sn) {
            receiverType = resolveReceiverTypeFallback(sn);
        } else if (receiver instanceof FieldAccess fa) {
            receiverType = resolveReceiverTypeFallback(fa.getName());
        } else if (receiver instanceof MethodInvocation innerMi) {
            receiverType = resolveMethodInvocationReturnType(innerMi);
        }
        
        if (receiverType == null) {
            ITypeBinding binding = receiver != null ? receiver.resolveTypeBinding() : null;
            if (binding != null) {
                receiverType = binding.getErasure().getQualifiedName();
            }
        }
        
        if (receiverType != null) {
            String rawType = receiverType;
            if (rawType.contains("<")) {
                rawType = rawType.substring(0, rawType.indexOf('<'));
            }
            if (rawType.equals("java.util.Map") || rawType.equals("Map")) {
                String valType = resolveMapValueType(mi);
                if (valType != null) return valType;
            }
            if (receiverType.contains("<")) {
                receiverType = receiverType.substring(0, receiverType.indexOf('<'));
            }
            String methodName = mi.getName().getIdentifier();
            TypeDeclaration td = context.getTypeDeclaration(receiverType);
            if (td == null && receiverType.contains(".")) {
                String simpleClass = receiverType.substring(receiverType.lastIndexOf('.') + 1);
                td = context.getTypeDeclaration(simpleClass);
            }
            if (td != null) {
                MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
                if (md != null && md.getReturnType2() != null) {
                    return resolveTypeToFqn(md.getReturnType2(), mi);
                }
            }
        }
        return null;
    }
}
