package click.kamil.springstatemachineexporter.ast.common;

import org.eclipse.jdt.core.dom.*;
import java.util.*;

public class ReachingDefinitions {
    private final ControlFlowGraph cfg;
    private final Map<ControlFlowGraph.CfgNode, Set<ASTNode>> inMap = new HashMap<>();
    private final Map<ControlFlowGraph.CfgNode, Set<ASTNode>> outMap = new HashMap<>();

    public ReachingDefinitions(ControlFlowGraph cfg) {
        this.cfg = cfg;
        analyze();
    }

    private void analyze() {
        List<ControlFlowGraph.CfgNode> cfgNodes = cfg.getNodes();
        Map<ControlFlowGraph.CfgNode, ASTNode> nodeDefs = new HashMap<>();
        Map<String, Set<ASTNode>> varToDefs = new HashMap<>();

        for (ControlFlowGraph.CfgNode node : cfgNodes) {
            ASTNode ast = node.getAstNode();
            if (ast == null) continue;

            ASTNode def = findDefinition(ast);
            if (def != null) {
                nodeDefs.put(node, def);
                String varName = getDefinedVariableName(def);
                if (varName != null) {
                    varToDefs.computeIfAbsent(varName, k -> new HashSet<>()).add(def);
                }
            }
        }

        for (ControlFlowGraph.CfgNode node : cfgNodes) {
            inMap.put(node, new HashSet<>());
            outMap.put(node, new HashSet<>());
        }

        boolean changed = true;
        int iterations = 0;
        int maxIterations = 5000;
        while (changed && iterations < maxIterations) {
            iterations++;
            changed = false;
            for (ControlFlowGraph.CfgNode node : cfgNodes) {
                Set<ASTNode> newIn = new HashSet<>();
                for (ControlFlowGraph.CfgNode pred : node.getPredecessors()) {
                    newIn.addAll(outMap.get(pred));
                }

                inMap.put(node, newIn);

                Set<ASTNode> newOut = new HashSet<>(newIn);
                ASTNode def = nodeDefs.get(node);
                if (def != null) {
                    String varName = getDefinedVariableName(def);
                    if (varName != null) {
                        Set<ASTNode> allDefs = varToDefs.get(varName);
                        if (allDefs != null) {
                            newOut.removeAll(allDefs);
                        }
                    }
                    newOut.add(def);
                }

                Set<ASTNode> oldOut = outMap.put(node, newOut);
                if (!newOut.equals(oldOut)) {
                    changed = true;
                }
            }
        }
        if (iterations >= maxIterations) {
            // Log warning or print to stderr
            System.err.println("Warning: ReachingDefinitions analysis reached max iterations (" + maxIterations + ") and terminated early.");
        }
    }

    public Set<ASTNode> getReachingDefinitions(ASTNode usageNode, String varName) {
        ControlFlowGraph.CfgNode cfgNode = findCfgNodeForUsage(usageNode);
        if (cfgNode == null) {
            return Collections.emptySet();
        }

        Set<ASTNode> reaching = new HashSet<>();
        Set<ASTNode> inDefs = inMap.get(cfgNode);
        if (inDefs != null) {
            for (ASTNode def : inDefs) {
                if (varName.equals(getDefinedVariableName(def))) {
                    reaching.add(def);
                }
            }
        }
        return reaching;
    }

    private ControlFlowGraph.CfgNode findCfgNodeForUsage(ASTNode usageNode) {
        ASTNode current = usageNode;
        while (current != null) {
            for (ControlFlowGraph.CfgNode node : cfg.getNodes()) {
                if (node.getAstNode() == current) {
                    return node;
                }
            }
            current = current.getParent();
        }
        return null;
    }

    private ASTNode findDefinition(ASTNode ast) {
        if (ast instanceof VariableDeclarationStatement vds) {
            if (!vds.fragments().isEmpty()) {
                return (ASTNode) vds.fragments().get(0);
            }
        }
        if (ast instanceof VariableDeclarationFragment) {
            return ast;
        }
        if (ast instanceof SingleVariableDeclaration) {
            return ast;
        }
        if (ast instanceof Assignment) {
            return ast;
        }
        if (ast instanceof ExpressionStatement es) {
            return findDefinition(es.getExpression());
        }
        return null;
    }

    public static String getDefinedVariableName(ASTNode def) {
        if (def instanceof VariableDeclarationFragment frag) {
            return frag.getName().getIdentifier();
        }
        if (def instanceof SingleVariableDeclaration svd) {
            return svd.getName().getIdentifier();
        }
        if (def instanceof Assignment assign) {
            if (assign.getLeftHandSide() instanceof SimpleName sn) {
                return sn.getIdentifier();
            }
        }
        return null;
    }

    public static Expression getDefinitionExpression(ASTNode def) {
        if (def instanceof VariableDeclarationFragment frag) {
            return frag.getInitializer();
        }
        if (def instanceof Assignment assign) {
            return assign.getRightHandSide();
        }
        return null;
    }
}
