// WebGPU host for the megakernel: device setup (Dawn in Node via the
// optional `webgpu` package; in a browser this maps onto navigator.gpu),
// buffer upload, the dispatch queue, and readback.

export async function installNodeWebGpu() {
  const mod = await import('webgpu');
  Object.assign(globalThis, mod.globals);
  const dawnFlags = process.env.DAWN_FLAGS?.split(';').filter(Boolean) ?? [];
  Object.defineProperty(globalThis, 'navigator', {
    value: { gpu: mod.create(dawnFlags) },
    configurable: true,
    writable: true,
  });
}

export async function requestDevice() {
  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) throw new Error('no WebGPU adapter');
  const L = adapter.limits;
  const device = await adapter.requestDevice({
    requiredLimits: {
      maxComputeInvocationsPerWorkgroup: L.maxComputeInvocationsPerWorkgroup,
      maxComputeWorkgroupSizeX: L.maxComputeWorkgroupSizeX,
      maxComputeWorkgroupStorageSize: L.maxComputeWorkgroupStorageSize,
    },
  });
  device.addEventListener?.('uncapturederror', (e) => {
    console.error('WebGPU uncaptured error:', e.error?.message ?? e);
  });
  const info = adapter.info
    ? `${adapter.info.vendor ?? '?'} ${adapter.info.architecture ?? ''} ${adapter.info.description ?? ''}`.trim()
    : 'unknown adapter';
  return { device, info };
}

// Multi-dispatch backend: 8 pipelines, one bind group with a dynamic-offset
// uniform carrying the within-chunk step index; the host writes the chunk
// base into state cell 0 before each chunk and keeps two chunks in flight.
export async function runMultiDispatch(device, L, stages, onProgress) {
  const nsteps = L.nsteps;
  const bgl = device.createBindGroupLayout({
    entries: [
      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
      { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
      { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform', hasDynamicOffset: true, minBindingSize: 16 } },
    ],
  });
  const layout = device.createPipelineLayout({ bindGroupLayouts: [bgl] });

  const pipelines = [];
  const simpleBGs = [];
  for (const name of stages.order) {
    const module = device.createShaderModule({ code: stages.sources[name] });
    const errs = (await module.getCompilationInfo()).messages.filter((m) => m.type === 'error');
    if (errs.length) {
      throw new Error(`stage ${name} WGSL errors:\n` +
        errs.map((m) => `  ${m.lineNum}:${m.linePos} ${m.message}`).join('\n'));
    }
    pipelines.push(device.createComputePipeline({ layout, compute: { module, entryPoint: 'main' } }));
  }

  const chunk = Number(process.env.CHUNK ?? 1024);
  const mkBuf = (size, usage) => device.createBuffer({ size, usage });
  const opsBuf = mkBuf(L.opsData.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
  device.queue.writeBuffer(opsBuf, 0, L.opsData);
  const stBuf = mkBuf(L.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST);
  const outBuf = mkBuf(2 * nsteps * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
  const uniBuf = mkBuf(chunk * 256, GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
  {
    const uni = new Uint32Array(chunk * 64);
    for (let s = 0; s < chunk; s++) uni[s * 64] = s;
    device.queue.writeBuffer(uniBuf, 0, uni);
  }
  const readOut = mkBuf(2 * nsteps * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
  const readSt = mkBuf(L.stLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);

  const bindGroup = device.createBindGroup({
    layout: bgl,
    entries: [
      { binding: 0, resource: { buffer: opsBuf } },
      { binding: 1, resource: { buffer: stBuf } },
      { binding: 2, resource: { buffer: outBuf } },
      { binding: 3, resource: { buffer: uniBuf, offset: 0, size: 16 } },
    ],
  });

  const t0 = performance.now();
  const pending = [];
  let kbase = 0;
  let recMs = 0;
  const inflight = Number(process.env.INFLIGHT ?? 2);
  while (kbase < nsteps) {
    const n = Math.min(chunk, nsteps - kbase);
    device.queue.writeBuffer(stBuf, 0, new Float32Array([kbase]));
    const tr = performance.now();
    const enc = device.createCommandEncoder();
    const pass = enc.beginComputePass();

    for (let s = 0; s < n; s++) {
      pass.setBindGroup(0, bindGroup, [s * 256]);
      for (const p of pipelines) {
        pass.setPipeline(p);
        pass.dispatchWorkgroups(stages.nwg);
      }
    }
    pass.end();
    const cb = enc.finish();
    recMs += performance.now() - tr;
    device.queue.submit([cb]);
    pending.push(device.queue.onSubmittedWorkDone());
    if (pending.length >= inflight) await pending.shift();
    kbase += n;
    if (onProgress) onProgress(kbase, nsteps, performance.now() - t0);
  }
  while (pending.length) await pending.shift();
  const gpuMs = performance.now() - t0;
  console.log(`\n  [recording ${(1000 * recMs / nsteps).toFixed(1)} us/step of ${(1000 * gpuMs / nsteps).toFixed(1)} wall]`);

  const enc = device.createCommandEncoder();
  enc.copyBufferToBuffer(outBuf, 0, readOut, 0, 2 * nsteps * 4);
  enc.copyBufferToBuffer(stBuf, 0, readSt, 0, L.stLen * 4);
  device.queue.submit([enc.finish()]);
  await readOut.mapAsync(GPUMapMode.READ);
  await readSt.mapAsync(GPUMapMode.READ);
  const outData = new Float32Array(readOut.getMappedRange().slice(0));
  const stData = new Float32Array(readSt.getMappedRange().slice(0));
  readOut.unmap();
  readSt.unmap();
  for (const b of [opsBuf, stBuf, outBuf, uniBuf, readOut, readSt]) b.destroy();

  return {
    acc: outData.subarray(0, nsteps),
    fhist: outData.subarray(nsteps, 2 * nsteps),
    state: stData,
    gpuMs,
  };
}


// Persistent-kernel backend: identical dispatches, each integrating CHUNK
// steps with global spin barriers between stages. sync[] holds the barrier
// atomics; state cell 0 (kbase) advances GPU-side at the end of each chunk.
export async function runPersistent(device, L, pk, onProgress) {
  const nsteps = L.nsteps;
  const module = device.createShaderModule({ code: pk.code });
  const errs = (await module.getCompilationInfo()).messages.filter((m) => m.type === 'error');
  if (errs.length) {
    throw new Error('persistent kernel WGSL errors:\n' +
      errs.map((m) => `  ${m.lineNum}:${m.linePos} ${m.message}`).join('\n'));
  }
  const pipeline = device.createComputePipeline({ layout: 'auto', compute: { module, entryPoint: 'main' } });

  const mkBuf = (size, usage) => device.createBuffer({ size, usage });
  const opsBuf = mkBuf(L.opsData.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
  device.queue.writeBuffer(opsBuf, 0, L.opsData);
  const stBuf = mkBuf(L.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
  const outBuf = mkBuf(2 * nsteps * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
  const syncBuf = mkBuf(16, GPUBufferUsage.STORAGE);
  const readOut = mkBuf(2 * nsteps * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
  const readSt = mkBuf(L.stLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);

  const bindGroup = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: { buffer: opsBuf } },
      { binding: 1, resource: { buffer: stBuf } },
      { binding: 2, resource: { buffer: outBuf } },
      { binding: 3, resource: { buffer: syncBuf } },
    ],
  });

  const totalDispatches = Math.ceil(nsteps / pk.chunk);
  const perSubmit = 8;
  const t0 = performance.now();
  let done = 0;
  const pending = [];
  while (done < totalDispatches) {
    const n = Math.min(perSubmit, totalDispatches - done);
    const enc = device.createCommandEncoder();
    for (let i = 0; i < n; i++) {
      const pass = enc.beginComputePass();
      pass.setPipeline(pipeline);
      pass.setBindGroup(0, bindGroup);
      pass.dispatchWorkgroups(pk.nwg);
      pass.end();
    }
    device.queue.submit([enc.finish()]);
    pending.push(device.queue.onSubmittedWorkDone());
    if (pending.length >= 2) await pending.shift();
    done += n;
    if (onProgress) onProgress(Math.min(done * pk.chunk, nsteps), nsteps, performance.now() - t0);
  }
  while (pending.length) await pending.shift();
  const gpuMs = performance.now() - t0;

  const enc = device.createCommandEncoder();
  enc.copyBufferToBuffer(outBuf, 0, readOut, 0, 2 * nsteps * 4);
  enc.copyBufferToBuffer(stBuf, 0, readSt, 0, L.stLen * 4);
  device.queue.submit([enc.finish()]);
  await readOut.mapAsync(GPUMapMode.READ);
  await readSt.mapAsync(GPUMapMode.READ);
  const outData = new Float32Array(readOut.getMappedRange().slice(0));
  const stData = new Float32Array(readSt.getMappedRange().slice(0));
  readOut.unmap();
  readSt.unmap();
  for (const b of [opsBuf, stBuf, outBuf, syncBuf, readOut, readSt]) b.destroy();

  return {
    acc: outData.subarray(0, nsteps),
    fhist: outData.subarray(nsteps, 2 * nsteps),
    state: stData,
    gpuMs,
  };
}

export async function runMegakernel(device, kernel, nsteps, onProgress) {
  const module = device.createShaderModule({ code: kernel.code });
  const compInfo = await module.getCompilationInfo();
  const errors = compInfo.messages.filter((m) => m.type === 'error');
  if (errors.length) {
    throw new Error(
      'WGSL compile errors:\n' +
        errors.map((m) => `  ${m.lineNum}:${m.linePos} ${m.message}`).join('\n'),
    );
  }

  const pipeline = device.createComputePipeline({
    layout: 'auto',
    compute: { module, entryPoint: 'main' },
  });

  const mkBuf = (size, usage) => device.createBuffer({ size, usage });
  const opsBuf = mkBuf(kernel.opsData.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
  device.queue.writeBuffer(opsBuf, 0, kernel.opsData);
  const stBuf = mkBuf(kernel.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
  const outBuf = mkBuf(2 * nsteps * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
  const readOut = mkBuf(2 * nsteps * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
  const readSt = mkBuf(kernel.stLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);

  const bindGroup = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: { buffer: opsBuf } },
      { binding: 1, resource: { buffer: stBuf } },
      { binding: 2, resource: { buffer: outBuf } },
    ],
  });

  const totalDispatches = Math.ceil(nsteps / kernel.chunk);
  const perSubmit = 16;
  const t0 = performance.now();
  let done = 0;
  while (done < totalDispatches) {
    const n = Math.min(perSubmit, totalDispatches - done);
    const enc = device.createCommandEncoder();
    for (let i = 0; i < n; i++) {
      const pass = enc.beginComputePass();
      pass.setPipeline(pipeline);
      pass.setBindGroup(0, bindGroup);
      pass.dispatchWorkgroups(1);
      pass.end();
    }
    device.queue.submit([enc.finish()]);
    await device.queue.onSubmittedWorkDone();
    done += n;
    if (onProgress) onProgress(Math.min(done * kernel.chunk, nsteps), nsteps, performance.now() - t0);
  }
  const gpuMs = performance.now() - t0;

  const enc = device.createCommandEncoder();
  enc.copyBufferToBuffer(outBuf, 0, readOut, 0, 2 * nsteps * 4);
  enc.copyBufferToBuffer(stBuf, 0, readSt, 0, kernel.stLen * 4);
  device.queue.submit([enc.finish()]);
  await readOut.mapAsync(GPUMapMode.READ);
  await readSt.mapAsync(GPUMapMode.READ);
  const outData = new Float32Array(readOut.getMappedRange().slice(0));
  const stData = new Float32Array(readSt.getMappedRange().slice(0));
  readOut.unmap();
  readSt.unmap();
  for (const b of [opsBuf, stBuf, outBuf, readOut, readSt]) b.destroy();

  return {
    acc: outData.subarray(0, nsteps),
    fhist: outData.subarray(nsteps, 2 * nsteps),
    state: stData,
    gpuMs,
  };
}