package click.kamil.springstatemachineexporter.analysis.service;

import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import lombok.extern.slf4j.Slf4j;
import org.eclipse.jdt.core.dom.*;
import java.util.*;
import java.util.function.Function;

@Slf4j
public class ConstantExtractor {
    private final CodebaseContext context;
    private final ConstantResolver constantResolver;
    private final Function<MethodInvocation, String> calledMethodResolver;
    private VariableTracer variableTracer;
    private ConstructorAnalyzer constructorAnalyzer;

    public ConstantExtractor(CodebaseContext context, ConstantResolver constantResolver, Function<MethodInvocation, String> calledMethodResolver) {
        this.context = context;
        this.constantResolver = constantResolver;
        this.calledMethodResolver = calledMethodResolver;
    }

    public void setVariableTracer(VariableTracer variableTracer) {
        this.variableTracer = variableTracer;
    }

    public void setConstructorAnalyzer(ConstructorAnalyzer constructorAnalyzer) {
        this.constructorAnalyzer = constructorAnalyzer;
    }

    public void extractConstantsFromExpression(Expression expr, List<String> constants) {
        if (constantResolver != null) {
            String resolved = constantResolver.resolve(expr, context);
            if (resolved != null) {
                if (resolved.startsWith("ENUM_SET:")) {
                    for (String eVal : resolved.substring(9).split(",")) {
                        String parsed = parseEnumSetElement(eVal);
                        if (!constants.contains(parsed)) constants.add(parsed);
                    }
                } else {
                    constants.add(resolved);
                }
                return;
            }
        }

        if (expr instanceof QualifiedName qn) {
            constants.add(qn.toString());
        } else if (expr instanceof StringLiteral sl) {
            constants.add(sl.getLiteralValue());
        } else if (expr instanceof ConditionalExpression ce) {
            extractConstantsFromExpression(ce.getThenExpression(), constants);
            extractConstantsFromExpression(ce.getElseExpression(), constants);
        } else if (expr instanceof ParenthesizedExpression pe) {
            extractConstantsFromExpression(pe.getExpression(), constants);
        } else if (expr instanceof CastExpression ce) {
            extractConstantsFromExpression(ce.getExpression(), constants);
        } else if (expr instanceof Assignment assignment) {
            extractConstantsFromExpression(assignment.getRightHandSide(), constants);
        } else if (expr instanceof SwitchExpression se) {
            se.accept(new ASTVisitor() {
                @Override
                public boolean visit(YieldStatement ys) {
                    extractConstantsFromExpression(ys.getExpression(), constants);
                    return super.visit(ys);
                }
                @Override
                public boolean visit(ExpressionStatement es) {
                    extractConstantsFromExpression(es.getExpression(), constants);
                    return super.visit(es);
                }
                @Override
                public boolean visit(ReturnStatement rs) {
                    extractConstantsFromExpression(rs.getExpression(), constants);
                    return super.visit(rs);
                }
            });
        } else if (expr instanceof ArrayAccess aa) {
            extractConstantsFromExpression(aa.getArray(), constants);
        } else if (expr instanceof ArrayCreation ac && ac.getInitializer() != null) {
            for (Object expObj : ac.getInitializer().expressions()) {
                extractConstantsFromArgument((Expression) expObj, constants);
            }
        } else if (expr instanceof ArrayInitializer ai) {
            for (Object expObj : ai.expressions()) {
                extractConstantsFromArgument((Expression) expObj, constants);
            }
        } else if (expr instanceof MethodInvocation mi) {
            String methodName = mi.getName().getIdentifier();
            
            log.trace("Processing mi: {}", mi);
            if (methodName.equals("get") || methodName.equals("getOrDefault")) {
                if (mi.getExpression() != null) {
                    if (variableTracer != null) {
                        List<Expression> traced = variableTracer.traceVariableAll(mi.getExpression());
                        for (Expression t : traced) {
                            extractConstantsFromExpression(t, constants);
                        }
                    } else {
                        extractConstantsFromExpression(mi.getExpression(), constants);
                    }
                }
                return;
            }
            if (methodName.equals("of") || methodName.equals("ofEntries") || methodName.equals("asList") || methodName.equals("entry") ||
                methodName.equals("just") || methodName.equals("withPayload") || methodName.equals("success")) {
                for (Object argObj : mi.arguments()) {
                    extractConstantsFromExpression((Expression) argObj, constants);
                }
                return;
            }
            
            if ((methodName.startsWith("get") || methodName.equals("type") || methodName.equals("event")) && mi.getExpression() instanceof SimpleName sn) {
                String varName = sn.getIdentifier();
                String propName = methodName.startsWith("get") ? methodName.substring(3) : methodName;
                
                Block block = findEnclosingBlock(mi);
                if (block != null) {
                    for (Object stmtObj : block.statements()) {
                        if (stmtObj == mi.getParent() || stmtObj == mi) break;
                        if (stmtObj instanceof ExpressionStatement es) {
                            if (es.getExpression() instanceof MethodInvocation setterMi) {
                                if (setterMi.getName().getIdentifier().equalsIgnoreCase("set" + propName) || setterMi.getName().getIdentifier().equalsIgnoreCase(propName)) {
                                    if (setterMi.getExpression() instanceof SimpleName setterSn && setterSn.getIdentifier().equals(varName)) {
                                        if (!setterMi.arguments().isEmpty()) {
                                            extractConstantsFromExpression((Expression) setterMi.arguments().get(0), constants);
                                            return;
                                        }
                                    }
                                }
                            } else if (es.getExpression() instanceof Assignment assignment) {
                                if (assignment.getLeftHandSide() instanceof FieldAccess fa) {
                                    if (fa.getExpression() instanceof SimpleName faSn && faSn.getIdentifier().equals(varName)) {
                                        if (fa.getName().getIdentifier().equalsIgnoreCase(propName)) {
                                            extractConstantsFromExpression(assignment.getRightHandSide(), constants);
                                            return;
                                        }
                                    }
                                } else if (assignment.getLeftHandSide() instanceof QualifiedName qqn) {
                                    if (qqn.getQualifier().getFullyQualifiedName().equals(varName)) {
                                        if (qqn.getName().getIdentifier().equalsIgnoreCase(propName)) {
                                            extractConstantsFromExpression(assignment.getRightHandSide(), constants);
                                            return;
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
            
            List<Expression> receivers = new ArrayList<>();
            if (mi.getExpression() != null) {
                if (mi.getExpression() instanceof ClassInstanceCreation) {
                    receivers.add(mi.getExpression());
                } else if (variableTracer != null) {
                    receivers.addAll(variableTracer.traceVariableAll(mi.getExpression()));
                }
            }
            
            for (Expression receiverExpr : receivers) {
                if (receiverExpr instanceof ClassInstanceCreation cic) {
                    if (cic.getAnonymousClassDeclaration() != null) {
                        for (Object declObj : cic.getAnonymousClassDeclaration().bodyDeclarations()) {
                            if (declObj instanceof MethodDeclaration mdecl) {
                                if (mdecl.getName().getIdentifier().equals(methodName) && mdecl.getBody() != null) {
                                    mdecl.getBody().accept(new ASTVisitor() {
                                        @Override
                                        public boolean visit(ReturnStatement rs) {
                                            if (rs.getExpression() != null) {
                                                extractConstantsFromExpression(rs.getExpression(), constants);
                                            }
                                            return super.visit(rs);
                                        }
                                    });
                                }
                            }
                        }
                    } else if (constructorAnalyzer != null) {
                        boolean handledByLombok = false;
                        if (methodName.startsWith("get") && methodName.length() > 3) {
                            String typeName = click.kamil.springstatemachineexporter.ast.common.AstUtils.extractSimpleTypeName(cic.getType());
                            TypeDeclaration typeDecl = context.getTypeDeclaration(typeName);
                            if (typeDecl != null) {
                                String fieldName = Character.toLowerCase(methodName.charAt(3)) + methodName.substring(4);
                                int fieldIndex = constructorAnalyzer.findConstructorArgumentIndexForLombokField(typeDecl, fieldName);
                                if (fieldIndex >= 0 && fieldIndex < cic.arguments().size()) {
                                    extractConstantsFromArgument((Expression) cic.arguments().get(fieldIndex), constants);
                                    handledByLombok = true;
                                }
                            }
                        }
                        if (!handledByLombok && methodName != null && !methodName.isEmpty()) {
                            List<String> narrowed = constructorAnalyzer.resolveInlineInstantiationGetterResult(
                                    cic, methodName, context, new HashSet<>());
                            if (narrowed.isEmpty()) {
                                String typeName = click.kamil.springstatemachineexporter.ast.common.AstUtils.extractSimpleTypeName(cic.getType());
                                CompilationUnit contextCu = expr.getRoot() instanceof CompilationUnit ? (CompilationUnit) expr.getRoot() : null;
                                TypeDeclaration typeDecl = contextCu != null ? context.getTypeDeclaration(typeName, contextCu) : context.getTypeDeclaration(typeName);
                                if (typeDecl == null) typeDecl = context.getTypeDeclaration(typeName);
                                if (typeDecl != null) {
                                    narrowed = resolveMethodReturnConstant(context.getFqn(typeDecl), methodName, 0, new HashSet<>(), contextCu);
                                }
                            }
                            if (narrowed != null && !narrowed.isEmpty()) {
                                constants.addAll(narrowed);
                            }
                        }
                    }
                }
            }
            
            Expression currentMi = mi;
            while (currentMi instanceof MethodInvocation chainMi) {
                String chainName = chainMi.getName().getIdentifier();
                if (chainName.equals("setHeader") && chainMi.arguments().size() == 2) {
                    extractConstantsFromExpression((Expression) chainMi.arguments().get(1), constants);
                } else if ((chainName.equalsIgnoreCase("event") || chainName.equalsIgnoreCase("type") || chainName.equalsIgnoreCase("eventType")) && chainMi.arguments().size() == 1) {
                    extractConstantsFromExpression((Expression) chainMi.arguments().get(0), constants);
                }
                currentMi = chainMi.getExpression();
            }
            
            TypeDeclaration td = findEnclosingType(mi);
            if (td != null && (mi.getExpression() == null || mi.getExpression() instanceof ThisExpression)) {
                List<String> values = resolveMethodReturnConstant(context.getFqn(td), methodName, 0, new HashSet<>(), null);
                if (values != null) constants.addAll(values);
            }
        }
    }

    public void extractConstantsFromArgument(Expression argObj, List<String> constants) {
        if (argObj instanceof ClassInstanceCreation innerCic) {
            String typeName = click.kamil.springstatemachineexporter.ast.common.AstUtils.extractSimpleTypeName(innerCic.getType());
            if ("String".equals(typeName)) {
                extractConstantsFromExpression(argObj, constants);
            }
        } else {
            extractConstantsFromExpression(argObj, constants);
        }
    }

    public List<String> resolveClassConstantReturns(String className, CodebaseContext context, CompilationUnit contextCu) {
        if (className != null && className.contains(".")) {
            String prefix = className.substring(0, className.lastIndexOf('.'));
            String suffix = className.substring(className.lastIndexOf('.') + 1);
            TypeDeclaration prefixTd = contextCu != null ? context.getTypeDeclaration(prefix, contextCu) : context.getTypeDeclaration(prefix);
            if (prefixTd == null) prefixTd = context.getTypeDeclaration(prefix);
            boolean isKnownType = prefixTd != null || context.getEnumValues(prefix) != null;
            if (isKnownType) {
                String simplePrefix = prefix.contains(".") ? prefix.substring(prefix.lastIndexOf('.') + 1) : prefix;
                return Collections.singletonList(simplePrefix + "." + suffix);
            }
        }

        TypeDeclaration td = contextCu != null ? context.getTypeDeclaration(className, contextCu) : context.getTypeDeclaration(className);
        if (td == null) td = context.getTypeDeclaration(className);
        if (td == null) return null;
        
        final List<String> resolvedConstants = new ArrayList<>();
        for (MethodDeclaration md : td.getMethods()) {
            if (md.getBody() != null) {
                md.getBody().accept(new ASTVisitor() {
                    public boolean visit(ReturnStatement node) {
                        if (variableTracer != null) {
                            List<Expression> tracedReturns = variableTracer.traceVariableAll(node.getExpression());
                            for (Expression tracedRet : tracedReturns) {
                                extractConstantsFromExpression(tracedRet, resolvedConstants);
                            }
                        }
                        return super.visit(node);
                    }
                });
            }
            if (!resolvedConstants.isEmpty()) {
                return resolvedConstants;
            }
        }
        return null;
    }

    public List<String> resolveMethodReturnConstant(String className, String methodName, int depth, Set<String> visited, CompilationUnit contextCu) {
        log.debug("ConstantExtractor.resolveMethodReturnConstant CALLED: className={}, methodName={}", className, methodName);
        log.debug("Entering resolveMethodReturnConstant with className={}, methodName={}, depth={}", className, methodName, depth);
        if (depth > 20) {
            log.debug("Exiting resolveMethodReturnConstant early (depth limit) for {}.{}", className, methodName);
            return Collections.emptyList();
        }
        String fqn = className + "." + methodName;
        if (!visited.add(fqn)) {
            log.debug("Exiting resolveMethodReturnConstant early (cycle detected) for fqn={}", fqn);
            return Collections.emptyList();
        }

        List<String> constants = new ArrayList<>();
        TypeDeclaration tempTd = contextCu != null ? context.getTypeDeclaration(className, contextCu) : null;
        if (tempTd == null) {
            tempTd = context.getTypeDeclaration(className);
        }
        final TypeDeclaration td = tempTd;
        if (td != null) {
            MethodDeclaration md = context.findMethodDeclaration(td, methodName, true);
            if (md != null && md.getBody() != null) {
                md.getBody().accept(new ASTVisitor() {
                    @Override
                    public boolean visit(ReturnStatement node) {
                        Expression retExpr = node.getExpression();
                        if (retExpr != null) {
                            boolean handled = false;
                            if (retExpr instanceof MethodInvocation mi && calledMethodResolver != null) {
                                String called = calledMethodResolver.apply(mi);
                                if (called != null && called.contains(".")) {
                                    if (visited.contains(called)) {
                                        handled = true;
                                    } else {
                                        String cName = called.substring(0, called.lastIndexOf('.'));
                                        String mName = called.substring(called.lastIndexOf('.') + 1);
                                        
                                        TypeDeclaration targetTd = contextCu != null ? context.getTypeDeclaration(cName, contextCu) : context.getTypeDeclaration(cName);
                                        if (targetTd == null) targetTd = context.getTypeDeclaration(cName);
                                        
                                        if (targetTd != null && context.findMethodDeclaration(targetTd, mName, true) != null) {
                                            List<String> delegationResult = resolveMethodReturnConstant(cName, mName, depth + 1, visited, contextCu);
                                            constants.addAll(delegationResult);
                                            handled = true;
                                        }
                                    }
                                }
                            } else if (retExpr instanceof SuperMethodInvocation smi) {
                                String superFqn = context.getSuperclassFqn(td);
                                if (superFqn != null) {
                                    String called = superFqn + "." + smi.getName().getIdentifier();
                                    if (visited.contains(called)) {
                                        handled = true;
                                    } else {
                                        String cName = superFqn;
                                        String mName = smi.getName().getIdentifier();
                                        TypeDeclaration targetTd = contextCu != null ? context.getTypeDeclaration(cName, contextCu) : context.getTypeDeclaration(cName);
                                        if (targetTd == null) targetTd = context.getTypeDeclaration(cName);
                                        
                                        if (targetTd != null && context.findMethodDeclaration(targetTd, mName, true) != null) {
                                            List<String> delegationResult = resolveMethodReturnConstant(cName, mName, depth + 1, visited, contextCu);
                                            constants.addAll(delegationResult);
                                            handled = true;
                                        }
                                    }
                                }
                            }
                            if (!handled && variableTracer != null && constructorAnalyzer != null) {
                                List<Expression> tracedRetExprs = variableTracer.traceVariableAll(retExpr);
                                if (tracedRetExprs.isEmpty() && retExpr != null) {
                                    tracedRetExprs = List.of(retExpr);
                                }
                                for (Expression tracedRetExpr : tracedRetExprs) {
                                    extractConstantsFromExpression(tracedRetExpr, constants);
                                    String val = constantResolver.resolve(tracedRetExpr, context);
                                    if (val != null) {
                                        if (val.startsWith("ENUM_SET:")) {
                                            for (String eVal : val.substring(9).split(",")) {
                                                String parsed = parseEnumSetElement(eVal);
                                                if (!constants.contains(parsed)) constants.add(parsed);
                                            }
                                        } else {
                                            if (!constants.contains(val)) constants.add(val);
                                        }
                                    } else if (tracedRetExpr instanceof SimpleName sn) {
                                        List<String> consts = constructorAnalyzer.traceFieldInConstructors(td, sn.getIdentifier(), context, visited);
                                        if (!consts.isEmpty()) {
                                            constants.addAll(consts);
                                        }
                                    } else if (tracedRetExpr instanceof FieldAccess fa) {
                                        List<String> consts = constructorAnalyzer.traceFieldInConstructors(td, fa.getName().getIdentifier(), context, visited);
                                        if (!consts.isEmpty()) {
                                            constants.addAll(consts);
                                        }
                                    }
                                }
                            }
                        }
                        return super.visit(node);
                    }
                });
            } else {
                List<String> impls = context.getImplementations(className);
                if (impls != null) {
                    for (String implName : impls) {
                        List<String> delegationResult = resolveMethodReturnConstant(implName, methodName, depth + 1, visited, contextCu);
                        constants.addAll(delegationResult);
                    }
                }
            }
        }
        visited.remove(fqn);
        log.debug("resolveMethodReturnConstant className={} methodName={} returns: {}", className, methodName, constants);
        log.debug("resolveMethodReturnConstant for class={}, method={} resolved constants: {}", className, methodName, constants);
        return constants;
    }

    public String parseEnumSetElement(String eVal) {
        return eVal.contains(".") ? eVal.substring(eVal.lastIndexOf('.', eVal.lastIndexOf('.') - 1) + 1) : eVal;
    }

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

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