const addr = "http://localhost:6969/ws";

let ws = new WebSocket(addr);
ws.binaryType = "arraybuffer";

let wasmModule;
let wasmInstance;
let isCompiling = false;

ws.onmessage = async (message) => {
    if (!wasmInstance) {
        if (isCompiling) return; 
        isCompiling = true;
        try {
            const { instance, module } = await WebAssembly.instantiate(message.data);
            wasmInstance = instance;
            wasmModule = module; 
            ws.send(JSON.stringify({
                success: true,
                message: "WASM module loaded successfully.",
            }));
        } catch (err) {
            ws.send(JSON.stringify({
                success: false,
                message: "WASM module loaded successfully.",
            }));
        } finally {
            isCompiling = false;
        }
        return;
    }

    try {
        const exports = wasmInstance.exports;
        const inputPtr = passArrayBufferToWasm(exports, message.data);

        const outputPtr = exports.entry(inputPtr);
        const resultString = getStringFromWasm(exports.memory, outputPtr);

        ws.send(JSON.stringify({
            success: true,
            message: resultString,
        }));

        exports.deallocate(inputPtr);
        exports.deallocate(outputPtr); 
    } catch (error) {
        console.error(error);
        ws.send(JSON.stringify({
            success: true,
            message: `Failed to run wasm: ${error.message}`
        }));
    }
};

function passArrayBufferToWasm(exports, arrayBuffer) {
    const srcBytes = new Uint8Array(arrayBuffer);
    const len = srcBytes.length;
    const ptr = exports.allocate(len + 1);
    const heap = new Uint8Array(exports.memory.buffer);

    heap.set(srcBytes, ptr);
    heap[ptr + len] = 0; 
    
    return ptr;
}

function getStringFromWasm(memory, ptr) {
    const bytes = new Uint8Array(memory.buffer);
    let endPtr = ptr;
    
    while (endPtr < bytes.length && bytes[endPtr] !== 0) {
        endPtr++;
    }
    
    return new TextDecoder("utf-8").decode(bytes.slice(ptr, endPtr));
}