package p.studio.compiler.backend.irvm;

import p.studio.compiler.backend.bytecode.BytecodeEmitter;

import java.util.ArrayDeque;
import java.util.HashMap;
import java.util.HashSet;

public class IRVMValidator {
    private final IRVMProfileFeatureGate profileFeatureGate;

    public IRVMValidator() {
        this(new IRVMProfileFeatureGate());
    }

    IRVMValidator(final IRVMProfileFeatureGate profileFeatureGate) {
        this.profileFeatureGate = profileFeatureGate;
    }

    public void validate(
            final IRVMProgram program,
            final boolean rejectInternalOpcodes) {
        final var input = program == null ? IRVMProgram.empty() : program;
        final BytecodeEmitter.EmissionPlan emissionPlan;
        if (input.emissionPlan() == null || input.emissionPlan().functions().isEmpty()) {
            emissionPlan = null;
        } else {
            emissionPlan = input.coherentEmissionPlan();
        }
        for (var functionIndex = 0; functionIndex < input.module().functions().size(); functionIndex++) {
            final var functionPlan = emissionPlan == null
                    ? null
                    : emissionPlan.functions().get(functionIndex);
            validateFunction(input.module(), functionPlan, functionIndex, rejectInternalOpcodes);
        }
    }

    public void validate(
            final IRVMModule module,
            final boolean rejectInternalOpcodes) {
        final var input = module == null ? IRVMModule.empty() : module;
        for (var functionIndex = 0; functionIndex < input.functions().size(); functionIndex++) {
            validateFunction(input, null, functionIndex, rejectInternalOpcodes);
        }
    }

    private void validateFunction(
            final IRVMModule module,
            final BytecodeEmitter.FunctionPlan functionPlan,
            final int functionIndex,
            final boolean rejectInternalOpcodes) {
        final var function = module.functions().get(functionIndex);
        final var instructions = function.instructions();
        if (instructions.isEmpty()) {
            return;
        }

        final var pcByIndex = new int[instructions.size()];
        final var indexByPc = new HashMap<Integer, Integer>();
        var pc = 0;
        for (var i = 0; i < instructions.size(); i++) {
            final var instr = instructions.get(i);
            pcByIndex[i] = pc;
            indexByPc.put(pc, i);
            pc += instr.encodedSize();
            if (rejectInternalOpcodes && instr.op().internal()) {
                throw new IRVMValidationException(
                        IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_INTERNAL_OPCODE_RESIDUAL,
                        "internal opcode must be eliminated before emission",
                        functionIndex,
                        pcByIndex[i]);
            }
            if (!profileFeatureGate.isAllowed(module.vmProfile(), instr.op())) {
                throw new IRVMValidationException(
                        IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_OPCODE_NOT_ALLOWED_FOR_PROFILE,
                        "opcode '%s' is not allowed for vm profile '%s'".formatted(instr.op().name(), module.vmProfile()),
                        functionIndex,
                        pcByIndex[i]);
            }
        }

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

        while (!worklist.isEmpty()) {
            final var index = worklist.removeFirst();
            visited.add(index);
            final var instr = instructions.get(index);
            final var instrPc = pcByIndex[index];
            var inHeight = incomingHeights.get(index);

            if (instr.op() == IRVMOp.CALL) {
                final var targetFn = instr.immediate();
                if (targetFn == null || targetFn < 0 || targetFn >= module.functions().size()) {
                    throw new IRVMValidationException(
                            IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_INVALID_FUNC_ID,
                            "invalid callee function id: " + targetFn,
                            functionIndex,
                            instrPc);
                }
                final var callee = module.functions().get(targetFn);
                final var outHeight = applyStackEffect(
                        inHeight,
                        callee.paramSlots(),
                        callee.returnSlots(),
                        function.maxStackSlots(),
                        functionIndex,
                        instrPc);
                enqueueSuccessors(
                        module,
                        functionIndex,
                        function,
                        index,
                        instr,
                        instrPc,
                        pcByIndex,
                        indexByPc,
                        outHeight,
                        incomingHeights,
                        worklist);
                continue;
            }

            final var stackEffect = resolveStackEffect(
                    functionPlan,
                    index,
                    instr,
                    functionIndex,
                    instrPc);
            final var outHeight = applyStackEffect(
                    inHeight,
                    stackEffect.pops(),
                    stackEffect.pushes(),
                    function.maxStackSlots(),
                    functionIndex,
                    instrPc);
            if (instr.op() == IRVMOp.RET && outHeight != function.returnSlots()) {
                throw new IRVMValidationException(
                        IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_BAD_RET_STACK_HEIGHT,
                        "ret stack height mismatch. expected=%d actual=%d".formatted(function.returnSlots(), outHeight),
                        functionIndex,
                        instrPc);
            }
            enqueueSuccessors(
                    module,
                    functionIndex,
                    function,
                    index,
                    instr,
                    instrPc,
                    pcByIndex,
                    indexByPc,
                    outHeight,
                    incomingHeights,
                    worklist);
        }

        if (visited.isEmpty()) {
            return;
        }
    }

    private int applyStackEffect(
            final int inHeight,
            final int pops,
            final int pushes,
            final int maxStackSlots,
            final int functionIndex,
            final int pc) {
        if (inHeight < pops) {
            throw new IRVMValidationException(
                    IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_STACK_UNDERFLOW,
                    "stack underflow. need=%d have=%d".formatted(pops, inHeight),
                    functionIndex,
                    pc);
        }
        final var outHeight = inHeight - pops + pushes;
        if (outHeight > maxStackSlots) {
            throw new IRVMValidationException(
                    IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_STACK_OVERFLOW,
                    "stack overflow. max=" + maxStackSlots + " actual=" + outHeight,
                    functionIndex,
                    pc);
        }
        return outHeight;
    }

    private void enqueueSuccessors(
            final IRVMModule module,
            final int functionIndex,
            final IRVMFunction function,
            final int index,
            final IRVMInstruction instr,
            final int instrPc,
            final int[] pcByIndex,
            final HashMap<Integer, Integer> indexByPc,
            final int outHeight,
            final HashMap<Integer, Integer> incomingHeights,
            final ArrayDeque<Integer> worklist) {
        if (instr.op() == IRVMOp.HALT || instr.op() == IRVMOp.RET) {
            return;
        }

        if (instr.op() == IRVMOp.JMP || instr.op() == IRVMOp.JMP_IF_TRUE || instr.op() == IRVMOp.JMP_IF_FALSE) {
            final var targetPc = instr.immediate();
            final var targetIndex = indexByPc.get(targetPc);
            if (targetIndex == null) {
                throw new IRVMValidationException(
                        IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_INVALID_JUMP_TARGET,
                        "invalid jump target pc=" + targetPc,
                        functionIndex,
                        instrPc);
            }
            mergeIncomingHeight(
                    functionIndex,
                    pcByIndex[targetIndex],
                    targetIndex,
                    outHeight,
                    incomingHeights,
                    worklist);
            if (instr.op() == IRVMOp.JMP) {
                return;
            }
        }

        final var nextIndex = index + 1;
        if (nextIndex >= function.instructions().size()) {
            throw new IRVMValidationException(
                    IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_UNTERMINATED_PATH,
                    "reachable path ends without terminator",
                    functionIndex,
                    instrPc);
        }
        mergeIncomingHeight(
                functionIndex,
                pcByIndex[nextIndex],
                nextIndex,
                outHeight,
                incomingHeights,
                worklist);
    }

    private void mergeIncomingHeight(
            final int functionIndex,
            final int pc,
            final int index,
            final int newHeight,
            final HashMap<Integer, Integer> incomingHeights,
            final ArrayDeque<Integer> worklist) {
        final var existing = incomingHeights.get(index);
        if (existing == null) {
            incomingHeights.put(index, newHeight);
            worklist.add(index);
            return;
        }
        if (existing != newHeight) {
            throw new IRVMValidationException(
                    IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_STACK_MISMATCH_JOIN,
                    "stack height mismatch at join. existing=%d new=%d".formatted(existing, newHeight),
                    functionIndex,
                    pc);
        }
    }

    private StackEffect resolveStackEffect(
            final BytecodeEmitter.FunctionPlan functionPlan,
            final int instructionIndex,
            final IRVMInstruction instruction,
            final int functionIndex,
            final int pc) {
        final BytecodeEmitter.Operation operation;
        if (functionPlan == null || instructionIndex >= functionPlan.operations().size()) {
            operation = null;
        } else {
            operation = functionPlan.operations().get(instructionIndex);
        }

        if (instruction.op() == IRVMOp.HOSTCALL) {
            if (operation == null || operation.kind() != BytecodeEmitter.OperationKind.HOSTCALL) {
                return new StackEffect(instruction.op().pops(), instruction.op().pushes());
            }
            final Integer pops = operation.expectedArgSlots() != null
                    ? operation.expectedArgSlots()
                    : operation.syscallDecl() == null ? null : operation.syscallDecl().argSlots();
            final Integer pushes = operation.expectedRetSlots() != null
                    ? operation.expectedRetSlots()
                    : operation.syscallDecl() == null ? null : operation.syscallDecl().retSlots();
            if (pops == null || pushes == null) {
                throw new IRVMValidationException(
                        IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_HOSTCALL_METADATA_MISSING,
                        "hostcall stack metadata is missing",
                        functionIndex,
                        pc);
            }
            return new StackEffect(pops, pushes);
        }

        if (instruction.op() == IRVMOp.INTRINSIC) {
            if (operation == null || operation.kind() != BytecodeEmitter.OperationKind.INTRINSIC) {
                return new StackEffect(instruction.op().pops(), instruction.op().pushes());
            }
            if (operation.expectedArgSlots() == null || operation.expectedRetSlots() == null) {
                throw new IRVMValidationException(
                        IRVMValidationErrorCode.MARSHAL_VERIFY_PRECHECK_INTRINSIC_SIGNATURE_MISSING,
                        "intrinsic stack signature metadata is missing",
                        functionIndex,
                        pc);
            }
            return new StackEffect(operation.expectedArgSlots(), operation.expectedRetSlots());
        }

        return new StackEffect(instruction.op().pops(), instruction.op().pushes());
    }

    private record StackEffect(
            int pops,
            int pushes) {
    }
}
