package p.studio.compiler.backend.bytecode;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.Set;

public class BytecodePreloadVerifierService {
    private static final Set<Integer> TERMINATORS = Set.of(0x01, 0x02, 0x51);

    public void verify(final BytecodeModule module) {
        final var input = module == null ? BytecodeModule.empty() : module;
        for (final var function : input.functions()) {
            verifyFunction(input.code(), function);
        }
    }

    private void verifyFunction(final byte[] code, final BytecodeModule.FunctionMeta function) {
        final var start = function.codeOffset();
        final var end = start + function.codeLen();
        if (end <= start) {
            return;
        }

        var pc = start;
        var lastOpcode = -1;
        while (pc < end) {
            if (pc + 2 > end) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_TRUNCATED_INSTRUCTION,
                        "truncated opcode in function at pc=" + pc);
            }
            final var opcode = readU16(code, pc);
            if (opcode == 0x70) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_LINKAGE_RAW_SYSCALL_IN_PRELOAD,
                        "raw syscall is forbidden in pre-load artifact");
            }
            final var size = sizeForOpcode(opcode, pc);
            if (pc + size > end) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_TRUNCATED_INSTRUCTION,
                        "truncated instruction in function at pc=" + pc);
            }
            lastOpcode = opcode;
            pc += size;
        }

        if (!TERMINATORS.contains(lastOpcode)) {
            throw new BytecodeMarshalingException(
                    BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_UNTERMINATED_FUNCTION,
                    "function ends without terminator opcode");
        }
    }

    private int sizeForOpcode(
            final int opcode,
            final int pc) {
        try {
            return BytecodeOpcodeLayout.sizeForOpcode(opcode);
        } catch (BytecodeMarshalingException ignored) {
            throw new BytecodeMarshalingException(
                    BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_UNKNOWN_OPCODE,
                    "unknown opcode 0x%04X at pc=%d".formatted(opcode, pc));
        }
    }

    private int readU16(final byte[] code, final int offset) {
        return ByteBuffer.wrap(code, offset, 2).order(ByteOrder.LITTLE_ENDIAN).getShort() & 0xFFFF;
    }
}
