// WebAssembly SIMD GEMM for the onnx-lite interpreter: X[m, nin] @ W[nin, nout].
// Weights are copied into wasm memory once (keyed by the weight's data array).

import { MATVEC_WASM_BASE64 } from './kernel-wasm.ts';

export interface KernelExports {
  matvec(w: number, x: number, y: number, nin: number, nout: number): void;
  gemm(w: number, x: number, y: number, m: number, nin: number, nout: number): void;
}

export function wasmSimdSupported(): boolean {
  const probe = new Uint8Array([0, 97, 115, 109, 1, 0, 0, 0, 1, 5, 1, 96, 0, 1, 123, 3, 2, 1, 0, 10, 10, 1, 8, 0, 65, 0, 253, 15, 253, 98, 11]);
  try {
    return typeof WebAssembly === 'object' && WebAssembly.validate(probe);
  } catch {
    return false;
  }
}

export function instantiateKernel(memory: WebAssembly.Memory): KernelExports {
  const bytes = Uint8Array.from(atob(MATVEC_WASM_BASE64), (c) => c.charCodeAt(0));
  return new WebAssembly.Instance(new WebAssembly.Module(bytes), { env: { memory } }).exports as unknown as KernelExports;
}

export class WasmGemm {
  memory: WebAssembly.Memory;
  kernel: KernelExports;
  top: number;
  weights: WeakMap<Float32Array, number>;
  scratch: { off: number; bytes: number };
  constructor() {
    this.memory = new WebAssembly.Memory({ initial: 64 });
    this.kernel = instantiateKernel(this.memory);
    this.top = 0;
    this.weights = new WeakMap();
    this.scratch = { off: 0, bytes: 0 };
  }
  #reserve(bytes: number): number {
    const off = (this.top + 15) & ~15;
    const end = off + bytes;
    const have = this.memory.buffer.byteLength;
    if (end > have) this.memory.grow(Math.ceil((end - have) / 65536));
    this.top = end;
    return off;
  }
  // Returns false when the shape is unsupported (caller falls back to JS).
  supports(nin: number, nout: number): boolean {
    return nin % 2 === 0 && nout % 4 === 0;
  }
  // Like matmul, but B is transient (copied each call instead of cached).
  matmulTransient(A: Float32Array, aOff: number, m: number, B: Float32Array, bOff: number, nin: number, nout: number, out: Float32Array, outOff: number): void {
    const need = (m * nin + nin * nout + m * nout) * 4;
    if (this.scratch.bytes < need) this.scratch = { off: this.#reserve(need * 2), bytes: need * 2 };
    const x = this.scratch.off;
    const w = x + m * nin * 4;
    const y = w + nin * nout * 4;
    const buf = this.memory.buffer;
    new Float32Array(buf, x, m * nin).set(A.subarray(aOff, aOff + m * nin));
    new Float32Array(buf, w, nin * nout).set(B.subarray(bOff, bOff + nin * nout));
    this.kernel.gemm(w, x, y, m, nin, nout);
    out.set(new Float32Array(buf, y, m * nout), outOff);
  }
  matmul(A: Float32Array, m: number, B: Float32Array, nin: number, nout: number, out: Float32Array): void {
    let w = this.weights.get(B);
    if (w === undefined) {
      w = this.#reserve(B.length * 4);
      new Float32Array(this.memory.buffer, w, B.length).set(B);
      this.weights.set(B, w);
      this.scratch.bytes = 0; // weights were appended after scratch; re-reserve scratch
    }
    const need = (m * nin + m * nout) * 4;
    if (this.scratch.bytes < need) {
      this.scratch = { off: this.#reserve(need * 2), bytes: need * 2 };
    }
    const x = this.scratch.off;
    const y = x + m * nin * 4;
    new Float32Array(this.memory.buffer, x, m * nin).set(A.subarray(0, m * nin));
    this.kernel.gemm(w, x, y, m, nin, nout);
    out.set(new Float32Array(this.memory.buffer, y, m * nout));
  }
}