package p.studio.compiler.pbs.lowering;

import p.studio.compiler.models.IRBackendExecutableFunction;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.*;

final class PbsExecutableStackAnalyzer {

    int analyzeMaxStackSlots(
            final ReadOnlyList<IRBackendExecutableFunction.Instruction> instructions,
            final int returnSlots) {
        if (instructions.isEmpty()) {
            return Math.max(1, returnSlots);
        }

        final var labelToIndex = new HashMap<String, Integer>();
        for (var i = 0; i < instructions.size(); i++) {
            final var instruction = instructions.get(i);
            if (instruction.kind() == IRBackendExecutableFunction.InstructionKind.LABEL) {
                if (labelToIndex.putIfAbsent(instruction.label(), i) != null) {
                    throw new ExecutableLoweringAnalysisException("duplicate label during stack analysis: " + instruction.label());
                }
            }
        }

        final var incomingHeightByInstruction = new HashMap<Integer, Integer>();
        final var worklist = new ArrayDeque<Integer>();
        final var visited = new HashSet<Integer>();
        incomingHeightByInstruction.put(0, 0);
        worklist.add(0);

        var maxHeight = 0;
        while (!worklist.isEmpty()) {
            final var index = worklist.removeFirst();
            visited.add(index);
            final var instruction = instructions.get(index);
            var outHeight = incomingHeightByInstruction.get(index);

            switch (instruction.kind()) {
                case PUSH_I32, PUSH_BOOL, PUSH_CONST, GET_LOCAL, GET_GLOBAL -> outHeight += 1;
                case POP, SET_LOCAL, SET_GLOBAL -> {
                    if (outHeight <= 0) {
                        throw new ExecutableLoweringAnalysisException(
                                "stack underflow at " + instruction.kind().name().toLowerCase() + ": need=1 have=" + outHeight);
                    }
                    outHeight -= 1;
                }
                case ADD, SUB, MUL, DIV, MOD, EQ, NEQ, LT, LTE, GT, GTE, AND, OR -> {
                    if (outHeight < 2) {
                        throw new ExecutableLoweringAnalysisException(
                                "stack underflow at " + instruction.kind().name().toLowerCase() + ": need=2 have=" + outHeight);
                    }
                    outHeight -= 1;
                }
                case NOT, NEG -> {
                    if (outHeight <= 0) {
                        throw new ExecutableLoweringAnalysisException(
                                "stack underflow at " + instruction.kind().name().toLowerCase() + ": need=1 have=" + outHeight);
                    }
                }
                case CALL_FUNC -> {
                    final var argSlots = instruction.expectedArgSlots() == null ? 0 : instruction.expectedArgSlots();
                    final var retSlots = instruction.expectedRetSlots() == null ? 0 : instruction.expectedRetSlots();
                    if (outHeight < argSlots) {
                        throw new ExecutableLoweringAnalysisException(
                                "stack underflow at call_func: need=%d have=%d".formatted(argSlots, outHeight));
                    }
                    outHeight = outHeight - argSlots + retSlots;
                }
                case CALL_HOST -> {
                    final var argSlots = instruction.hostCall() == null ? 0 : instruction.hostCall().argSlots();
                    final var retSlots = instruction.hostCall() == null ? 0 : instruction.hostCall().retSlots();
                    if (outHeight < argSlots) {
                        throw new ExecutableLoweringAnalysisException(
                                "stack underflow at call_host: need=%d have=%d".formatted(argSlots, outHeight));
                    }
                    outHeight = outHeight - argSlots + retSlots;
                }
                case CALL_INTRINSIC -> {
                    final var argSlots = instruction.expectedArgSlots() == null ? 0 : instruction.expectedArgSlots();
                    final var retSlots = instruction.expectedRetSlots() == null ? 0 : instruction.expectedRetSlots();
                    if (outHeight < argSlots) {
                        throw new ExecutableLoweringAnalysisException(
                                "stack underflow at call_intrinsic: need=%d have=%d".formatted(argSlots, outHeight));
                    }
                    outHeight = outHeight - argSlots + retSlots;
                }
                case HALT, LABEL, JMP, RET -> {
                }
                case JMP_IF_TRUE, JMP_IF_FALSE -> {
                    if (outHeight <= 0) {
                        throw new ExecutableLoweringAnalysisException("conditional jump requires stack height > 0");
                    }
                    outHeight -= 1;
                }
            }

            maxHeight = Math.max(maxHeight, outHeight);
            if (instruction.kind() == IRBackendExecutableFunction.InstructionKind.RET
                    || instruction.kind() == IRBackendExecutableFunction.InstructionKind.HALT) {
                continue;
            }

            final var successors = new ArrayList<Integer>(2);
            if (instruction.kind() == IRBackendExecutableFunction.InstructionKind.JMP
                    || instruction.kind() == IRBackendExecutableFunction.InstructionKind.JMP_IF_TRUE
                    || instruction.kind() == IRBackendExecutableFunction.InstructionKind.JMP_IF_FALSE) {
                final var targetIndex = labelToIndex.get(instruction.targetLabel());
                if (targetIndex == null) {
                    throw new ExecutableLoweringAnalysisException("missing jump target label: " + instruction.targetLabel());
                }
                successors.add(targetIndex);
                if (instruction.kind() == IRBackendExecutableFunction.InstructionKind.JMP) {
                    mergeIncomingHeights(incomingHeightByInstruction, worklist, targetIndex, outHeight);
                    continue;
                }
            }

            final var nextIndex = index + 1;
            if (nextIndex >= instructions.size()) {
                throw new ExecutableLoweringAnalysisException("reachable path ends without terminal instruction");
            }
            successors.add(nextIndex);
            for (final var successor : successors) {
                mergeIncomingHeights(incomingHeightByInstruction, worklist, successor, outHeight);
            }
        }

        if (!visited.contains(instructions.size() - 1)) {
            throw new ExecutableLoweringAnalysisException("entrypoint does not reach function tail");
        }
        return Math.max(1, Math.max(maxHeight, returnSlots));
    }

    private void mergeIncomingHeights(
            final Map<Integer, Integer> incomingHeightByInstruction,
            final ArrayDeque<Integer> worklist,
            final int targetIndex,
            final int newHeight) {
        final var existingHeight = incomingHeightByInstruction.get(targetIndex);
        if (existingHeight == null) {
            incomingHeightByInstruction.put(targetIndex, newHeight);
            worklist.add(targetIndex);
            return;
        }
        if (existingHeight != newHeight) {
            throw new ExecutableLoweringAnalysisException(
                    "stack mismatch at join: existing=%d new=%d".formatted(existingHeight, newHeight));
        }
    }

    static final class ExecutableLoweringAnalysisException extends RuntimeException {
        private ExecutableLoweringAnalysisException(final String message) {
            super(message);
        }
    }
}
