// Persistent-kernel backend: ONE dispatch of NWG workgroups runs CHUNK time
// steps, with a global spin barrier on storage atomics between stages in
// place of dispatch boundaries. This removes the per-dispatch launch bubble
// that keeps an idle-clocked GPU at its minimum frequency, and it uses the
// whole GPU, unlike the one-workgroup megakernel.
//
// WebGPU does not guarantee cross-workgroup forward progress, so this mode
// assumes all NWG workgroups are co-resident (13 workgroups of 128 lanes is
// far below any real device's residency) and is validated against the
// multi-dispatch backend, which remains the portable fallback. No thread
// ever returns early: every guard is non-returning so all threads reach
// every barrier.

import { fmt } from './kernel.mjs';

export function buildPersistent(L, opts = {}) {
  const wg = opts.wg ?? 128;
  const chunk = opts.chunk ?? 512;
  const { nrt, nth, N, nf, K, nsteps, off, A, c } = L;
  const { ST_CVEL, ST_D, ST_U, ST_UH } = L.layout;
  const nwg = Math.ceil(N / wg);

  const constLines = Object.entries(c)
    .map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`)
    .join('\n');

  const code = /* wgsl */ `
const NRT = ${nrt}u;
const NTH = ${nth}u;
const N = ${N}u;
const NF = ${nf}u;
const NSTEPS = ${nsteps}u;
const KS = ${K}u;
const NRU = ${L.Nr}u;
const NWG = ${nwg}u;
const CHUNK = ${chunk}u;
const LS0 = ${L.LS0}u; const JS0 = ${L.JS0}u;
const LP0 = ${L.LP0}u; const JP0 = ${L.JP0}u;
const TIPI = ${L.TIPI0}u;
const NLR0 = ${L.NLR0}u;
const HOLDSTEPS = ${L.HOLDSTEPS};
${constLines}
${Object.entries(off).map(([k, v]) => `const OP_${k} = ${v}u;`).join('\n')}
const ST_CVEL = ${ST_CVEL}u;
const ST_D = ${ST_D}u;
const ST_U = ${ST_U}u;
const ST_UH = ${ST_UH}u;
${L.scratchNames.map((n) => `const A_${n} = ${A[n]}u;`).join('\n')}

@group(0) @binding(0) var<storage, read> ops: array<f32>;
@group(0) @binding(1) var<storage, read_write> st: array<f32>;
@group(0) @binding(2) var<storage, read_write> outb: array<f32>;
@group(0) @binding(3) var<storage, read_write> sync: array<atomic<u32>, 2>;

var<workgroup> wgF: f32;
var<workgroup> wgHA: f32;
var<workgroup> wgHB: f32;
var<workgroup> wgHC: f32;
var<workgroup> shD: array<f32, ${nf}>;
var<workgroup> shD1: array<f32, ${nf}>;

fn thm(j: u32) -> u32 { return (j + NTH - 1u) % NTH; }
fn thp(j: u32) -> u32 { return (j + 1u) % NTH; }

var<private> P_lid: u32;

// Classic two-phase global barrier: sync[0] arrive counter, sync[1] epoch.
fn globalBarrier() {
  workgroupBarrier();
  storageBarrier();
  if (P_lid == 0u) {
    let gen = atomicLoad(&sync[1]);
    let n = atomicAdd(&sync[0], 1u);
    if (n == NWG - 1u) {
      atomicStore(&sync[0], 0u);
      atomicAdd(&sync[1], 1u);
    } else {
      loop {
        if (atomicLoad(&sync[1]) != gen) { break; }
      }
    }
  }
  workgroupBarrier();
  storageBarrier();
}

fn u_dr(b: u32, l: u32, j: u32) -> f32 {
  if (l < 1u || l > NRT - 2u) { return 0.0; }
  return (st[b + (l + 1u) + j * NRT] - st[b + (l - 1u) + j * NRT]) * C1R;
}

fn p_dr(l: u32, j: u32) -> f32 {
  if (l < 1u || l > NRT - 2u) { return 0.0; }
  return (st[A_PHI + (l + 1u) + j * NRT] - st[A_PHI + (l - 1u) + j * NRT]) * C1R;
}

@compute @workgroup_size(${wg})
fn main(@builtin(global_invocation_id) gid: vec3<u32>,
        @builtin(local_invocation_index) lid: u32,
        @builtin(workgroup_id) wgid: vec3<u32>) {
  P_lid = lid;
  let i = gid.x;
  let l = i % NRT;
  let j = i / NRT;
  let onGrid = i < N;
  let kbase = u32(st[0u]);

  for (var s: u32 = 0u; s < CHUNK; s = s + 1u) {
    let k = kbase + s;
    let live = k < NSTEPS;
    let uS = ST_U + (k % 3u) * N;

    // ---- AB: hand + contact (wg0 lane0), beam (wg0), stencils (all) ----
    if (wgid.x == 0u) {
      if (lid == 0u && live) {
        var shellA = u32(st[1u]); var inC = u32(st[2u]); var phase = u32(st[3u]);
        var strokeIdx = u32(st[4u]); var holdLeft = i32(st[5u]); var vtar = st[6u];
        var vH = st[7u]; var uHm = st[8u]; var uH0 = st[9u]; var uHp = st[10u];
        var peakXi = st[12u]; var peakDisp = st[13u];
        let dCur = ST_D + ((k + 2u) % 3u) * NF;
        let dOld = ST_D + ((k + 1u) % 3u) * NF;
        if (shellA != 0u && k > 0u && (k % 512u) == 0u) {
          var pd: f32 = 0.0;
          for (var l2: u32 = 0u; l2 < NRT; l2 = l2 + 1u) {
            pd = max(pd, abs(st[uS + l2 + JS0 * NRT]));
          }
          if (pd > peakDisp) { peakDisp = pd; }
        }
        let xi = st[uS + LS0 + JS0 * NRT] - st[dCur + TIPI] - GAP;
        var fM: f32 = 0.0;
        if (xi > 0.0) {
          fM = KM * pow(xi, ALM);
          shellA = 1u;
          if (inC == 0u && strokeIdx < KS && st[ST_CVEL + strokeIdx] == 0.0) {
            st[ST_CVEL + strokeIdx] = (st[dCur + TIPI] - st[dOld + TIPI]) / HT;
          }
          inC = 1u;
          if (xi > peakXi) { peakXi = xi; }
        } else {
          inC = 0u;
        }
        outb[NSTEPS + k] = fM;
        wgF = fM;
        wgHA = -(uHp - 2.0 * uH0 + uHm);
        wgHB = -(HT * 0.5) * (uHp - uHm);
        wgHC = -HT * HT * (PS * uH0 + ((1.0 - PS) * 0.5) * (uHp + uHm));
        if (phase == 0u) {
          if (strokeIdx < KS && f32(k) >= ops[OP_KGO + strokeIdx]) {
            phase = 1u; vtar = ops[OP_VSTR + strokeIdx];
          }
        } else if (phase == 1u) {
          vH = max(vH - AH * HT, -vtar);
          if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; }
          else if (vH <= -vtar) { phase = 2u; }
        } else if (phase == 2u) {
          if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; }
          else if (uHp <= -MISSD) { phase = 4u; }
        } else if (phase == 3u) {
          holdLeft = holdLeft - 1;
          if (holdLeft <= 0) { phase = 4u; }
        } else if (phase == 4u) {
          vH = min(vH + AHR * HT, 0.5 * vtar);
          if (vH > 0.0 && uHp >= -vH * vH / (2.0 * AHR)) { phase = 5u; }
        } else {
          vH = vH - AHR * HT;
          if (vH <= 0.0 || uHp >= 0.0) { vH = 0.0; phase = 0u; strokeIdx = strokeIdx + 1u; }
        }
        uHm = uH0; uH0 = uHp; uHp = uHp + vH * HT;
        st[1u] = f32(shellA); st[2u] = f32(inC); st[3u] = f32(phase);
        st[4u] = f32(strokeIdx); st[5u] = f32(holdLeft); st[6u] = vtar;
        st[7u] = vH; st[8u] = uHm; st[9u] = uH0; st[10u] = uHp;
        st[12u] = peakXi; st[13u] = peakDisp; st[14u] = fM;
      }
      if (lid < NF && live) {
        shD[lid] = st[ST_D + ((k + 2u) % 3u) * NF + lid];
        shD1[lid] = st[ST_D + ((k + 1u) % 3u) * NF + lid];
      }
      workgroupBarrier();
      if (lid < NF && live) {
        let r = lid;
        var sum: f32 = ops[OP_VM + r] * wgHA + ops[OP_VR + r] * wgHB + ops[OP_VK + r] * wgHC
          + (HT * HT * wgF) * ops[OP_TIPCOL + r];
        for (var p: u32 = 0u; p < NF; p = p + 1u) {
          sum = sum + ops[OP_ASIBT + p + r * NF] * shD[p]
                    + ops[OP_ASICT + p + r * NF] * shD1[p];
        }
        st[ST_D + (k % 3u) * NF + r] = sum;
      }
    } else {
      workgroupBarrier();   // match wg0's barrier count
    }
    // The shell-active latch for THIS step is decided by wg0 above; other
    // workgroups may see the pre-latch zero here, which only skips work on a
    // still-zero field, exactly as in the multi-dispatch backend.
    var act = onGrid && live && st[1u] != 0.0;
    if (act) {
      var ur: f32 = 0.0; var urr: f32 = 0.0;
      if (l >= 1u && l <= NRT - 2u) {
        let um = st[uS + (l - 1u) + j * NRT];
        let u0 = st[uS + l + j * NRT];
        let up = st[uS + (l + 1u) + j * NRT];
        ur = (up - um) * C1R;
        urr = (um - 2.0 * u0 + up) * C2R;
      }
      let ujm = st[uS + l + thm(j) * NRT];
      let uj0 = st[uS + l + j * NRT];
      let ujp = st[uS + l + thp(j) * NRT];
      let v1 = ops[OP_RD + l] * ur + ops[OP_R2D + l] * ((ujm - 2.0 * uj0 + ujp) * C2T);
      let urt = (u_dr(uS, l, thp(j)) - u_dr(uS, l, thm(j))) * C1T;
      let v2 = ops[OP_RD + l] * urt - ops[OP_R2D + l] * ((ujp - ujm) * C1T);
      st[A_URR + i] = urr;
      st[A_V1 + i] = v1;
      st[A_V2 + i] = v2;
      var rv: f32 = 0.0;
      if (l >= 3u && l <= NRU - 2u) {
        rv = EHRS * (urr + v1) - EH2 * 2.0 * (urr * v1 - v2 * v2);
      }
      st[A_RHSA + i] = rv;
    }
    globalBarrier();
    // After the barrier every workgroup sees the latch for this step.
    act = onGrid && live && st[1u] != 0.0;

    // ---- C: RHSAF = RHSA * F ----
    if (act) {
      var sum: f32 = 0.0;
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        sum = sum + st[A_RHSA + l + q * NRT] * ops[OP_F + q + j * NTH];
      }
      st[A_RHSAF + i] = sum;
    }
    globalBarrier();

    // ---- D: Airy solve per wavenumber ----
    if (act) {
      let base = OP_FI + j * NRT * NRT;
      var sum: f32 = 0.0;
      for (var p: u32 = 0u; p < NRT; p = p + 1u) {
        sum = sum + ops[base + l + p * NRT] * st[A_RHSAF + p + j * NRT];
      }
      st[A_PH3 + i] = sum;
    }
    globalBarrier();

    // ---- E: PHI = PH3 * Gi ----
    if (act) {
      var sum: f32 = 0.0;
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        sum = sum + st[A_PH3 + l + q * NRT] * ops[OP_GI + q + j * NTH];
      }
      st[A_PHI + i] = sum;
    }
    globalBarrier();

    // ---- FG: nonlinear load ----
    if (act) {
      var nl: f32 = 0.0;
      if (l >= NLR0 && l < NRU) {
        let pm = st[A_PHI + (l - 1u) + j * NRT];
        let p0 = st[A_PHI + l + j * NRT];
        let pp = st[A_PHI + (l + 1u) + j * NRT];
        let pr = (pp - pm) * C1R;
        let prr = (pm - 2.0 * p0 + pp) * C2R;
        let pjm = st[A_PHI + l + thm(j) * NRT];
        let pjp = st[A_PHI + l + thp(j) * NRT];
        let w1 = ops[OP_RD + l] * pr + ops[OP_R2D + l] * ((pjm - 2.0 * p0 + pjp) * C2T);
        let prt = (p_dr(l, thp(j)) - p_dr(l, thm(j))) * C1T;
        let w2 = ops[OP_RD + l] * prt - ops[OP_R2D + l] * ((pjp - pjm) * C1T);
        nl = prr * st[A_V1 + i] + st[A_URR + i] * w1 - 2.0 * w2 * st[A_V2 + i]
           - (prr + w1) * RSI;
        if (l == LS0 && j == JS0) { nl = nl - st[14u] * FSC; }
        nl = nl * SNL;
      }
      st[A_NL + i] = nl;
    }
    globalBarrier();

    // ---- H: NLF = NL * F ----
    if (act) {
      var sum: f32 = 0.0;
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        sum = sum + st[A_NL + l + q * NRT] * ops[OP_F + q + j * NTH];
      }
      st[A_NLF + i] = sum;
    }
    globalBarrier();

    // ---- I: implicit update per wavenumber ----
    if (act) {
      let h3S = ST_UH + (k % 3u) * N;
      let h13S = ST_UH + ((k + 2u) % 3u) * N;
      let h3nS = ST_UH + ((k + 1u) % 3u) * N;
      let base = j * NRT * NRT;
      var sum: f32 = 0.0;
      for (var p: u32 = 0u; p < NRT; p = p + 1u) {
        let pj = p + j * NRT;
        sum = sum + ops[OP_AIB + base + l + p * NRT] * st[h3S + pj]
                  + ops[OP_AIC + base + l + p * NRT] * st[h13S + pj]
                  + ops[OP_AI + base + l + p * NRT] * st[A_NLF + pj];
      }
      st[h3nS + i] = sum;
    }
    globalBarrier();

    // ---- J: synthesis, pickup acceleration ----
    if (act) {
      let h3nS = ST_UH + ((k + 1u) % 3u) * N;
      let u1S = ST_U + ((k + 2u) % 3u) * N;
      let unS = ST_U + ((k + 1u) % 3u) * N;
      var sum: f32 = 0.0;
      for (var q: u32 = 0u; q < NTH; q = q + 1u) {
        sum = sum + st[h3nS + l + q * NRT] * ops[OP_GI + q + j * NTH];
      }
      st[unS + i] = sum;
      if (l == LP0 && j == JP0) {
        outb[k] = (sum - 2.0 * st[uS + i] + st[u1S + i]) * HT2I;
      }
    }
    globalBarrier();
  }

  if (wgid.x == 0u && lid == 0u) {
    st[0u] = f32(min(kbase + CHUNK, NSTEPS));
  }
}
`;

  return { code, chunk, wg, nwg };
}