import {
  createCachedFloat32BufferWrite,
  writeFloat32BufferIfChanged,
} from '../../utils/graphics/cached-buffer-write';
import { setUpFullScreenQuad } from '../../utils/graphics/full-screen-quad';
import { smartCompile } from '../../utils/graphics/smart-compile';
import { rgbChannelToUnit, type RgbColor } from '../../utils/rgb-color';
import { CommonState } from '../common-state/common-state';
import { RenderSettings } from './render-settings';
import shader from './render.wgsl?raw';

export class RenderPipeline {
  private static readonly UNIFORM_COUNT = 20;

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

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

  public constructor(
    private readonly context: GPUCanvasContext,
    private readonly device: GPUDevice,
    private readonly commonState: CommonState
  ) {
    this.bindGroupLayout = device.createBindGroupLayout(RenderPipeline.bindGroupLayout);

    const vertex = setUpFullScreenQuad(device);

    const format = navigator.gpu.getPreferredCanvasFormat();
    this.pipeline = this.createPipeline(format, vertex, 'fragment');
    this.noSourcePipeline = this.createPipeline(format, vertex, 'fragmentNoSource');

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

  private createPipeline(
    format: GPUTextureFormat,
    vertex: GPUVertexState,
    fragmentEntryPoint: string
  ): GPURenderPipeline {
    return this.device.createRenderPipeline({
      layout: this.device.createPipelineLayout({
        bindGroupLayouts: [this.commonState.bindGroupLayout, this.bindGroupLayout],
      }),
      vertex,
      fragment: {
        module: smartCompile(this.device, CommonState.shaderCode, shader),
        entryPoint: fragmentEntryPoint,
        targets: [
          {
            format,
          },
        ],
      },
      primitive: {
        topology: 'triangle-list',
      },
    });
  }

  public setParameters({
    channelColors,
    backgroundColor,
    clarity,
    renderTraceNormalizationFloor,
    renderBrushColorBase,
    renderBrushColorStrengthMultiplier,
    backgroundGrainStrength,
  }: RenderSettings & {
    channelColors: [RgbColor, RgbColor, RgbColor];
    backgroundColor: RgbColor;
  }) {
    const [a, b, c] = channelColors;
    this.uniformValues[0] = rgbChannelToUnit(a[0]);
    this.uniformValues[1] = rgbChannelToUnit(a[1]);
    this.uniformValues[2] = rgbChannelToUnit(a[2]);
    this.uniformValues[3] = 0;
    this.uniformValues[4] = rgbChannelToUnit(b[0]);
    this.uniformValues[5] = rgbChannelToUnit(b[1]);
    this.uniformValues[6] = rgbChannelToUnit(b[2]);
    this.uniformValues[7] = 0;
    this.uniformValues[8] = rgbChannelToUnit(c[0]);
    this.uniformValues[9] = rgbChannelToUnit(c[1]);
    this.uniformValues[10] = rgbChannelToUnit(c[2]);
    this.uniformValues[11] = 0;
    this.uniformValues[12] = rgbChannelToUnit(backgroundColor[0]);
    this.uniformValues[13] = rgbChannelToUnit(backgroundColor[1]);
    this.uniformValues[14] = rgbChannelToUnit(backgroundColor[2]);
    this.uniformValues[15] = clarity;
    this.uniformValues[16] = renderTraceNormalizationFloor;
    this.uniformValues[17] = renderBrushColorBase;
    this.uniformValues[18] = renderBrushColorStrengthMultiplier;
    this.uniformValues[19] = backgroundGrainStrength;
    writeFloat32BufferIfChanged(
      this.device,
      this.uniforms,
      this.uniformValues,
      this.uniformCache
    );
  }

  public execute(
    commandEncoder: GPUCommandEncoder,
    colorTexture: GPUTextureView,
    sourceTexture: GPUTextureView,
    useSourceTexture = true
  ): GPUTexture {
    const bindGroup = this.getBindGroup(colorTexture, sourceTexture);
    const canvasTexture = this.context.getCurrentTexture();

    const renderPassDescriptor: GPURenderPassDescriptor = {
      colorAttachments: [
        {
          view: canvasTexture.createView(),
          clearValue: { r: 0, g: 0, b: 0, a: 1 },
          loadOp: 'clear',
          storeOp: 'store',
        },
      ],
    };
    const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
    passEncoder.setPipeline(useSourceTexture ? this.pipeline : this.noSourcePipeline);
    this.commonState.execute(passEncoder);
    passEncoder.setBindGroup(1, bindGroup);
    passEncoder.draw(3, 1);
    passEncoder.end();

    return canvasTexture;
  }

  public executeToView(
    commandEncoder: GPUCommandEncoder,
    colorTexture: GPUTextureView,
    sourceTexture: GPUTextureView,
    outputTexture: GPUTextureView
  ) {
    const bindGroup = this.getBindGroup(colorTexture, sourceTexture);

    const passEncoder = commandEncoder.beginRenderPass({
      colorAttachments: [
        {
          view: outputTexture,
          clearValue: { r: 0, g: 0, b: 0, a: 1 },
          loadOp: 'clear',
          storeOp: 'store',
        },
      ],
    });
    passEncoder.setPipeline(this.pipeline);
    this.commonState.execute(passEncoder);
    passEncoder.setBindGroup(1, bindGroup);
    passEncoder.draw(3, 1);
    passEncoder.end();
  }

  private getBindGroup(
    colorTexture: GPUTextureView,
    sourceTexture: GPUTextureView
  ): GPUBindGroup {
    let sourceTextureCache = this.bindGroupsByTexture.get(colorTexture);
    if (!sourceTextureCache) {
      sourceTextureCache = new WeakMap<GPUTextureView, GPUBindGroup>();
      this.bindGroupsByTexture.set(colorTexture, sourceTextureCache);
    }

    const cached = sourceTextureCache.get(sourceTexture);
    if (cached) {
      return cached;
    }

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

    sourceTextureCache.set(sourceTexture, bindGroup);
    return bindGroup;
  }

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

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