package click.kamil.springstatemachineexporter.ast.common;

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

public class CfgBuilder {

    public static ControlFlowGraph build(MethodDeclaration method) {
        ControlFlowGraph cfg = new ControlFlowGraph(method);
        if (method.getBody() == null) {
            cfg.getEntryNode().addSuccessor(cfg.getExitNode());
            return cfg;
        }

        // Build CFG recursively
        List<ControlFlowGraph.CfgNode> exits = buildFlow(method.getBody(), List.of(cfg.getEntryNode()), cfg);
        for (ControlFlowGraph.CfgNode exit : exits) {
            exit.addSuccessor(cfg.getExitNode());
        }

        return cfg;
    }

    private static List<ControlFlowGraph.CfgNode> buildFlow(
            ASTNode node, 
            List<ControlFlowGraph.CfgNode> sources, 
            ControlFlowGraph cfg) {
        
        if (node == null || sources.isEmpty()) {
            return sources;
        }

        if (node instanceof Block block) {
            List<ControlFlowGraph.CfgNode> currentSources = sources;
            for (Object stmtObj : block.statements()) {
                currentSources = buildFlow((ASTNode) stmtObj, currentSources, cfg);
            }
            return currentSources;
        }

        if (node instanceof IfStatement ifStmt) {
            // Condition node
            ControlFlowGraph.CfgNode condNode = new ControlFlowGraph.CfgNode(ifStmt.getExpression(), ControlFlowGraph.CfgNodeType.STATEMENT);
            cfg.addNode(condNode);
            for (ControlFlowGraph.CfgNode src : sources) {
                src.addSuccessor(condNode);
            }

            List<ControlFlowGraph.CfgNode> thenExits = buildFlow(ifStmt.getThenStatement(), List.of(condNode), cfg);
            List<ControlFlowGraph.CfgNode> elseExits;
            if (ifStmt.getElseStatement() != null) {
                elseExits = buildFlow(ifStmt.getElseStatement(), List.of(condNode), cfg);
            } else {
                elseExits = List.of(condNode);
            }

            List<ControlFlowGraph.CfgNode> combinedExits = new ArrayList<>();
            combinedExits.addAll(thenExits);
            combinedExits.addAll(elseExits);
            return combinedExits;
        }

        if (node instanceof WhileStatement whileStmt) {
            ControlFlowGraph.CfgNode condNode = new ControlFlowGraph.CfgNode(whileStmt.getExpression(), ControlFlowGraph.CfgNodeType.STATEMENT);
            cfg.addNode(condNode);
            for (ControlFlowGraph.CfgNode src : sources) {
                src.addSuccessor(condNode);
            }

            List<ControlFlowGraph.CfgNode> bodyExits = buildFlow(whileStmt.getBody(), List.of(condNode), cfg);
            for (ControlFlowGraph.CfgNode exit : bodyExits) {
                exit.addSuccessor(condNode); // Loop back
            }

            return List.of(condNode); // Exit loop when cond is false
        }

        if (node instanceof DoStatement doStmt) {
            ControlFlowGraph.CfgNode condNode = new ControlFlowGraph.CfgNode(doStmt.getExpression(), ControlFlowGraph.CfgNodeType.STATEMENT);
            cfg.addNode(condNode);

            List<ControlFlowGraph.CfgNode> bodySources = new ArrayList<>(sources);
            bodySources.add(condNode);

            List<ControlFlowGraph.CfgNode> bodyExits = buildFlow(doStmt.getBody(), bodySources, cfg);
            for (ControlFlowGraph.CfgNode exit : bodyExits) {
                exit.addSuccessor(condNode);
            }
            return List.of(condNode);
        }

        if (node instanceof ForStatement forStmt) {
            // Initializers
            List<ControlFlowGraph.CfgNode> currentSources = sources;
            for (Object initObj : forStmt.initializers()) {
                ControlFlowGraph.CfgNode initNode = new ControlFlowGraph.CfgNode((ASTNode) initObj, ControlFlowGraph.CfgNodeType.STATEMENT);
                cfg.addNode(initNode);
                for (ControlFlowGraph.CfgNode src : currentSources) {
                    src.addSuccessor(initNode);
                }
                currentSources = List.of(initNode);
            }

            // Condition
            ControlFlowGraph.CfgNode condNode;
            if (forStmt.getExpression() != null) {
                condNode = new ControlFlowGraph.CfgNode(forStmt.getExpression(), ControlFlowGraph.CfgNodeType.STATEMENT);
            } else {
                condNode = new ControlFlowGraph.CfgNode(forStmt, ControlFlowGraph.CfgNodeType.STATEMENT); // Dummy node
            }
            cfg.addNode(condNode);
            for (ControlFlowGraph.CfgNode src : currentSources) {
                src.addSuccessor(condNode);
            }

            List<ControlFlowGraph.CfgNode> bodyExits = buildFlow(forStmt.getBody(), List.of(condNode), cfg);

            // Updaters
            List<ControlFlowGraph.CfgNode> updaterSources = bodyExits;
            for (Object updateObj : forStmt.updaters()) {
                ControlFlowGraph.CfgNode updateNode = new ControlFlowGraph.CfgNode((ASTNode) updateObj, ControlFlowGraph.CfgNodeType.STATEMENT);
                cfg.addNode(updateNode);
                for (ControlFlowGraph.CfgNode src : updaterSources) {
                    src.addSuccessor(updateNode);
                }
                updaterSources = List.of(updateNode);
            }

            for (ControlFlowGraph.CfgNode src : updaterSources) {
                src.addSuccessor(condNode); // Loop back
            }

            return List.of(condNode);
        }

        if (node instanceof EnhancedForStatement effStmt) {
            ControlFlowGraph.CfgNode loopNode = new ControlFlowGraph.CfgNode(effStmt.getParameter(), ControlFlowGraph.CfgNodeType.STATEMENT);
            cfg.addNode(loopNode);
            for (ControlFlowGraph.CfgNode src : sources) {
                src.addSuccessor(loopNode);
            }

            List<ControlFlowGraph.CfgNode> bodyExits = buildFlow(effStmt.getBody(), List.of(loopNode), cfg);
            for (ControlFlowGraph.CfgNode exit : bodyExits) {
                exit.addSuccessor(loopNode); // Loop back
            }

            return List.of(loopNode);
        }

        if (node instanceof ReturnStatement returnStmt) {
            ControlFlowGraph.CfgNode retNode = new ControlFlowGraph.CfgNode(returnStmt, ControlFlowGraph.CfgNodeType.STATEMENT);
            cfg.addNode(retNode);
            for (ControlFlowGraph.CfgNode src : sources) {
                src.addSuccessor(retNode);
            }
            // Return flows to exit, not to subsequent sequential nodes
            retNode.addSuccessor(cfg.getExitNode());
            return Collections.emptyList();
        }

        // Default: general statement
        ControlFlowGraph.CfgNode stmtNode = new ControlFlowGraph.CfgNode(node, ControlFlowGraph.CfgNodeType.STATEMENT);
        cfg.addNode(stmtNode);
        for (ControlFlowGraph.CfgNode src : sources) {
            src.addSuccessor(stmtNode);
        }
        return List.of(stmtNode);
    }
}
