/**
 * The generic runner: everything the host does with an instrument file's
 * WGSL, with no knowledge of what the shaders compute.
 *
 * A phase (precompute or render) runs as: create its f32 storage buffers at
 * the sizes the derived expressions give, upload any named inputs, build
 * one compute pipeline per shader with the scope values the shader lists as
 * override constants, run the setup passes, run the loop, and read back the
 * requested buffers. Two conveniences the schedules lean on:
 *
 *  - `__tick`: a u32 the host zeroes before the loop and bumps after every
 *    iteration, so a shader knows which timestep or tile it is in without
 *    any per-iteration upload.
 *  - `rotate`: buffer groups cycled one place per iteration, which is how
 *    a leapfrog's state histories advance without copies. With a group
 *    [p, pm, pn], the buffer a pass sees as `p` at iteration k is the one
 *    it saw as `pn` at iteration k-1.
 *
 * Iterations are encoded in batches with an await between them, which
 * bounds how far the queue runs ahead and paces progress and cancellation.
 */
import { evalExpr } from '../expr.ts';
import type { PassSpec, PhaseSpec, ShaderSpec } from '../format.ts';

export const WORKGROUP_SIZE = 64;

/** Workgroup grids over 65535 spill into y; shaders linearize back with
 *  `gid.x + gid.y * 4194240u` (65535 * 64, one x-row of threads). */
const MAX_WG_X = 65535;

export class PhaseCancelled extends Error {
  constructor() {
    super('phase cancelled');
  }
}

export interface PhaseRunOptions {
  device: GPUDevice;
  shaders: Record<string, ShaderSpec>;
  phase: PhaseSpec;
  /** Fully evaluated scope: params, derived, and any runtime values. */
  scope: Record<string, number>;
  /** Uploaded into same-named declared buffers before setup. */
  inputs?: Record<string, Float32Array>;
  /** Buffer names to read back when the phase completes. */
  outputs: string[];
  label: string;
  onProgress?: (done: number, total: number) => void;
  cancelled?: () => boolean;
}

const bumpShader = `
@group(0) @binding(0) var<storage, read_write> tick: array<u32>;
@compute @workgroup_size(1)
fn main() { tick[0] = tick[0] + 1u; }
`;

export async function runPhase(opts: PhaseRunOptions): Promise<Record<string, Float32Array>> {
  const { device, shaders, phase, scope, label } = opts;
  const owned: GPUBuffer[] = [];
  const mk = (name: string, bytes: number, extra: GPUBufferUsageFlags = 0): GPUBuffer => {
    const b = device.createBuffer({
      label: `${label}-${name}`,
      size: Math.max(4, bytes),
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST | extra,
    });
    owned.push(b);
    return b;
  };

  try {
    // Buffers. Sizes are expressions; contents start zeroed (WebGPU
    // guarantees zero-initialized buffers), inputs overwrite by name.
    const buffers = new Map<string, GPUBuffer>();
    for (const b of phase.buffers) {
      const count = evalExpr(b.count, scope);
      if (!Number.isInteger(count) || count < 0) {
        throw new Error(`buffer '${b.name}': count '${b.count}' = ${count} is not a whole number`);
      }
      buffers.set(b.name, mk(b.name, 4 * count));
    }
    buffers.set('__tick', mk('__tick', 4));
    for (const [name, data] of Object.entries(opts.inputs ?? {})) {
      const buf = buffers.get(name);
      if (!buf) throw new Error(`input '${name}' has no declared buffer`);
      if (data.byteLength > buf.size) {
        throw new Error(`input '${name}' (${data.byteLength} B) exceeds its buffer (${buf.size} B)`);
      }
      device.queue.writeBuffer(buf, 0, data as Float32Array<ArrayBuffer>);
    }

    // Pipelines: one per distinct shader used, override constants from the
    // scope. An unknown constant name is a real authoring error; surface it.
    const passes: PassSpec[] = [...phase.setup, ...(phase.loop?.body ?? [])];
    const pipelines = new Map<string, GPUComputePipeline>();
    device.pushErrorScope('validation');
    for (const p of passes) {
      if (pipelines.has(p.shader)) continue;
      const s = shaders[p.shader];
      const constants: Record<string, number> = {};
      for (const c of s.constants) {
        if (!(c in scope)) throw new Error(`shader '${p.shader}': constant '${c}' is not in scope`);
        constants[c] = scope[c];
      }
      pipelines.set(
        p.shader,
        await device.createComputePipelineAsync({
          label: `${label}-${p.shader}`,
          layout: 'auto',
          compute: {
            module: device.createShaderModule({ code: s.code, label: `${label}-${p.shader}` }),
            entryPoint: 'main',
            constants,
          },
        }),
      );
    }
    const bumpPipe = await device.createComputePipelineAsync({
      label: `${label}-bump`,
      layout: 'auto',
      compute: { module: device.createShaderModule({ code: bumpShader }), entryPoint: 'main' },
    });
    const pipeErr = await device.popErrorScope();
    if (pipeErr) throw new Error(`${label}: shader compilation failed: ${pipeErr.message}`);

    // Rotation: per phase index, each logical name in a rotate group maps to
    // the group's physical buffer (index - k) mod n. Bind groups are built
    // per pass per rotation phase; passes outside the loop use phase 0.
    const rotate = phase.loop?.rotate ?? [];
    const nPhases = rotate.reduce((a, g) => lcm(a, g.length), 1);
    const physical = (name: string, k: number): GPUBuffer => {
      for (const g of rotate) {
        const l = g.indexOf(name);
        if (l >= 0) {
          const idx = ((l - k) % g.length + g.length) % g.length;
          const buf = buffers.get(g[idx]);
          if (!buf) throw new Error(`rotate group buffer '${g[idx]}' is not declared`);
          return buf;
        }
      }
      const buf = buffers.get(name);
      if (!buf) throw new Error(`pass binds undeclared buffer '${name}'`);
      return buf;
    };
    const bindOf = (p: PassSpec, k: number): GPUBindGroup =>
      device.createBindGroup({
        layout: pipelines.get(p.shader)!.getBindGroupLayout(0),
        entries: p.buffers.map((name, i) => ({
          binding: i,
          resource: { buffer: physical(name, k) },
        })),
      });
    const setupBinds = phase.setup.map((p) => bindOf(p, 0));
    const loopBinds = (phase.loop?.body ?? []).map((p) =>
      Array.from({ length: nPhases }, (_, k) => bindOf(p, k)),
    );
    const bumpBind = device.createBindGroup({
      layout: bumpPipe.getBindGroupLayout(0),
      entries: [{ binding: 0, resource: { buffer: buffers.get('__tick')! } }],
    });

    const dispatchCount = new Map<PassSpec, number>();
    for (const p of passes) dispatchCount.set(p, Math.ceil(evalExpr(p.dispatch, scope) / WORKGROUP_SIZE));

    const encodePass = (
      cp: GPUComputePassEncoder,
      p: PassSpec,
      bind: GPUBindGroup,
    ): void => {
      cp.setPipeline(pipelines.get(p.shader)!);
      cp.setBindGroup(0, bind);
      const wg = dispatchCount.get(p)!;
      cp.dispatchWorkgroups(Math.min(wg, MAX_WG_X), Math.ceil(wg / MAX_WG_X));
    };

    // Setup.
    {
      const enc = device.createCommandEncoder({ label: `${label}-setup` });
      const cp = enc.beginComputePass();
      phase.setup.forEach((p, i) => encodePass(cp, p, setupBinds[i]));
      cp.end();
      device.queue.submit([enc.finish()]);
    }

    // The loop, batched.
    if (phase.loop) {
      const total = evalExpr(phase.loop.count, scope);
      if (!Number.isInteger(total) || total < 0) {
        throw new Error(`loop count '${phase.loop.count}' = ${total} is not a whole number`);
      }
      const batch = phase.loop.batch ?? 256;
      const body = phase.loop.body;
      let done = 0;
      while (done < total) {
        if (opts.cancelled?.()) throw new PhaseCancelled();
        const n = Math.min(batch, total - done);
        const enc = device.createCommandEncoder({ label: `${label}-loop` });
        const cp = enc.beginComputePass();
        for (let j = 0; j < n; j++) {
          const k = (done + j) % nPhases;
          body.forEach((p, i) => encodePass(cp, p, loopBinds[i][k]));
          cp.setPipeline(bumpPipe);
          cp.setBindGroup(0, bumpBind);
          cp.dispatchWorkgroups(1);
        }
        cp.end();
        device.queue.submit([enc.finish()]);
        await device.queue.onSubmittedWorkDone();
        done += n;
        opts.onProgress?.(done, total);
      }
    }

    // Readback.
    const out: Record<string, Float32Array> = {};
    for (const name of opts.outputs) {
      const src = buffers.get(name);
      if (!src) throw new Error(`output '${name}' has no declared buffer`);
      const rb = device.createBuffer({
        label: `${label}-read-${name}`,
        size: src.size,
        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
      });
      owned.push(rb);
      const enc = device.createCommandEncoder();
      enc.copyBufferToBuffer(src, 0, rb, 0, src.size);
      device.queue.submit([enc.finish()]);
      await rb.mapAsync(GPUMapMode.READ);
      out[name] = new Float32Array(rb.getMappedRange().slice(0));
      rb.unmap();
    }
    return out;
  } finally {
    for (const b of owned) b.destroy();
  }
}

const gcd = (a: number, b: number): number => (b === 0 ? a : gcd(b, a % b));
const lcm = (a: number, b: number): number => (a / gcd(a, b)) * b;