package p.studio.compiler.backend.bytecode;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;

public class BytecodeLinkPrecheckService {

    public void validate(final BytecodeModule module) {
        final var input = module == null ? BytecodeModule.empty() : module;
        validateFunctionBounds(input);
        validateExports(input);
        validateHostcallIndices(input);
    }

    private void validateFunctionBounds(final BytecodeModule module) {
        var previousEnd = 0;
        for (final var function : module.functions()) {
            if (function.codeOffset() < 0 || function.codeLen() < 0) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_FUNCTION_BOUNDS_INVALID,
                        "negative function code bounds");
            }
            if (function.codeOffset() < previousEnd) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_FUNCTION_BOUNDS_INVALID,
                        "function offsets must be monotonic");
            }
            final var end = function.codeOffset() + function.codeLen();
            if (end < function.codeOffset() || end > module.code().length) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_FUNCTION_BOUNDS_INVALID,
                        "function code range exceeds module code length");
            }
            previousEnd = end;
        }
    }

    private void validateExports(final BytecodeModule module) {
        for (final var export : module.exports()) {
            if (export.funcIdx() < 0 || export.funcIdx() >= module.functions().size()) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_EXPORT_FUNC_INDEX_INVALID,
                        "export references invalid func_idx: " + export.funcIdx());
            }
        }
    }

    private void validateHostcallIndices(final BytecodeModule module) {
        var pc = 0;
        while (pc < module.code().length) {
            if (pc + 2 > module.code().length) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_TRUNCATED_INSTRUCTION,
                        "truncated opcode at pc=" + pc);
            }
            final var opcode = readU16(module.code(), pc);
            final var size = sizeForOpcode(opcode, pc);
            if (pc + size > module.code().length) {
                throw new BytecodeMarshalingException(
                        BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_TRUNCATED_INSTRUCTION,
                        "truncated instruction at pc=" + pc);
            }
            if (opcode == 0x71) {
                final var index = readU32(module.code(), pc + 2);
                if (index < 0 || index >= module.syscalls().size()) {
                    throw new BytecodeMarshalingException(
                            BytecodeMarshalingErrorCode.MARSHAL_VERIFY_PRECHECK_HOSTCALL_INDEX_INVALID,
                            "hostcall index out of bounds: " + index);
                }
            }
            pc += size;
        }
    }

    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;
    }

    private int readU32(final byte[] code, final int offset) {
        return ByteBuffer.wrap(code, offset, 4).order(ByteOrder.LITTLE_ENDIAN).getInt();
    }
}
