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';
import { createPipeline } from './wgpu-pipeline.js';
import { createState } from './wgpu-state.js';

const canvas = document.querySelector('canvas');
const context = canvas.getContext('webgpu');
const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
const state = createState(config);

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

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

    const shaderCode = await fetchShaderCode('shaders.wgsl');
    const vertexSize = config.floatsPerVertex * 4;

    state.pipeline = await createPipeline(state.device, presentationFormat, vertexSize, shaderCode);

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

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

    state.indexBuffer = state.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];
    });
    state.device.queue.writeBuffer(state.indexBuffer, 0, new Uint32Array(indices));

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

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

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

    state.matrix = state.uniformValues.subarray(0, 16);

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

    state.numGlyphs = numGlyphs;
    state.width = width;
    state.height = height;

    requestAnimationFrame((time) => render(time, context, state, RENDER_PASS_DESCRIPTOR));
}

function render(time, context, state, 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 = state.device.createCommandEncoder();
    const pass = encoder.beginRenderPass(RENDER_PASS_DESCRIPTOR);
    pass.setPipeline(state.pipeline);
    mat4.rotateY(viewProjectionMatrix, time, state.matrix);
    mat4.translate(state.matrix, [-state.width / 2, -state.height / 2, 0], state.matrix);
    state.device.queue.writeBuffer(state.uniformBuffer, 0, state.uniformValues);
    pass.setBindGroup(0, state.bindGroup);
    pass.setVertexBuffer(0, state.vertexBuffer);
    pass.setIndexBuffer(state.indexBuffer, 'uint32');
    pass.drawIndexed(state.numGlyphs * 6);
    pass.end();
    state.device.queue.submit([encoder.finish()]);
    requestAnimationFrame((t) => render(t, context, state, 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();