import shaderSource from './text.wgsl?raw'
import type { SlugData } from './slug-generator'
import type { TextLayoutResult } from './text-layout'

const TEXT_UNIFORM_FLOATS = 8 // resolution(2) + pad(2) + color(4)
const TEXT_UNIFORM_BYTES = TEXT_UNIFORM_FLOATS * 4

const QUAD_VERTICES = new Float32Array([
  -1.0, -1.0,
  -1.0, 1.0,
  1.0, 1.0,
  -1.0, -1.0,
  1.0, 1.0,
  1.0, -1.0
])

const MAX_GLYPHS = 4096
const INSTANCE_STRIDE = 4 * 4

export class TextRenderer {
  private device!: GPUDevice
  private pipeline!: GPURenderPipeline
  private bindGroup!: GPUBindGroup
  private uniformBuffer!: GPUBuffer
  private uniformData = new Float32Array(TEXT_UNIFORM_FLOATS)
  private quadBuffer!: GPUBuffer

  private scaleBiasBuffer!: GPUBuffer
  private glyphBandScaleBuffer!: GPUBuffer
  private bandMaxTexCoordsBuffer!: GPUBuffer

  private curvesTex!: GPUTexture
  private bandsTex!: GPUTexture
  private glyphCount = 0
  private initialized = false

  get ready(): boolean {
    return this.initialized
  }

  init(device: GPUDevice, format: GPUTextureFormat, slugData: SlugData): void {
    this.device = device

    this.createTextures(slugData)

    const module = device.createShaderModule({ code: shaderSource })

    const bindGroupLayout = device.createBindGroupLayout({
      entries: [
        { binding: 0, visibility: GPUShaderStage.VERTEX | GPUShaderStage.FRAGMENT, buffer: { type: 'uniform' } },
        { binding: 1, visibility: GPUShaderStage.FRAGMENT, texture: { sampleType: 'unfilterable-float' } },
        { binding: 2, visibility: GPUShaderStage.FRAGMENT, texture: { sampleType: 'uint' } }
      ]
    })

    const pipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] })

    this.pipeline = device.createRenderPipeline({
      layout: pipelineLayout,
      vertex: {
        module,
        entryPoint: 'vs_main',
        buffers: [
          {
            arrayStride: 2 * 4,
            attributes: [{ shaderLocation: 0, offset: 0, format: 'float32x2' }]
          },
          {
            arrayStride: INSTANCE_STRIDE,
            stepMode: 'instance',
            attributes: [{ shaderLocation: 1, offset: 0, format: 'float32x4' }]
          },
          {
            arrayStride: INSTANCE_STRIDE,
            stepMode: 'instance',
            attributes: [{ shaderLocation: 2, offset: 0, format: 'float32x4' }]
          },
          {
            arrayStride: INSTANCE_STRIDE,
            stepMode: 'instance',
            attributes: [{ shaderLocation: 3, offset: 0, format: 'uint32x4' }]
          }
        ]
      },
      fragment: {
        module,
        entryPoint: 'fs_main',
        targets: [{
          format,
          blend: {
            color: { srcFactor: 'src-alpha', dstFactor: 'one-minus-src-alpha', operation: 'add' },
            alpha: { srcFactor: 'one', dstFactor: 'one-minus-src-alpha', operation: 'add' }
          }
        }]
      },
      primitive: { topology: 'triangle-list' }
    })

    this.uniformBuffer = device.createBuffer({
      size: TEXT_UNIFORM_BYTES,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
    })

    this.quadBuffer = device.createBuffer({
      size: QUAD_VERTICES.byteLength,
      usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    })
    device.queue.writeBuffer(this.quadBuffer, 0, QUAD_VERTICES)

    this.scaleBiasBuffer = device.createBuffer({
      size: MAX_GLYPHS * INSTANCE_STRIDE,
      usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    })
    this.glyphBandScaleBuffer = device.createBuffer({
      size: MAX_GLYPHS * INSTANCE_STRIDE,
      usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    })
    this.bandMaxTexCoordsBuffer = device.createBuffer({
      size: MAX_GLYPHS * INSTANCE_STRIDE,
      usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST
    })

    this.bindGroup = device.createBindGroup({
      layout: bindGroupLayout,
      entries: [
        { binding: 0, resource: { buffer: this.uniformBuffer } },
        { binding: 1, resource: this.curvesTex.createView() },
        { binding: 2, resource: this.bandsTex.createView() }
      ]
    })

    this.initialized = true
  }

  private createTextures(slugData: SlugData): void {
    this.curvesTex = this.device.createTexture({
      size: { width: slugData.curvesWidth, height: slugData.curvesHeight },
      format: 'rgba32float',
      usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST
    })
    this.device.queue.writeTexture(
      { texture: this.curvesTex },
      slugData.curvesData.buffer as BufferSource,
      { bytesPerRow: slugData.curvesWidth * 4 * 4 },
      { width: slugData.curvesWidth, height: slugData.curvesHeight }
    )

    this.bandsTex = this.device.createTexture({
      size: { width: slugData.bandsWidth, height: slugData.bandsHeight },
      format: 'rg32uint',
      usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST
    })
    this.device.queue.writeTexture(
      { texture: this.bandsTex },
      slugData.bandsData.buffer as BufferSource,
      { bytesPerRow: slugData.bandsWidth * 2 * 4 },
      { width: slugData.bandsWidth, height: slugData.bandsHeight }
    )
  }

  updateText(layout: TextLayoutResult): void {
    if (!this.initialized) return
    const count = Math.min(layout.glyphCount, MAX_GLYPHS)
    this.glyphCount = count
    if (count === 0) return

    const sb = layout.scaleBias.subarray(0, count * 4)
    this.device.queue.writeBuffer(this.scaleBiasBuffer, 0, sb.buffer as BufferSource, sb.byteOffset, sb.byteLength)
    const gb = layout.glyphBandScale.subarray(0, count * 4)
    this.device.queue.writeBuffer(this.glyphBandScaleBuffer, 0, gb.buffer as BufferSource, gb.byteOffset, gb.byteLength)
    const bm = layout.bandMaxTexCoords.subarray(0, count * 4)
    this.device.queue.writeBuffer(this.bandMaxTexCoordsBuffer, 0, bm.buffer as BufferSource, bm.byteOffset, bm.byteLength)
  }

  renderPass(
    encoder: GPUCommandEncoder,
    canvasView: GPUTextureView,
    resolution: Float32Array,
    textColor: Float32Array
  ): void {
    if (!this.initialized || this.glyphCount === 0) return

    this.uniformData[0] = resolution[0]
    this.uniformData[1] = resolution[1]
    this.uniformData[4] = textColor[0]
    this.uniformData[5] = textColor[1]
    this.uniformData[6] = textColor[2]
    this.uniformData[7] = textColor[3]
    this.device.queue.writeBuffer(this.uniformBuffer, 0, this.uniformData)

    const pass = encoder.beginRenderPass({
      colorAttachments: [
        { view: canvasView, loadOp: 'load', storeOp: 'store' }
      ]
    })
    pass.setPipeline(this.pipeline)
    pass.setBindGroup(0, this.bindGroup)
    pass.setVertexBuffer(0, this.quadBuffer)
    pass.setVertexBuffer(1, this.scaleBiasBuffer)
    pass.setVertexBuffer(2, this.glyphBandScaleBuffer)
    pass.setVertexBuffer(3, this.bandMaxTexCoordsBuffer)
    pass.draw(6, this.glyphCount)
    pass.end()
  }
}