import * as fs from "node:fs"

function loadShader(file) {
    try {
        const data = fs.readFileSync(file, 'utf8');
        return data;
    } catch (err) {
        console.error('Failed to load shader:', err);
    }
}
const inputArray = new Uint8Array([150, 70, 43, 68, 72, 93, 63, 3, 183, 151, 92, 60, 76, 35, 220, 145, 218, 41, 20, 253, 167, 52, 142, 101, 213, 48, 121, 26, 214, 81, 179, 223, 193, 2, 227, 92, 253, 187, 11, 222, 231, 31, 209, 196, 41, 242, 17, 90])
const shader = loadShader('compute.wgsl');

async function main() {
    const adapter = await navigator.gpu.requestAdapter();
    const device = await adapter.requestDevice();
    const queue = device.queue
    
    // Setup for beginning a render pass
    const encoder = device.createCommandEncoder();
    const texture = device.createTexture({
        format: "bgra8unorm",
        label: "texture",
        size: {
            width: 12,
            height: 12,
            depthOrArrayLayers: 1
        },
        usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_DST | GPUTextureUsage.COPY_SRC,
        viewFormats: ["bgra8unorm"]
    });

    const view = texture.createView({
        dimension: "2d-array", // Works fine when it is "2d" 
    })

    const renderPassEncoder = encoder.beginRenderPass({
        colorAttachments: [{
            loadOp: "load",
            storeOp: "discard",
            view: view
        }],
        maxDrawCount: 10,
    });
    renderPassEncoder.end(); // This also affects the output of the compute pass 

    // Setup for compute pass
    const inputBuffer = device.createBuffer({
        size: 48,
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST | GPUBufferUsage.INDEX | GPUBufferUsage.QUERY_RESOLVE
    });
    const storageBuffer = device.createBuffer({
        size: 32,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.INDIRECT | GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.VERTEX | GPUBufferUsage.INDEX
    });

    const shaderModule = device.createShaderModule({
        code: shader
    });
    const computePipeline = device.createComputePipeline({
        compute: {
            entryPoint: "main",
            module: shaderModule
        },
        layout: "auto"
    });

    const computePassEncoder = encoder.beginComputePass();
    computePassEncoder.setPipeline(computePipeline);
    queue.writeBuffer(inputBuffer, 0, inputArray);

    const GPUBindGroup3 = device.createBindGroup({
        layout: computePipeline.getBindGroupLayout(0),
        entries: [{
            binding: 0,
            resource: {
                buffer: inputBuffer
            }
        }, {
            binding: 1,
            resource: {
                buffer: storageBuffer
            }
        }]
    });
    computePassEncoder.setBindGroup(0, GPUBindGroup3);
    computePassEncoder.dispatchWorkgroups(1);
    computePassEncoder.end();

    // Copy results of storage buffer to output buffer
    const outputBuffer = device.createBuffer({
        size: 32,
        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
    });
    encoder.copyBufferToBuffer(storageBuffer, 0, outputBuffer, 0, 32);

    queue.submit([encoder.finish()]);

    // Print the output of the compute pass
    await outputBuffer.mapAsync(GPUMapMode.READ)
    const arrayBuffer = outputBuffer.getMappedRange();
    const computePassResult = arrayBuffer.slice(0)
    const readableArray = new Uint8Array(computePassResult)
    console.log("Compute pass has output:", ...readableArray);
}
main().catch(console.error);