import { vec2 } from 'gl-matrix';

import {
  createCachedFloat32BufferWrite,
  writeFloat32BufferIfChanged,
} from '../../utils/graphics/cached-buffer-write';
import { smartCompile } from '../../utils/graphics/smart-compile';
import shader from './copy.wgsl?raw';

export class CopyPipeline {
  private static readonly UNIFORM_COUNT = 2;
  private static readonly DEFAULT_SCALE = vec2.fromValues(1, 1);

  private readonly bindGroupLayout: GPUBindGroupLayout;
  private readonly pipeline: GPURenderPipeline;
  private readonly uniforms: GPUBuffer;
  private readonly uniformValues = new Float32Array(CopyPipeline.UNIFORM_COUNT);
  private readonly uniformCache = createCachedFloat32BufferWrite(
    CopyPipeline.UNIFORM_COUNT
  );
  private readonly sampler: GPUSampler;

  private readonly vertexBuffer: GPUBuffer;

  private readonly bindGroupsByInput = new WeakMap<GPUTextureView, GPUBindGroup>();

  public constructor(private readonly device: GPUDevice) {
    this.bindGroupLayout = device.createBindGroupLayout(CopyPipeline.bindGroupLayout);

    this.uniforms = this.device.createBuffer({
      size: CopyPipeline.UNIFORM_COUNT * Float32Array.BYTES_PER_ELEMENT,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });

    this.sampler = this.device.createSampler({
      magFilter: 'linear',
      minFilter: 'linear',
    });

    this.vertexBuffer = device.createBuffer({
      size: 2 * 4 * Float32Array.BYTES_PER_ELEMENT, // 4 x vec2<f32>
      usage: GPUBufferUsage.VERTEX,
      mappedAtCreation: true,
    });
    // prettier-ignore
    const vertexData = [
    // U    V
      0.0, 1.0,
      1.0, 1.0,
      0.0, 0.0,
      1.0, 0.0,
    ];
    new Float32Array(this.vertexBuffer.getMappedRange()).set(vertexData);
    this.vertexBuffer.unmap();

    const shaderModule = smartCompile(device, shader);
    this.pipeline = device.createRenderPipeline({
      layout: device.createPipelineLayout({
        bindGroupLayouts: [this.bindGroupLayout],
      }),
      vertex: {
        module: shaderModule,
        entryPoint: 'vertex',
        buffers: [
          {
            arrayStride: 2 * Float32Array.BYTES_PER_ELEMENT,
            stepMode: 'vertex',
            attributes: [
              {
                shaderLocation: 0,
                offset: 0,
                format: 'float32x2',
              },
            ],
          },
        ],
      },
      fragment: {
        module: shaderModule,
        entryPoint: 'fragment',
        targets: [
          {
            format: 'rgba16float',
          },
        ],
      },
      primitive: {
        topology: 'triangle-strip',
      },
    });
  }

  public execute(
    commandEncoder: GPUCommandEncoder,
    trailMapIn: GPUTextureView,
    trailMapOut: GPUTextureView,
    scale: vec2 = CopyPipeline.DEFAULT_SCALE
  ) {
    this.uniformValues[0] = scale[0];
    this.uniformValues[1] = scale[1];
    writeFloat32BufferIfChanged(
      this.device,
      this.uniforms,
      this.uniformValues,
      this.uniformCache
    );

    const renderPassDescriptor: GPURenderPassDescriptor = {
      colorAttachments: [
        {
          view: trailMapOut,
          loadOp: 'clear',
          storeOp: 'store',
        },
      ],
    };

    const bindGroup = this.getBindGroup(trailMapIn);
    const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
    passEncoder.setPipeline(this.pipeline);
    passEncoder.setBindGroup(0, bindGroup);
    passEncoder.setVertexBuffer(0, this.vertexBuffer);
    passEncoder.draw(4, 1);
    passEncoder.end();
  }

  public destroy() {
    this.vertexBuffer.destroy();
    this.uniforms.destroy();
  }

  private getBindGroup(trailMapIn: GPUTextureView): GPUBindGroup {
    const cached = this.bindGroupsByInput.get(trailMapIn);
    if (cached) {
      return cached;
    }

    const bindGroup = this.device.createBindGroup({
      layout: this.bindGroupLayout,
      entries: [
        {
          binding: 0,
          resource: {
            buffer: this.uniforms,
          },
        },
        {
          binding: 1,
          resource: this.sampler,
        },
        {
          binding: 2,
          resource: trailMapIn,
        },
      ],
    });

    this.bindGroupsByInput.set(trailMapIn, bindGroup);
    return bindGroup;
  }

  private static get bindGroupLayout(): GPUBindGroupLayoutDescriptor {
    return {
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.VERTEX,
          buffer: {
            type: 'uniform',
          },
        },
        {
          binding: 1,
          visibility: GPUShaderStage.FRAGMENT,
          sampler: {
            type: 'filtering',
          },
        },
        {
          binding: 2,
          visibility: GPUShaderStage.FRAGMENT,
          texture: {
            sampleType: 'float',
          },
        },
      ],
    };
  }
}