package click.kamil.springstatemachineexporter.analysis.service;

import click.kamil.springstatemachineexporter.analysis.model.CallEdge;
import click.kamil.springstatemachineexporter.analysis.resolver.BooleanConstraintEvaluator;
import click.kamil.springstatemachineexporter.analysis.resolver.ConstantResolver;
import click.kamil.springstatemachineexporter.analysis.resolver.TruncatedEnumRefReconstructor;
import click.kamil.springstatemachineexporter.ast.common.AstUtils;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.StringJoiner;

/**
 * Walks a call path forward, propagating known parameter / local bindings so switch-case
 * edge constraints can prune impossible branches (e.g. {@code dispatch} arms).
 */
public class PathBindingEvaluator {

    private final CodebaseContext context;
    private final VariableTracer variableTracer;
    private final ConstantResolver constantResolver;
    private final TypeResolver typeResolver;
    private final PathBindingExpressionResolver expressionResolver;
    private final Map<String, Map<String, String>> traceBindingsCache = new HashMap<>();
    /** Optional machine/event enum FQN used to disambiguate truncated package+CONSTANT refs. */
    private String preferredEnumTypeFqn;

    public PathBindingEvaluator(
            CodebaseContext context,
            VariableTracer variableTracer,
            ConstantResolver constantResolver,
            TypeResolver typeResolver) {
        this(context, variableTracer, constantResolver, typeResolver, new PathBindingExpressionResolver(context, variableTracer, constantResolver));
    }

    PathBindingEvaluator(
            CodebaseContext context,
            VariableTracer variableTracer,
            ConstantResolver constantResolver,
            TypeResolver typeResolver,
            PathBindingExpressionResolver expressionResolver) {
        this.context = context;
        this.variableTracer = variableTracer;
        this.constantResolver = constantResolver;
        this.typeResolver = typeResolver;
        this.expressionResolver = expressionResolver;
    }

    public void setPreferredEnumTypeFqn(String preferredEnumTypeFqn) {
        this.preferredEnumTypeFqn = preferredEnumTypeFqn;
    }

    /** Visible for tests: bindings accumulated while walking a path forward (ignores constraint rejection). */
    Map<String, String> traceBindings(List<String> path, Map<String, List<CallEdge>> callGraph, CallGraphPathFinder pathFinder) {
        if (path == null || path.size() < 2) {
            return Map.of();
        }
        String cacheKey = pathCacheKey(path);
        Map<String, String> cached = traceBindingsCache.get(cacheKey);
        if (cached != null) {
            return new HashMap<>(cached);
        }
        Map<String, String> bindings = computeTraceBindings(path, callGraph, pathFinder);
        traceBindingsCache.put(cacheKey, Map.copyOf(bindings));
        return new HashMap<>(bindings);
    }

    public void clearAnalysisCaches() {
        traceBindingsCache.clear();
    }

    private Map<String, String> computeTraceBindings(
            List<String> path, Map<String, List<CallEdge>> callGraph, CallGraphPathFinder pathFinder) {
        context.setAnalysisCallPath(path);
        try {
            Map<String, String> bindings = new HashMap<>();
            for (int i = 0; i < path.size() - 1; i++) {
                String caller = path.get(i);
                String target = path.get(i + 1);
                CallEdge edge = findEdge(caller, target, callGraph, pathFinder);
                if (edge == null) {
                    continue;
                }
                enrichBindings(caller, target, edge, path, i + 1, callGraph, pathFinder, bindings);
            }
            return bindings;
        } finally {
            context.clearAnalysisCallPath();
        }
    }

    private static String pathCacheKey(List<String> path) {
        StringJoiner joiner = new StringJoiner("\0");
        for (String hop : path) {
            joiner.add(hop);
        }
        return joiner.toString();
    }

    public boolean isPathCompatible(
            List<String> path,
            Map<String, List<CallEdge>> callGraph,
            CallGraphPathFinder pathFinder,
            Map<String, String> initialBindings) {
        if (path == null || path.size() < 2) {
            return true;
        }
        Map<String, String> bindings = new HashMap<>();
        if (initialBindings != null) {
            bindings.putAll(initialBindings);
        }
        for (int i = 0; i < path.size() - 1; i++) {
            String caller = path.get(i);
            String target = path.get(i + 1);
            CallEdge edge = findEdge(caller, target, callGraph, pathFinder);
            if (edge == null) {
                continue;
            }
            if (edge.getConstraint() != null
                    && !edge.getConstraint().isBlank()
                    && !BooleanConstraintEvaluator.isCompatibleWithBindings(
                            edge.getConstraint(), bindings, preferredEnumTypeFqn, context)) {
                return false;
            }
            enrichBindings(caller, target, edge, path, i + 1, callGraph, pathFinder, bindings);
        }
        return true;
    }

    public boolean isPathCompatible(List<String> path, Map<String, List<CallEdge>> callGraph, CallGraphPathFinder pathFinder) {
        return isPathCompatible(path, callGraph, pathFinder, Map.of());
    }

    /**
     * Returns {@code true} when {@code caller}→{@code callee} via {@code edge} satisfies the edge constraint
     * under {@code bindings}; on success, {@code bindings} is updated with propagated parameter values.
     */
    public boolean tryTraverseEdge(
            String caller,
            String callee,
            CallEdge edge,
            List<String> pathWithCallee,
            int calleeIndex,
            Map<String, List<CallEdge>> callGraph,
            CallGraphPathFinder pathFinder,
            Map<String, String> bindings) {
        if (edge.getConstraint() != null
                && !edge.getConstraint().isBlank()
                && !BooleanConstraintEvaluator.isCompatibleWithBindings(
                        edge.getConstraint(), bindings, preferredEnumTypeFqn, context)) {
            return false;
        }
        if (variableTracer == null || typeResolver == null) {
            return true;
        }
        enrichBindings(caller, callee, edge, pathWithCallee, calleeIndex, callGraph, pathFinder, bindings);
        return true;
    }

    private void enrichBindings(
            String caller,
            String target,
            CallEdge edge,
            List<String> path,
            int pathIndex,
            Map<String, List<CallEdge>> callGraph,
            CallGraphPathFinder pathFinder,
            Map<String, String> bindings) {
        if (FunctionalInterfaceTypes.isFunctionalSamMethod(target) && pathIndex >= 2) {
            int paramIndex = variableTracer.findFunctionalParameterIndex(caller);
            if (paramIndex < 0) {
                paramIndex = 0;
            }
            String functionalCaller = path.get(pathIndex - 2);
            CallEdge functionalEdge = CallSiteMatcher.selectEdge(
                    functionalCaller, caller, callGraph.get(functionalCaller), bindings);
            List<String> constants = variableTracer.resolveConstantsFromCallSiteArgument(
                    functionalCaller, caller, paramIndex, functionalEdge);
            for (String constant : constants) {
                if (constant != null && !constant.isBlank()) {
                    bindings.putIfAbsent(constant, constant);
                }
            }
        }
        Map<String, String> paramValues = variableTracer.buildParameterValuesMap(
                caller, target, callGraph, path, pathIndex, bindings);
        mergeResolvedBindings(bindings, paramValues, caller);

        List<String> targetParams = typeResolver.getParameterNames(target);
        if (targetParams == null || edge.getArguments() == null) {
            return;
        }
        for (int i = 0; i < edge.getArguments().size() && i < targetParams.size(); i++) {
            String arg = edge.getArguments().get(i);
            String paramName = targetParams.get(i);
            if (typeResolver != null
                    && FunctionalInterfaceTypes.isFunctionalInterface(typeResolver.getParameterType(target, i))) {
                List<String> constants = variableTracer.resolveConstantsFromCallSiteArgument(caller, target, i, edge);
                if (constants.size() == 1) {
                    String constant = maybeReconstruct(constants.get(0));
                    if (shouldStoreBinding(paramName, constant)) {
                        bindings.put(paramName, maybeReconstruct(constant));
                    }
                    continue;
                }
            }
            String resolved = resolveBindingValue(caller, arg, bindings);
            if (shouldStoreBinding(paramName, resolved)) {
                bindings.put(paramName, maybeReconstruct(resolved));
            }
        }
    }

    private boolean shouldStoreBinding(String paramName, String resolved) {
        if (resolved == null || resolved.isBlank() || resolved.equals(paramName)) {
            return false;
        }
        if (AstUtils.isAnalysisToken(resolved) || !AstUtils.isConcreteConstant(resolved)) {
            return false;
        }
        if (BooleanConstraintEvaluator.isIncompleteDottedName(resolved)) {
            return false;
        }
        // Reconstruct with preferred type before rejecting truncated package+CONSTANT refs.
        String candidate = maybeReconstruct(resolved);
        if (TruncatedEnumRefReconstructor.looksLikePackageTruncatedEnumRef(candidate)) {
            return false;
        }
        if (candidate.contains("(") && !AstUtils.isEnumLikeConstantName(candidate)) {
            return false;
        }
        return true;
    }

    private void mergeResolvedBindings(Map<String, String> bindings, Map<String, String> paramValues, String caller) {
        if (paramValues == null || paramValues.isEmpty()) {
            return;
        }
        for (Map.Entry<String, String> entry : paramValues.entrySet()) {
            String resolved = resolveBindingValue(caller, entry.getValue(), bindings);
            if (shouldStoreBinding(entry.getKey(), resolved)) {
                bindings.put(entry.getKey(), maybeReconstruct(resolved));
            }
        }
    }

    private String resolveBindingValue(String callerFqn, String rawValue, Map<String, String> bindings) {
        if (rawValue == null || rawValue.isBlank()) {
            return null;
        }
        if (AstUtils.isAnalysisToken(rawValue)) {
            return null;
        }
        if (rawValue.startsWith("\"") || AstUtils.isEnumLikeConstantName(rawValue)) {
            return maybeReconstruct(normalizeLiteral(rawValue));
        }
        if (bindings.containsKey(rawValue)) {
            String bound = bindings.get(rawValue);
            // Analysis tokens are not concrete proofs — fall through to re-resolve under path bindings.
            if (bound != null && !AstUtils.isAnalysisToken(bound)) {
                return maybeReconstruct(bound);
            }
        }
        String traced = variableTracer.traceLocalVariable(callerFqn, rawValue, bindings);
        if (traced != null && !traced.isBlank()) {
            if (AstUtils.isAnalysisToken(traced)) {
                String fromCall = expressionResolver.resolveMethodCallFromSource(callerFqn, rawValue, bindings);
                if (fromCall != null && !AstUtils.isAnalysisToken(fromCall)) {
                    return maybeReconstruct(fromCall);
                }
                return null;
            }
            if (traced.startsWith("\"") || AstUtils.isEnumLikeConstantName(traced)) {
                return maybeReconstruct(traced);
            }
            String fromCall = expressionResolver.resolveMethodCallFromSource(callerFqn, rawValue, bindings);
            if (fromCall != null && !AstUtils.isAnalysisToken(fromCall)) {
                return maybeReconstruct(fromCall);
            }
            if (!traced.equals(rawValue)) {
                return maybeReconstruct(traced);
            }
        }
        return maybeReconstruct(rawValue);
    }

    private String maybeReconstruct(String value) {
        if (value == null || context == null) {
            return value;
        }
        return TruncatedEnumRefReconstructor.reconstruct(value, preferredEnumTypeFqn, context);
    }

    private static String normalizeLiteral(String value) {
        if (value != null && value.startsWith("\"") && value.endsWith("\"")) {
            return value.substring(1, value.length() - 1);
        }
        return value;
    }

    private static CallEdge findEdge(
            String caller, String target, 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(target) || pathFinder.isHeuristicMatch(edge.getTargetMethod(), target)) {
                return edge;
            }
        }
        return null;
    }
}
