import type { MatmulMethod } from './types'

const TILE = 16

// Tiled sgemm: each workgroup computes one TILE×TILE block of C, staging
// TILE×TILE tiles of A and B through workgroup (shared) memory.
const SHADER = /* wgsl */ `
struct Dims { n: u32 }
@group(0) @binding(0) var<uniform> dims: Dims;
@group(0) @binding(1) var<storage, read> A: array<f32>;
@group(0) @binding(2) var<storage, read> B: array<f32>;
@group(0) @binding(3) var<storage, read_write> C: array<f32>;

var<workgroup> tileA: array<array<f32, ${TILE}>, ${TILE}>;
var<workgroup> tileB: array<array<f32, ${TILE}>, ${TILE}>;

@compute @workgroup_size(${TILE}, ${TILE})
fn main(
  @builtin(global_invocation_id) gid: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>,
) {
  let n = dims.n;
  let row = gid.y;
  let col = gid.x;
  var sum: f32 = 0.0;
  let numTiles = (n + ${TILE}u - 1u) / ${TILE}u;

  for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
    let aCol = t * ${TILE}u + lid.x;
    let bRow = t * ${TILE}u + lid.y;
    tileA[lid.y][lid.x] = select(0.0, A[row * n + aCol], row < n && aCol < n);
    tileB[lid.y][lid.x] = select(0.0, B[bRow * n + col], bRow < n && col < n);
    workgroupBarrier();
    for (var k: u32 = 0u; k < ${TILE}u; k = k + 1u) {
      sum = sum + tileA[lid.y][k] * tileB[k][lid.x];
    }
    workgroupBarrier();
  }

  if (row < n && col < n) {
    C[row * n + col] = sum;
  }
}
`

let devicePromise: Promise<GPUDevice> | null = null

function getDevice(): Promise<GPUDevice> {
  if (!devicePromise) {
    devicePromise = (async () => {
      if (!navigator.gpu) throw new Error('WebGPU not supported')
      const adapter = await navigator.gpu.requestAdapter()
      if (!adapter) throw new Error('No WebGPU adapter available')
      return adapter.requestDevice()
    })()
  }
  return devicePromise
}

let pipelinePromise: Promise<{ device: GPUDevice; pipeline: GPUComputePipeline }> | null = null

function getPipeline() {
  if (!pipelinePromise) {
    pipelinePromise = getDevice().then((device) => {
      const module = device.createShaderModule({ code: SHADER })
      const pipeline = device.createComputePipeline({
        layout: 'auto',
        compute: { module, entryPoint: 'main' },
      })
      return { device, pipeline }
    })
  }
  return pipelinePromise
}

export const webgpuMatmul: MatmulMethod = {
  id: 'webgpu',
  label: 'WebGPU',
  precision: 'f32',
  note: 'single precision only — WGSL has no f64',
  available: () => typeof navigator !== 'undefined' && !!navigator.gpu,
  async run(n, af64, bf64) {
    const { device, pipeline } = await getPipeline()

    const a = Float32Array.from(af64)
    const b = Float32Array.from(bf64)
    const bytes = n * n * 4

    const dimsBuf = device.createBuffer({ size: 4, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST })
    device.queue.writeBuffer(dimsBuf, 0, new Uint32Array([n]))

    const aBuf = device.createBuffer({ size: bytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST })
    const bBuf = device.createBuffer({ size: bytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST })
    const cBuf = device.createBuffer({ size: bytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC })
    const readBuf = device.createBuffer({ size: bytes, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ })

    device.queue.writeBuffer(aBuf, 0, a)
    device.queue.writeBuffer(bBuf, 0, b)

    const bindGroup = device.createBindGroup({
      layout: pipeline.getBindGroupLayout(0),
      entries: [
        { binding: 0, resource: { buffer: dimsBuf } },
        { binding: 1, resource: { buffer: aBuf } },
        { binding: 2, resource: { buffer: bBuf } },
        { binding: 3, resource: { buffer: cBuf } },
      ],
    })

    const t0 = performance.now()
    const encoder = device.createCommandEncoder()
    const pass = encoder.beginComputePass()
    pass.setPipeline(pipeline)
    pass.setBindGroup(0, bindGroup)
    const groups = Math.ceil(n / TILE)
    pass.dispatchWorkgroups(groups, groups)
    pass.end()
    encoder.copyBufferToBuffer(cBuf, 0, readBuf, 0, bytes)
    device.queue.submit([encoder.finish()])

    await readBuf.mapAsync(GPUMapMode.READ)
    // Timing includes the GPU->CPU readback, since that's what a real caller pays.
    const ms = performance.now() - t0
    const sample = new Float32Array(readBuf.getMappedRange())[0]
    readBuf.unmap()

    dimsBuf.destroy()
    aBuf.destroy()
    bBuf.destroy()
    cBuf.destroy()
    readBuf.destroy()

    return { ms, sample }
  },
}