package click.kamil.springstatemachineexporter.analysis.service;

import click.kamil.springstatemachineexporter.analysis.model.CallEdge;
import org.eclipse.jdt.core.dom.ASTNode;
import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.MethodInvocation;
import org.eclipse.jdt.core.dom.SimpleName;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * Resolves call-site arguments across stack frames (e.g. {@code Supplier.get()} → caller lambda).
 */
public final class CallSiteArgumentResolver {

    public interface ExpressionHooks {
        /** Unwrap lambda / enum literal argument expressions to a stable string form. */
        String resolveArgument(Expression expr);

        /** Parse a stringified expression back to AST for constant extraction. */
        Expression parseExpression(String expr);
    }

    private CallSiteArgumentResolver() {
    }

    /**
     * When {@code supplierParam.get()} is seen inside {@code scopeMethod}, walk to the caller frame
     * and extract enum/constants from the argument bound to {@code supplierParam}.
     */
    public static List<String> resolveFunctionalGetFromCallerFrame(
            MethodInvocation getCall,
            List<String> path,
            Map<String, List<CallEdge>> callGraph,
            String scopeMethod,
            TypeResolver typeResolver,
            CallGraphPathFinder pathFinder,
            ConstantExtractor constantExtractor,
            ExpressionHooks hooks) {
        if (path == null || path.size() < 2 || scopeMethod == null || getCall.getExpression() == null) {
            return List.of();
        }
        if (!(getCall.getExpression() instanceof SimpleName supplierParam)) {
            return List.of();
        }
        int scopeIndex = indexOfMethod(path, scopeMethod);
        if (scopeIndex <= 0) {
            return List.of();
        }
        String callee = path.get(scopeIndex);
        String caller = path.get(scopeIndex - 1);
        int paramIndex = typeResolver.getParameterIndex(callee, supplierParam.getIdentifier());
        if (paramIndex < 0) {
            return List.of();
        }
        String argValue = findCallSiteArgument(caller, callee, paramIndex, callGraph, pathFinder);
        if (argValue == null) {
            return List.of();
        }
        return extractConstantsFromCallSiteArgument(argValue, constantExtractor, hooks);
    }

    public static String findCallSiteArgument(
            String caller,
            String callee,
            int paramIndex,
            Map<String, List<CallEdge>> callGraph,
            CallGraphPathFinder pathFinder) {
        List<CallEdge> edges = callGraph.get(caller);
        if (edges == null) {
            return null;
        }
        for (CallEdge edge : edges) {
            if (!edge.getTargetMethod().equals(callee) && !pathFinder.isHeuristicMatch(edge.getTargetMethod(), callee)) {
                continue;
            }
            if (paramIndex < edge.getArguments().size()) {
                return edge.getArguments().get(paramIndex);
            }
        }
        return null;
    }

    public static int indexOfMethod(List<String> path, String methodFqn) {
        for (int i = path.size() - 1; i >= 0; i--) {
            if (methodFqn.equals(path.get(i))) {
                return i;
            }
        }
        return -1;
    }

    public static List<String> extractConstantsFromCallSiteArgument(
            String argValue,
            ConstantExtractor constantExtractor,
            ExpressionHooks hooks) {
        if (argValue == null || argValue.isBlank()) {
            return List.of();
        }
        Expression parsed = hooks.parseExpression(argValue);
        if (parsed == null) {
            if (FunctionalInterfaceTypes.looksLikeEnumConstant(argValue)) {
                return List.of(argValue);
            }
            return List.of();
        }
        String resolved = hooks.resolveArgument(parsed);
        Expression resolvedExpr = hooks.parseExpression(resolved);
        List<String> constants = new ArrayList<>();
        if (resolvedExpr != null) {
            constantExtractor.extractConstantsFromExpression(resolvedExpr, constants);
        }
        if (constants.isEmpty() && FunctionalInterfaceTypes.looksLikeEnumConstant(resolved)) {
            constants.add(resolved);
        }
        return constants;
    }
}
