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.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.*;

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 Map<String, Map<String, String>> traceBindingsCache = new HashMap<>();

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

    /** 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) {
        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;
    }

    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)) {
                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)) {
            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) {
        Map<String, String> paramValues = variableTracer.buildParameterValuesMap(caller, target, callGraph, path, pathIndex);
        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);
            String resolved = resolveBindingValue(caller, arg, bindings);
            if (shouldStoreBinding(paramName, resolved)) {
                bindings.put(paramName, resolved);
            }
        }
    }

    private boolean shouldStoreBinding(String paramName, String resolved) {
        if (resolved == null || resolved.isBlank() || resolved.equals(paramName)) {
            return false;
        }
        if (resolved.contains("(") && !isEnumLikeConstant(resolved)) {
            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(), resolved);
            }
        }
    }

    private String resolveBindingValue(String callerFqn, String rawValue, Map<String, String> bindings) {
        if (rawValue == null || rawValue.isBlank()) {
            return null;
        }
        if (rawValue.startsWith("\"") || isEnumLikeConstant(rawValue)) {
            return normalizeLiteral(rawValue);
        }
        if (bindings.containsKey(rawValue)) {
            return bindings.get(rawValue);
        }
        String traced = variableTracer.traceLocalVariable(callerFqn, rawValue, bindings);
        if (traced != null && !traced.isBlank()) {
            if (traced.startsWith("\"") || isEnumLikeConstant(traced)) {
                return traced;
            }
            String fromCall = resolveMethodCallFromSource(callerFqn, rawValue, bindings);
            if (fromCall != null) {
                return fromCall;
            }
            if (!traced.equals(rawValue)) {
                return traced;
            }
        }
        return rawValue;
    }

    private String resolveMethodCallFromSource(String callerFqn, String varName, Map<String, String> bindings) {
        MethodDeclaration md = findMethodDeclaration(callerFqn);
        if (md == null || md.getBody() == null) {
            return null;
        }
        final Expression[] initializer = new Expression[1];
        md.getBody().accept(new ASTVisitor() {
            @Override
            public boolean visit(VariableDeclarationFragment node) {
                if (node.getName().getIdentifier().equals(varName) && node.getInitializer() != null) {
                    initializer[0] = node.getInitializer();
                }
                return super.visit(node);
            }
        });
        if (initializer[0] == null) {
            return null;
        }
        return evaluateExpressionValue(initializer[0], callerFqn, bindings);
    }

    private String evaluateExpressionValue(Expression expr, String callerFqn, Map<String, String> bindings) {
        if (expr instanceof StringLiteral sl) {
            return "\"" + sl.getLiteralValue() + "\"";
        }
        if (expr instanceof QualifiedName qn) {
            return qn.getFullyQualifiedName();
        }
        if (expr instanceof SimpleName sn) {
            return bindings.getOrDefault(sn.getIdentifier(), sn.getIdentifier());
        }
        if (expr instanceof MethodInvocation mi) {
            return evaluateMethodInvocationReturn(mi, callerFqn, bindings);
        }
        if (expr instanceof SwitchExpression se) {
            return constantResolver.evaluateSwitchWithParams(se, bindings, context);
        }
        String resolved = constantResolver.resolve(expr, context);
        return resolved;
    }

    private String evaluateMethodInvocationReturn(MethodInvocation mi, String callerFqn, Map<String, String> bindings) {
        String ownerFqn = resolveMethodOwnerFqn(mi, callerFqn);
        if (ownerFqn == null) {
            return null;
        }
        String methodName = mi.getName().getIdentifier();
        TypeDeclaration owner = context.getTypeDeclaration(ownerFqn);
        if (owner == null) {
            return null;
        }
        MethodDeclaration md = context.findMethodDeclaration(owner, methodName, true);
        if (md == null || md.getBody() == null) {
            return null;
        }

        Map<String, String> callBindings = new HashMap<>(bindings);
        for (int i = 0; i < md.parameters().size() && i < mi.arguments().size(); i++) {
            SingleVariableDeclaration param = (SingleVariableDeclaration) md.parameters().get(i);
            String paramName = param.getName().getIdentifier();
            String argValue = evaluateExpressionValue((Expression) mi.arguments().get(i), callerFqn, callBindings);
            if (argValue != null) {
                callBindings.put(paramName, argValue);
            }
        }

        final String[] result = new String[1];
        md.getBody().accept(new ASTVisitor() {
            @Override
            public boolean visit(ReturnStatement node) {
                if (node.getExpression() instanceof SwitchExpression se) {
                    result[0] = constantResolver.evaluateSwitchWithParams(se, callBindings, context);
                } else if (node.getExpression() != null) {
                    result[0] = evaluateExpressionValue(node.getExpression(), ownerFqn + "." + methodName, callBindings);
                }
                return super.visit(node);
            }
        });
        return result[0];
    }

    private String resolveFieldTypeFqn(String callerFqn, String fieldName) {
        if (callerFqn == null || !callerFqn.contains(".")) {
            return null;
        }
        String className = callerFqn.substring(0, callerFqn.lastIndexOf('.'));
        TypeDeclaration td = context.getTypeDeclaration(className);
        if (td == null) {
            return null;
        }
        for (FieldDeclaration field : td.getFields()) {
            for (Object fragmentObj : field.fragments()) {
                VariableDeclarationFragment fragment = (VariableDeclarationFragment) fragmentObj;
                if (fragment.getName().getIdentifier().equals(fieldName)) {
                    IVariableBinding binding = fragment.resolveBinding();
                    if (binding != null && binding.getType() != null) {
                        return binding.getType().getErasure().getQualifiedName();
                    }
                    String simpleType = field.getType().toString();
                    TypeDeclaration resolved = context.getTypeDeclaration(simpleType);
                    if (resolved != null) {
                        return context.getFqn(resolved);
                    }
                    if (td.getRoot() instanceof CompilationUnit cu && cu.getPackage() != null) {
                        String pkg = cu.getPackage().getName().getFullyQualifiedName();
                        resolved = context.getTypeDeclaration(pkg + "." + simpleType);
                        if (resolved != null) {
                            return context.getFqn(resolved);
                        }
                    }
                    return simpleType;
                }
            }
        }
        return null;
    }

    private String resolveMethodOwnerFqn(MethodInvocation mi, String callerFqn) {
        IMethodBinding methodBinding = mi.resolveMethodBinding();
        if (methodBinding != null && methodBinding.getDeclaringClass() != null) {
            return methodBinding.getDeclaringClass().getErasure().getQualifiedName();
        }
        Expression receiver = mi.getExpression();
        if (receiver instanceof SimpleName sn) {
            return resolveFieldTypeFqn(callerFqn, sn.getIdentifier());
        }
        if (receiver instanceof FieldAccess fa) {
            IVariableBinding fieldBinding = fa.resolveFieldBinding();
            if (fieldBinding != null && fieldBinding.getType() != null) {
                return fieldBinding.getType().getErasure().getQualifiedName();
            }
            return resolveFieldTypeFqn(callerFqn, fa.getName().getIdentifier());
        }
        return null;
    }

    private MethodDeclaration findMethodDeclaration(String methodFqn) {
        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) {
            return null;
        }
        return context.findMethodDeclaration(td, methodName, true);
    }

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

    private static boolean isEnumLikeConstant(String value) {
        if (value == null || value.isEmpty()) {
            return false;
        }
        int dot = value.lastIndexOf('.');
        if (dot <= 0 || dot >= value.length() - 1) {
            return false;
        }
        return Character.isUpperCase(value.charAt(dot + 1));
    }

    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;
    }
}
