if (!navigator.gpu) {
  throw new Error("WebGPU not supported on this browser.");
}

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) {
  throw new Error("No WebGPU adapter");
}

const device = await adapter.requestDevice();

const canvas = document.querySelector("canvas");
canvas.width = canvas.clientWidth
canvas.height = canvas.clientHeight

const context = canvas.getContext("webgpu");
const canvasFormat = navigator.gpu.getPreferredCanvasFormat();
//console.log(canvasFormat);
context.configure({
  device: device,
  format: canvasFormat,
});

const encoder = device.createCommandEncoder();

const renderPass = encoder.beginRenderPass({
  colorAttachments: [{
    view: context.getCurrentTexture().createView(),
    loadOp: "clear",
    clearValue: { r: 0.9, g: 0.5, b: 0.2, a: 1 },
    storeOp: "store",
  }]
});
renderPass.end();

const commandBuffer = encoder.finish();
device.queue.submit([commandBuffer]);

const vertices = new Float32Array([
//   X,    Y,
  -0.8, -0.8, // Triangle 1 (Blue)
   0.8, -0.8,
   0.8,  0.8,

  -0.8, -0.8, // Triangle 2 (Red)
   0.8,  0.8,
  -0.8,  0.8,
]);

const vertexBuffer = device.createBuffer({
  label: "cell vertices",
  size: vertices.byteLength,
  usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(vertexBuffer, 0, vertices);

const vertexBufferLayout = {
  arrayStride: 8,
  attributes: [{
    format: "float32x2",
    offset: 0,
    shaderLocation: 0,
  }],
  stepMode: "vertex",
};

function getPath(path)
{
  return fetch(path).then((response) => {
    if (!response.ok) {
      throw new Error(`${path}: ${response.status}`)
    }
    return response.blob()
  }).then((blob) => { return blob.text() })
}

const cellShaderModule = device.createShaderModule({
  label: "cell shader",
  code: await getPath("index.wgsl")
});

const bindGroupLayout = device.createBindGroupLayout({
  label: "bind group layout",
  entries: [{
    binding: 0,
    visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
    buffer: { type: "uniform" }
  }, {
    binding: 1,
    visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
    buffer: { type: "read-only-storage" }
  }, {
    binding: 2,
    visibility: GPUShaderStage.COMPUTE,
    buffer: { type: "storage" }
  }]
});

const pipelineLayout = device.createPipelineLayout({
  label: "pipeline layout",
  bindGroupLayouts: [ bindGroupLayout ],
});

const cellPipeline = device.createRenderPipeline({
  label: "cell pipeline",
  layout: pipelineLayout,
  vertex: {
    module: cellShaderModule,
    entryPoint: "vertexMain",
    buffers: [vertexBufferLayout],
  },
  fragment: {
    module: cellShaderModule,
    entryPoint: "fragmentMain",
    targets: [{
      format: canvasFormat,
    }]
  }
});

const computeShaderModule = device.createShaderModule({
  label: "compute shader",
  code: await getPath("compute.wgsl")
});

const computePipeline = device.createComputePipeline({
  label: "compute pipeline",
  layout: pipelineLayout,
  compute: {
    module: computeShaderModule,
    entryPoint: "computeMain",
  },
});

const gridSize = 32;
const workGroupSize = 16;

const uniformArray = new Float32Array([gridSize, gridSize, 1, 1]);
const uniformBuffer = device.createBuffer({
  label: "grid uniform",
  size: uniformArray.byteLength,
  usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(uniformBuffer, 0, uniformArray);

const cellStateArray = new Uint32Array(gridSize * gridSize);

const cellStateStorage = [
  device.createBuffer({
    label: "cell state 0",
    size: cellStateArray.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
  }),
  device.createBuffer({
    label: "cell state 1",
    size: cellStateArray.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
  }),
];

for (let i = 0; i < cellStateArray.length; i += 1) {
  cellStateArray[i] = Math.random() > 0.6 ? 1 : 0;
}
device.queue.writeBuffer(cellStateStorage[0], 0, cellStateArray);
device.queue.writeBuffer(cellStateStorage[1], 0, cellStateArray);

const bindGroups = [
  device.createBindGroup({
    label: "cell bind group 0",
    layout: bindGroupLayout,
    entries: [{
      binding: 0,
      resource: { buffer: uniformBuffer },
    }, {
      binding: 1,
      resource: { buffer: cellStateStorage[0] },
    }, {
      binding: 2,
      resource: { buffer: cellStateStorage[1] },
    }],
  }),
  device.createBindGroup({
    label: "cell bind group 1",
    layout: bindGroupLayout,
    entries: [{
      binding: 0,
      resource: { buffer: uniformBuffer },
    }, {
      binding: 1,
      resource: { buffer: cellStateStorage[1] },
    }, {
      binding: 2,
      resource: { buffer: cellStateStorage[0] },
    }],
  })
];

let tick = 0;
let which = 0;

function render()
{
  const encoder = device.createCommandEncoder();

  const computePass = encoder.beginComputePass();
  computePass.setPipeline(computePipeline);
  computePass.setBindGroup(0, bindGroups[which]);
  const workgroupCount = Math.ceil((gridSize * gridSize) / workGroupSize);
  computePass.dispatchWorkgroups(workgroupCount);
  computePass.end();

  tick += 1;
  if ((tick % 5) == 0) {
    which ^= 1;
  }

  canvas.width = canvas.clientWidth
  canvas.height = canvas.clientHeight

  const aspect = canvas.height / canvas.width;
  if (canvas.width > canvas.height) {
    uniformArray[2] = aspect;
    uniformArray[3] = 1;
  } else {
    uniformArray[2] = 1;
    uniformArray[3] = 1 / aspect;
  }
  device.queue.writeBuffer(uniformBuffer, 0, uniformArray, 0);

  const renderPass = encoder.beginRenderPass({
    colorAttachments: [{
      view: context.getCurrentTexture().createView(),
      loadOp: "clear",
      clearValue: { r: 0.2, g: 0.2, b: 0.4, a: 1 },
      storeOp: "store",
    }]
  });

  renderPass.setPipeline(cellPipeline);
  renderPass.setVertexBuffer(0, vertexBuffer);
  renderPass.setBindGroup(0, bindGroups[which]);
  renderPass.draw(vertices.length / 2, gridSize * gridSize);

  renderPass.end();

  const commandBuffer = encoder.finish();
  device.queue.submit([commandBuffer]);

  requestAnimationFrame(render)
}
requestAnimationFrame(render)