import { mat4 } from 'https://webgpufundamentals.org/3rdparty/wgpu-matrix.module.js';
import { fetchShaderCode, generateGlyphTextureAtlas } from './utility.js';
import { CONFIG } from './config.js';
import { CANVAS, CTX, COLORS, RENDER_PASS_DESCRIPTOR } from './constants.js';

const canvas = document.querySelector('canvas');
const context = canvas.getContext('webgpu');
const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
let device;
let pipeline;
let vertexBuffer;
let indexBuffer;
let uniformBuffer;
let texture;
let sampler;
let bindGroup;

async function main() {
    const adapter = await navigator.gpu?.requestAdapter();
    device = await adapter?.requestDevice();
    if (!device) {
        alert('need a browser that supports WebGPU');
        return;
    }

    context.configure({
        device,
        format: presentationFormat,
    });

    const shaderCode = await fetchShaderCode('shaders.wgsl');
    const module = device.createShaderModule({
        label: 'textured quad shaders',
        code: shaderCode,
    });

    const glyphCanvas = generateGlyphTextureAtlas(CANVAS, CTX, CONFIG);
    document.body.appendChild(glyphCanvas);
    glyphCanvas.style.backgroundColor = '#222';

    const vertexSize = CONFIG.floatsPerVertex * 4;
    const vertexBufferSize = CONFIG.maxGlyphs * CONFIG.vertsPerGlyph * vertexSize;
    vertexBuffer = device.createBuffer({
        label: 'vertices',
        size: vertexBufferSize,
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
    });

    indexBuffer = device.createBuffer({
        label: 'indices',
        size: CONFIG.maxGlyphs * CONFIG.vertsPerGlyph * 4,
        usage: GPUBufferUsage.INDEX | GPUBufferUsage.COPY_DST,
    });

    const indices = Array.from({ length: CONFIG.maxGlyphs * 6 }, (_, i) => {
        const ndx = Math.floor(i / 6) * 4;
        return (i % 6 < 3 ? [ndx, ndx + 1, ndx + 2] : [ndx + 2, ndx + 1, ndx + 3])[i % 3];
    });
    device.queue.writeBuffer(indexBuffer, 0, new Uint32Array(indices));

    const { vertexData, numGlyphs, width, height } = generateGlyphVerticesForText('Hello\nworld!\nText in\nWebGPU!', COLORS, glyphCanvas);
    device.queue.writeBuffer(vertexBuffer, 0, vertexData);

    pipeline = device.createRenderPipeline({
        label: 'textured quad pipeline',
        layout: 'auto',
        vertex: {
            module,
            entryPoint: 'vs',
            buffers: [
                {
                    arrayStride: vertexSize,
                    attributes: [
                        { shaderLocation: 0, offset: 0, format: 'float32x2' },  // pos
                        { shaderLocation: 1, offset: 8, format: 'float32x2' },  // tex
                        { shaderLocation: 2, offset: 16, format: 'float32x4' }  // col
                    ],
                },
            ],
        },
        fragment: {
            module,
            entryPoint: 'fs',
            targets: [{
                format: presentationFormat,
                blend: {
                    color: { srcFactor: 'one', dstFactor: 'one-minus-src-alpha', operation: 'add' },
                    alpha: { srcFactor: 'one', dstFactor: 'one-minus-src-alpha', operation: 'add' }
                },
            }],
        },
    });

    texture = createTextureFromSource(device, glyphCanvas, { mips: true });
    sampler = device.createSampler({
        minFilter: 'linear',
        magFilter: 'linear'
    });

    uniformBuffer = device.createBuffer({
        label: 'uniforms for quad',
        size: CONFIG.uniformBufferSize,
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
    });

    const uniformValues = new Float32Array(CONFIG.uniformBufferSize / 4);
    const matrix = uniformValues.subarray(0, 16);

    bindGroup = device.createBindGroup({
        layout: pipeline.getBindGroupLayout(0),
        entries: [
            { binding: 0, resource: sampler },
            { binding: 1, resource: texture.createView() },
            { binding: 2, resource: { buffer: uniformBuffer } },
        ],
    });

    function render(time, context, pipeline, uniformBuffer, uniformValues, bindGroup, vertexBuffer, indexBuffer, RENDER_PASS_DESCRIPTOR) {
        time *= CONFIG.time.phase;
        const fov = 60 * Math.PI / 180;
        const aspect = canvas.clientWidth / canvas.clientHeight;
        const projectionMatrix = mat4.perspective(fov, aspect, CONFIG.render.zNear, CONFIG.render.zFar);
        const viewMatrix = mat4.lookAt([0, 0, 5], [0, 0, 0], [0, 1, 0]);
        const viewProjectionMatrix = mat4.multiply(projectionMatrix, viewMatrix);
        RENDER_PASS_DESCRIPTOR.colorAttachments[0].view = context.getCurrentTexture().createView();
        const encoder = device.createCommandEncoder();
        const pass = encoder.beginRenderPass(RENDER_PASS_DESCRIPTOR);
        pass.setPipeline(pipeline);
        mat4.rotateY(viewProjectionMatrix, time, matrix);
        mat4.translate(matrix, [-width / 2, -height / 2, 0], matrix);
        device.queue.writeBuffer(uniformBuffer, 0, uniformValues);
        pass.setBindGroup(0, bindGroup);
        pass.setVertexBuffer(0, vertexBuffer);
        pass.setIndexBuffer(indexBuffer, 'uint32');
        pass.drawIndexed(numGlyphs * 6);
        pass.end();
        device.queue.submit([encoder.finish()]);
        requestAnimationFrame((time) => render(time, context, pipeline, uniformBuffer, uniformValues, bindGroup, vertexBuffer, indexBuffer, RENDER_PASS_DESCRIPTOR));
    }
    
    requestAnimationFrame((time) => render(time, context, pipeline, uniformBuffer, uniformValues, bindGroup, vertexBuffer, indexBuffer, RENDER_PASS_DESCRIPTOR));
}

function generateGlyphVerticesForText(s, COLORS, glyphCanvas) {
    const vertexData = new Float32Array(CONFIG.maxGlyphs * CONFIG.floatsPerVertex * CONFIG.vertsPerGlyph);
    const glyphUVWidth = CONFIG.glyphWidth / glyphCanvas.width;
    const glyphUVHeight = CONFIG.glyphHeight / glyphCanvas.height;
    let offset = 0, x0 = 0, y0 = 0, x1 = 1, y1 = 1, width = 0;
    let colorNdx = 0;

    const addVertex = (x, y, u, v, color) => {
        vertexData.set([x, y, u, v, ...color], offset);
        offset += 8;
    };

    for (let i = 0; i < s.length; ++i) {
        const c = s.charCodeAt(i);
        if (c >= 33) {
            const cIndex = c - 33;
            const glyphX = cIndex % CONFIG.glyphsAcrossTexture;
            const glyphY = Math.floor(cIndex / CONFIG.glyphsAcrossTexture);
            const u0 = glyphX * CONFIG.glyphWidth / glyphCanvas.width;
            const v1 = glyphY * CONFIG.glyphHeight / glyphCanvas.height;
            const u1 = u0 + glyphUVWidth;
            const v0 = v1 + glyphUVHeight;
            width = Math.max(x1, width);
            addVertex(x0, y0, u0, v0, COLORS[colorNdx]);
            addVertex(x1, y0, u1, v0, COLORS[colorNdx]);
            addVertex(x0, y1, u0, v1, COLORS[colorNdx]);
            addVertex(x1, y1, u1, v1, COLORS[colorNdx]);
        } else {
            colorNdx = (colorNdx + 1) % COLORS.length;
            if (c === 10) { // Newline
                x0 = 0; x1 = 1; y0--; y1 = y0 + 1;
                continue;
            }
        }
        x0 += 0.55; x1 = x0 + 1;
    }
    return { vertexData, numGlyphs: offset / CONFIG.floatsPerVertex, width, height: y1 };
}

function createTextureFromSource(device, source, options = {}) {
    const texture = device.createTexture({
        format: 'rgba8unorm',
        size: [source.width, source.height],
        usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST | GPUTextureUsage.RENDER_ATTACHMENT,
    });

    device.queue.copyExternalImageToTexture(
        { source, flipY: options.flipY },
        { texture, premultipliedAlpha: true },
        { width: source.width, height: source.height }
    );

    return texture;
}

main();