import { createCompute } from "./gpu.js";
import shader from "./world.wgsl?raw";

export class World {
  constructor(canvas, gpu) {
    this.ctx = canvas.getContext("2d");
    this.gpu = gpu;
    this.power = 8.0;
    this.time = 0;

    this.image = this.ctx.createImageData(256, 256);

    if (gpu.mode === "gpu") this.initGPU();
  }

  async initGPU() {
    const d = this.gpu.device;
    this.pipeline = createCompute(d, shader);

    this.param = d.createBuffer({ size: 8, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
    this.buf = d.createBuffer({ size: 256*256*4*4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC });
    this.read = d.createBuffer({ size: 256*256*4*4, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
  }

  update(dt) {
    this.time += dt;
    if (this.gpu.mode !== "gpu") return;

    const d = this.gpu.device;
    d.queue.writeBuffer(this.param, 0, new Float32Array([this.time, this.power]));

    const enc = d.createCommandEncoder();
    const pass = enc.beginComputePass();

    pass.setPipeline(this.pipeline);
    pass.setBindGroup(0, d.createBindGroup({
      layout: this.pipeline.getBindGroupLayout(0),
      entries: [
        { binding: 0, resource: { buffer: this.param } },
        { binding: 1, resource: { buffer: this.buf } }
      ]
    }));

    pass.dispatchWorkgroups(32, 32);
    pass.end();

    enc.copyBufferToBuffer(this.buf, 0, this.read, 0, 256*256*4*4);
    d.queue.submit([enc.finish()]);
  }

  async render() {
    if (this.gpu.mode !== "gpu") return;

    await this.read.mapAsync(GPUMapMode.READ);
    const f = new Float32Array(this.read.getMappedRange());

    for (let i = 0; i < f.length; i++) {
      this.image.data[i] = Math.min(255, f[i] * 255);
    }

    this.read.unmap();
    this.ctx.putImageData(this.image, 0, 0);
  }

  apply(cmd) {
    if (cmd.payload?.power) this.power = cmd.payload.power;
  }
}