import { getPath } from "./common.js";

class Compute {
  constructor(device, canvasFormat, wgsl)
  {
    const label = "compute";
    this.label = label;

    //////////////////////////////////////////////////////////////////////
    // shader module
    //////////////////////////////////////////////////////////////////////

    const shaderModule = device.createShaderModule({
      label: `${label} shader`,
      code: wgsl,
    });

    //////////////////////////////////////////////////////////////////////
    // bind group layout
    //////////////////////////////////////////////////////////////////////

    this.bindGroupLayout = {
      compute: device.createBindGroupLayout({
        label: `${label} bind group layout`,
        entries: [
          {
            binding: 0,
            visibility: GPUShaderStage.COMPUTE,
            storageTexture: {
              access: "write-only",
              format: canvasFormat,
              viewDimension: "2d",
            }
          }
        ]
      }),
    };

    //////////////////////////////////////////////////////////////////////
    // pipeline
    //////////////////////////////////////////////////////////////////////

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

    this.computePipeline = device.createComputePipeline({
      label: `${label} compute pipeline`,
      layout: pipelineLayout,
      compute: {
        module: shaderModule,
      },
    });

    this.textures = undefined;
  }

  recreateTextures(device, canvasTexture)
  {
    const label = this.label;

    if (this.textures !== undefined && this.textures[0].width == canvasTexture.width && this.textures[0].height == canvasTexture.height) {
      return;
    }
    if (this.textures !== undefined) {
      for (let i = 0; i < 2; i++) {
        this.textures[i].destroy();
      }
    }

    //////////////////////////////////////////////////////////////////////
    // textures
    //////////////////////////////////////////////////////////////////////

    this.textures = [];
    for (let i = 0; i < 2; i++) {
      const texture = device.createTexture({
        label: `${label} texture ${i}`,
        format: canvasTexture.format,
        usage: GPUTextureUsage.STORAGE_BINDING | GPUTextureUsage.COPY_SRC,
        size: [canvasTexture.width, canvasTexture.height],
      });
      this.textures.push(texture);
    }

    //////////////////////////////////////////////////////////////////////
    // bind group
    //////////////////////////////////////////////////////////////////////

    this.bindGroups = []
    for (let i = 0; i < 2; i++) {
      const bindGroup = device.createBindGroup({
        label: `${label} bind group ${i}`,
        layout: this.bindGroupLayout.compute,
        entries: [{
          binding: 0,
          resource: this.textures[i].createView(),
        }],
      });
      this.bindGroups.push(bindGroup);
    }
  }

  render(device, frameNumber, canvasTexture)
  {
    this.recreateTextures(device, canvasTexture);

    const label = this.label;
    const encoder = device.createCommandEncoder({ label: `${label} command encoder` });
    const computePass = encoder.beginComputePass({ label: `${label} compute pass` });

    computePass.setPipeline(this.computePipeline);
    computePass.setBindGroup(0, this.bindGroups[frameNumber]);
    computePass.dispatchWorkgroups(canvasTexture.width / 16, canvasTexture.height / 16);
    computePass.end();

    encoder.copyTextureToTexture({
      texture: this.textures[frameNumber],
    }, {
      texture: canvasTexture
    }, {
      width: canvasTexture.width,
      height: canvasTexture.height,
      depthOrArrayLayers: 1,
    });

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

async function loadCompute(device, canvasFormat)
{
  const canvasStorage = `@group(0) @binding(0) var out: texture_storage_2d<${canvasFormat}, write>;\n`
  const wgsl = canvasStorage + await getPath("compute.wgsl");

  const compute = new Compute(device, canvasFormat, wgsl);

  return compute;
}

export { loadCompute };