import { buildMiniCPMKernels, H, KV, VOCAB } from "./minicpm-kernels.js";
export function buildDSparkKernels(options) {
  const kernels = buildMiniCPMKernels({
    ...options,
    weightFormat: "bf16",
    draft: true,
  });
  kernels.attention = kernels.attention.replace(
    "let pos=step.position+token;",
    "let pos=step.position+step.count-1u;",
  );
  const pre = `${options.subgroups ? "enable subgroups;" : ""}\nstruct Step{position:u32,count:u32,chunk:u32,chunks:u32};@group(0)@binding(0)var<uniform>step:Step;\n`;
  const bind = (i, name, type, mode = "read") =>
    `@group(1)@binding(${i})var<storage,${mode}>${name}:array<${type}>;\n`;
  kernels.capture =
    pre +
    bind(0, "x", "f32") +
    bind(1, "residual", "f32") +
    bind(2, "features", "f32", "read_write") +
    "@group(1)@binding(3)var<uniform>slot:vec4<u32>;" +
    `
 @compute @workgroup_size(256)fn main(@builtin(global_invocation_id)g:vec3<u32>){if(g.x<2048u&&g.y<step.count){features[g.y*10240u+slot.x*2048u+g.x]=x[g.y*2048u+g.x]+residual[g.y*2048u+g.x];}}`;
  const normRope = (context) =>
    pre +
    bind(0, "input", "f32") +
    bind(1, "qNorm", "f32") +
    bind(2, "kNorm", "f32") +
    bind(3, "queries", "f32", "read_write") +
    bind(4, "keys", "f32", "read_write") +
    bind(5, "values", "f32", "read_write") +
    `
 var<workgroup> red:array<f32,128>;var<workgroup> normalized:array<f32,128>;
 @compute @workgroup_size(128)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){let d=l.x;let head=g.x;let t=g.y;let isQ=${context ? "false" : "head<16u"};let i=t*${context ? 512 : 2560}u+head*128u+d;let raw=input[i];red[d]=raw*raw;workgroupBarrier();for(var s=64u;s>0u;s>>=1u){if(d<s){red[d]+=red[d+s];}workgroupBarrier();}var weight=kNorm[d];if(isQ){weight=qNorm[d];}normalized[d]=raw*inverseSqrt(red[0]/128.0+0.000001)*weight;workgroupBarrier();
 let theta=f32(step.position+t)*pow(5000000.0,-f32(d%64u)/64.0);let pair=select(d-64u,d+64u,d<64u);let rotated=normalized[d]*cos(theta)+select(normalized[pair],-normalized[pair],d<64u)*sin(theta);
 if(isQ){queries[t*2048u+head*128u+d]=rotated;}else{let kh=${context ? "head" : "head-16u"};let out=t*256u+kh*128u+d;keys[out]=rotated;values[out]=input[t*${context ? 512 : 2560}u+${context ? 256 : 2304}u+kh*128u+d];}}
 `;
  kernels.noiseNormRope = normRope(false);
  kernels.contextNormRope = normRope(true);
  kernels.noise =
    pre +
    bind(0, "source", "u32") +
    bind(1, "tokens", "u32", "read_write") +
    `@compute @workgroup_size(8)fn main(@builtin(local_invocation_id)l:vec3<u32>){if(l.x<step.count){tokens[step.position+l.x]=select(75982u,source[step.position],l.x==0u);}}`;
  kernels.markovEmbed =
    pre +
    bind(0, "weights", "u32") +
    bind(1, "tokens", "u32") +
    bind(2, "out", "f32", "read_write") +
    `@compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>){let i=tokens[step.position]*256u+l.x;let w=weights[i>>1u];out[l.x]=bitcast<f32>(select(w<<16u,w&0xffff0000u,(i&1u)==1u));}`;
  kernels.addBias =
    pre +
    bind(0, "base", "f32") +
    bind(1, "bias", "f32") +
    bind(2, "out", "f32", "read_write") +
    `@compute @workgroup_size(256)fn main(@builtin(global_invocation_id)g:vec3<u32>){if(g.x<${VOCAB}u){out[g.x]=base[g.x]+bias[g.x];}}`;
  kernels.argmax =
    pre +
    bind(0, "logits", "f32") +
    bind(1, "out", "vec2<u32>", "read_write") +
    `
 var<workgroup> val:array<f32,256>;var<workgroup> ids:array<u32,256>;
 @compute @workgroup_size(256)fn main(@builtin(local_invocation_id)l:vec3<u32>,@builtin(workgroup_id)g:vec3<u32>){var v=-3.402823e38;var id=0xffffffffu;for(var i=g.x*256u+l.x;i<${VOCAB}u;i+=32768u){let x=logits[i];if(x>v||(x==v&&i<id)){v=x;id=i;}}val[l.x]=v;ids[l.x]=id;workgroupBarrier();for(var s=128u;s>0u;s>>=1u){if(l.x<s){let o=l.x+s;if(val[o]>val[l.x]||(val[o]==val[l.x]&&ids[o]<ids[l.x])){val[l.x]=val[o];ids[l.x]=ids[o];}}workgroupBarrier();}if(l.x==0u){out[g.x]=vec2<u32>(bitcast<u32>(val[0]),ids[0]);}}
 `;
  kernels.argmaxFinal =
    pre +
    bind(0, "input", "vec2<u32>") +
    bind(1, "tokens", "u32", "read_write") +
    `
 var<workgroup> val:array<f32,128>;var<workgroup> ids:array<u32,128>;
 @compute @workgroup_size(128)fn main(@builtin(local_invocation_id)l:vec3<u32>){val[l.x]=bitcast<f32>(input[l.x].x);ids[l.x]=input[l.x].y;workgroupBarrier();for(var s=64u;s>0u;s>>=1u){if(l.x<s){let o=l.x+s;if(val[o]>val[l.x]||(val[o]==val[l.x]&&ids[o]<ids[l.x])){val[l.x]=val[o];ids[l.x]=ids[o];}}workgroupBarrier();}if(l.x==0u){tokens[step.position+1u]=ids[0];}}
 `;
  return kernels;
}