import { mat4 } from './glm.js';

// Initialize WebGPU
const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();
const canvas = document.querySelector('canvas');
const context = canvas.getContext('webgpu');
const format = navigator.gpu.getPreferredCanvasFormat();
context.configure({ device, format });

// Create depth texture
const depthTexture = device.createTexture({
    size: [canvas.width, canvas.height],
    usage: GPUTextureUsage.RENDER_ATTACHMENT,
    format: 'depth24plus',
});

// Prepare color texture
// 1. fetch the texture from the server
const imageBitmap = await fetch('base.png')
    .then((response) => response.blob())
    .then((blob) => createImageBitmap(blob));

// 2. create a texture
const colorTexture = device.createTexture({
    size: [imageBitmap.width, imageBitmap.height],
    usage:
        GPUTextureUsage.TEXTURE_BINDING |
        GPUTextureUsage.RENDER_ATTACHMENT |
        GPUTextureUsage.COPY_DST,
    format: 'rgba8unorm',
    mipLevelCount: 2,
});

// 3. transfer data
device.queue.copyExternalImageToTexture(
    { source: imageBitmap },
    { texture: colorTexture },
    [imageBitmap.width, imageBitmap.height],
);

// 4. create sampler
const colorSampler = device.createSampler({
    mipmapFilter: 'linear',
});

// Create vertex buffer
// prettier-ignore
const vertices = new Float32Array([
    // position       // color
    -1, -1, -1,  1,   1, 0, 0, 1,
     1, -1, -1,  1,   0, 1, 0, 1,
    -1,  1, -1,  1,   0, 0, 1, 1,
     1,  1, -1,  1,   1, 1, 0, 1,

    -1, -1,  1,  1,   1, 0, 0, 1,
     1, -1,  1,  1,   0, 1, 0, 1,
    -1,  1,  1,  1,   0, 0, 1, 1,
     1,  1,  1,  1,   1, 1, 0, 1,
]);

const vertexBuffer = device.createBuffer({
    size: vertices.byteLength,
    usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
});

device.queue.writeBuffer(vertexBuffer, 0, vertices);

// Create index buffer
// prettier-ignore
const indices = new Uint32Array([
    0, 1, 2,   1, 2, 3,
    0, 1, 3,   1, 4, 5,
    0, 4, 2,   4, 2, 5,
    1, 5, 3,   5, 3, 7,
    2, 3, 6,   3, 6, 7,
    4, 5, 6,   5, 6, 7,
]);

const indexBuffer = device.createBuffer({
    size: indices.byteLength,
    usage: GPUBufferUsage.INDEX | GPUBufferUsage.COPY_DST,
});

device.queue.writeBuffer(indexBuffer, 0, indices);

// Fetch and compile shaders
const code = await fetch('shader.wgsl').then((response) => response.text());
const module = device.createShaderModule({ code });

// Create the pipeline
/** @type {GPUVertexBufferLayout} */
const vertexBufferLayout = {
    arrayStride: 32,
    attributes: [
        {
            shaderLocation: 0,
            offset: 0,
            format: 'float32x4',
        },
        {
            shaderLocation: 1,
            offset: 8,
            format: 'float32x4',
        },
    ],
};

const pipeline = device.createRenderPipeline({
    vertex: {
        module,
        buffers: [vertexBufferLayout],
    },
    fragment: {
        module,
        targets: [{ format }],
    },
    depthStencil: {
        depthWriteEnabled: true,
        depthCompare: 'less',
        format: 'depth24plus',
    },
    layout: 'auto',
});

// Create uniform buffer
const uniformBuffer = device.createBuffer({
    size: 4 * 4 * 4,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
});

// Create the bind group
const bindGroup = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
        { binding: 0, resource: { buffer: uniformBuffer } },
        { binding: 1, resource: colorTexture.createView() },
        { binding: 2, resource: colorSampler },
    ],
});

function update() {
    // Animate the square
    const time = performance.now() / 1000;

    const modelMatrix = mat4
        .create()
        .rotateX(time * 0.7)
        .rotateY(time);
    const viewMatrix = mat4.create().translate([0, 0, 5]).invert();
    const projectionMatrix = mat4.create().perspectiveZO(1, 1, 0.01, 100);

    const matrix = mat4
        .create()
        .multiply(projectionMatrix)
        .multiply(viewMatrix)
        .multiply(modelMatrix);

    // const radius = 0.5;
    // const frequency = 0.5;
    // const x = radius * Math.cos(frequency * time * 2 * Math.PI);
    // const y = radius * Math.sin(frequency * time * 2 * Math.PI);

    device.queue.writeBuffer(uniformBuffer, 0, matrix);
}

function render() {
    // Render the square
    const commandEncoder = device.createCommandEncoder();
    const renderPass = commandEncoder.beginRenderPass({
        colorAttachments: [
            {
                view: context.getCurrentTexture().createView(),
                loadOp: 'clear',
                clearValue: [1, 1, 1, 1],
                storeOp: 'store',
            },
        ],
        depthStencilAttachment: {
            view: depthTexture.createView(),
            depthLoadOp: 'clear',
            depthClearValue: 1,
            depthStoreOp: 'discard',
        },
    });
    renderPass.setPipeline(pipeline);
    renderPass.setVertexBuffer(0, vertexBuffer);
    renderPass.setIndexBuffer(indexBuffer, 'uint32');
    renderPass.setBindGroup(0, bindGroup);
    renderPass.drawIndexed(indices.length);
    renderPass.end();
    device.queue.submit([commandEncoder.finish()]);
}

function frame() {
    update();
    render();
    requestAnimationFrame(frame);
}

requestAnimationFrame(frame);