// In a Vite bundle the .wasm cannot be located next to the glue at runtime, so
// we resolve its URL via Vite's `?url` import and feed it to Emscripten's
// `locateFile` (same pattern qhull-wasm-demo uses for qhull-wasm).
import createMatmulModule, { type MatmulExports } from '../../wasm/dist/matmul.js'
import wasmUrl from '../../wasm/dist/matmul.wasm?url'
import type { MatmulMethod } from './types'

let modulePromise: Promise<MatmulExports> | null = null

function getModule(): Promise<MatmulExports> {
  if (!modulePromise) {
    modulePromise = createMatmulModule({ locateFile: () => wasmUrl })
  }
  return modulePromise
}

async function run(kernel: '_matmul_naive' | '_matmul_blocked', n: number, a: Float64Array, b: Float64Array) {
  const mod = await getModule()
  const bytes = n * n * 8
  const aPtr = mod._malloc(bytes)
  const bPtr = mod._malloc(bytes)
  const cPtr = mod._malloc(bytes)
  try {
    new Float64Array(mod.HEAPF64.buffer, aPtr, n * n).set(a)
    new Float64Array(mod.HEAPF64.buffer, bPtr, n * n).set(b)
    const t0 = performance.now()
    mod[kernel](aPtr, bPtr, cPtr, n)
    const ms = performance.now() - t0
    // Memory may have grown (and the buffer been replaced) during the call,
    // so re-view HEAPF64 rather than reuse a view captured before the call.
    const sample = new Float64Array(mod.HEAPF64.buffer, cPtr, n * n)[0]
    return { ms, sample }
  } finally {
    mod._free(aPtr)
    mod._free(bPtr)
    mod._free(cPtr)
  }
}

export const wasmNaiveMatmul: MatmulMethod = {
  id: 'wasm-naive',
  label: 'WASM (C, naive)',
  precision: 'f64',
  worker: true,
  available: () => typeof WebAssembly !== 'undefined',
  run: (n, a, b) => run('_matmul_naive', n, a, b),
}

export const wasmBlockedMatmul: MatmulMethod = {
  id: 'wasm-blocked',
  label: 'WASM (C, blocked+SIMD)',
  precision: 'f64',
  worker: true,
  available: () => typeof WebAssembly !== 'undefined',
  run: (n, a, b) => run('_matmul_blocked', n, a, b),
}

// The same blocked+SIMD kernel, parallelized over rows with WASM pthreads.
// Runs through threadedClient's classic worker (see blisMatmul for the shared
// mechanism); needs crossOriginIsolated for SharedArrayBuffer.
export const wasmBlockedMtMatmul: MatmulMethod = {
  id: 'wasm-blocked-mt',
  label: 'WASM (C, blocked+SIMD, threaded)',
  precision: 'f64',
  threadedKind: 'matmul-mt',
  note: 'WASM threads — needs cross-origin isolation (SharedArrayBuffer)',
  available: () => typeof WebAssembly !== 'undefined' && globalThis.crossOriginIsolated === true,
}