package click.kamil.springstatemachineexporter.analysis.pipeline;

import org.eclipse.jdt.core.dom.ASTNode;
import org.eclipse.jdt.core.dom.Block;
import org.eclipse.jdt.core.dom.Expression;
import org.eclipse.jdt.core.dom.ExpressionStatement;
import org.eclipse.jdt.core.dom.ReturnStatement;
import org.eclipse.jdt.core.dom.StringLiteral;
import org.eclipse.jdt.core.dom.SwitchCase;
import org.eclipse.jdt.core.dom.SwitchExpression;
import org.eclipse.jdt.core.dom.SwitchStatement;
import org.eclipse.jdt.core.dom.ThrowStatement;
import org.eclipse.jdt.core.dom.YieldStatement;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

/**
 * Shared walk of switch statement / expression arms (arrow, yield, return, default blocks).
 * Used by event-identity const maps, {@code ConstantExtractor}, and switch poly collection.
 */
public final class SwitchConstSupport {

    private SwitchConstSupport() {
    }

    public static List<?> statementsOf(ASTNode switchNode) {
        if (switchNode instanceof SwitchExpression se) {
            return se.statements();
        }
        if (switchNode instanceof SwitchStatement ss) {
            return ss.statements();
        }
        return List.of();
    }

    /**
     * Payload expressions for every arm, including {@code default} block yields.
     * Opaque / empty arms (throw-only, empty blocks) are omitted — use
     * {@link #hasOpaqueArm(List)} for fail-closed unique-agreement checks.
     */
    public static List<Expression> armPayloadExpressions(List<?> statements) {
        List<Expression> payloads = new ArrayList<>();
        if (statements == null) {
            return payloads;
        }
        boolean inArm = false;
        for (Object stmtObj : statements) {
            if (stmtObj instanceof SwitchCase) {
                inArm = true;
                continue;
            }
            if (!inArm) {
                continue;
            }
            Expression armExpr = directArmExpression(stmtObj);
            if (armExpr != null) {
                payloads.add(armExpr);
                inArm = false;
                continue;
            }
            Expression fromBlock = blockArmExpression(stmtObj);
            if (fromBlock != null) {
                payloads.add(fromBlock);
                inArm = false;
                continue;
            }
            if (isOpaqueArmStatement(stmtObj)) {
                inArm = false;
            }
        }
        return payloads;
    }

    public static List<Expression> armPayloadExpressions(ASTNode switchNode) {
        return armPayloadExpressions(statementsOf(switchNode));
    }

    /**
     * True when any switch arm has no resolvable payload (throw-only, empty block, etc.).
     * Callers that need fail-closed unique agreement must not treat partial payloads as unique.
     */
    public static boolean hasOpaqueArm(List<?> statements) {
        if (statements == null || statements.isEmpty()) {
            return false;
        }
        boolean inArm = false;
        boolean sawPayload = false;
        for (Object stmtObj : statements) {
            if (stmtObj instanceof SwitchCase) {
                if (inArm && !sawPayload) {
                    return true;
                }
                inArm = true;
                sawPayload = false;
                continue;
            }
            if (!inArm) {
                continue;
            }
            if (directArmExpression(stmtObj) != null || blockArmExpression(stmtObj) != null) {
                sawPayload = true;
                inArm = false;
                continue;
            }
            if (isOpaqueArmStatement(stmtObj)) {
                return true;
            }
        }
        return inArm && !sawPayload;
    }

    public static boolean hasOpaqueArm(ASTNode switchNode) {
        return hasOpaqueArm(statementsOf(switchNode));
    }

    /**
     * Non-default string/resolved case labels → arm payload expression.
     * Default arms are skipped (lookup maps must not invent a key for default).
     */
    public static Map<String, Expression> stringCasePayloads(
            List<?> statements,
            Function<Expression, String> caseKeyResolver) {
        Map<String, Expression> map = new LinkedHashMap<>();
        if (statements == null || caseKeyResolver == null) {
            return map;
        }
        String currentKey = null;
        for (Object stmtObj : statements) {
            if (stmtObj instanceof SwitchCase sc) {
                if (sc.isDefault()) {
                    currentKey = null;
                    continue;
                }
                currentKey = null;
                for (Object exprObj : sc.expressions()) {
                    if (exprObj instanceof StringLiteral sl) {
                        currentKey = sl.getLiteralValue();
                    } else if (exprObj instanceof Expression ex) {
                        String resolved = caseKeyResolver.apply(ex);
                        if (resolved != null && !resolved.isBlank()) {
                            currentKey = resolved;
                        }
                    }
                }
                continue;
            }
            if (currentKey == null) {
                continue;
            }
            Expression armExpr = directArmExpression(stmtObj);
            if (armExpr == null) {
                armExpr = blockArmExpression(stmtObj);
            }
            if (armExpr != null) {
                map.put(currentKey, armExpr);
                currentKey = null;
            }
        }
        return map;
    }

    public static Map<String, Expression> stringCasePayloads(
            ASTNode switchNode,
            Function<Expression, String> caseKeyResolver) {
        return stringCasePayloads(statementsOf(switchNode), caseKeyResolver);
    }

    private static boolean isOpaqueArmStatement(Object stmtObj) {
        if (stmtObj instanceof ThrowStatement) {
            return true;
        }
        if (stmtObj instanceof Block block) {
            if (block.statements().isEmpty()) {
                return true;
            }
            for (Object nested : block.statements()) {
                if (nested instanceof YieldStatement || nested instanceof ReturnStatement) {
                    return false;
                }
            }
            for (Object nested : block.statements()) {
                if (nested instanceof ThrowStatement) {
                    return true;
                }
            }
            return true;
        }
        return false;
    }

    private static Expression directArmExpression(Object stmtObj) {
        if (stmtObj instanceof ReturnStatement rs) {
            return rs.getExpression();
        }
        if (stmtObj instanceof YieldStatement ys) {
            return ys.getExpression();
        }
        if (stmtObj instanceof ExpressionStatement es) {
            return es.getExpression();
        }
        return null;
    }

    private static Expression blockArmExpression(Object stmtObj) {
        if (!(stmtObj instanceof Block block)) {
            return null;
        }
        for (Object nested : block.statements()) {
            if (nested instanceof YieldStatement ys) {
                return ys.getExpression();
            }
            if (nested instanceof ReturnStatement rs) {
                return rs.getExpression();
            }
        }
        return null;
    }
}
