import { MiniCPMRuntime, layoutTypes, cacheBytes } from "./minicpm-runtime.js";
import { buildDSparkKernels } from "./dspark-kernels.js";
import {
  H,
  F,
  KV,
  VOCAB,
  MAX_CHUNKS,
  PREFILL_BATCH,
} from "./minicpm-kernels.js";
export const DSPARK_LAYERS = [1, 10, 20, 30, 39];
const extraTypes = {
  projectBatch: ["r", "r", "w"],
  projectBlock: ["r", "r", "w"],
  contextBatch: ["r", "r", "w"],
  contextBlock: ["r", "r", "w"],
  capture: ["r", "r", "w", "u"],
  noiseNormRope: ["r", "r", "r", "w", "w", "w"],
  contextNormRope: ["r", "r", "r", "w", "w", "w"],
  noise: ["r", "w"],
  markovEmbed: ["r", "r", "w"],
  markov: ["r", "r", "w"],
  addBias: ["r", "r", "w"],
  argmax: ["r", "w"],
  argmaxFinal: ["r", "w"],
};
export class DSparkRuntime extends MiniCPMRuntime {
  async initialize(parent, blob, signal) {
    this.parent = parent;
    this.device = parent.device;
    this.adapter = parent.adapter;
    this.subgroups = parent.subgroups;
    this.maxContext = parent.maxContext;
    this.stepLayout = parent.stepLayout;
    this.settings = parent.settings;
    const all = buildDSparkKernels({
      subgroups: this.subgroups,
      kv: this.settings.cacheK,
    });
    const needed = [
      "norm",
      "qkvBlock",
      "outputBlock",
      "upGateBlock",
      "downBlock",
      "silu",
      "projectBatch",
      "projectBlock",
      "contextBatch",
      "contextBlock",
      "markov",
      "store",
      "attention",
      "combine",
      ...Object.keys(extraTypes),
    ];
    await Promise.all(
      [...new Set(needed)].map(async (name) => {
        const module = this.device.createShaderModule({
          code: all[name],
          label: "dspark " + name,
        });
        const info = await module.getCompilationInfo();
        const errors = info.messages.filter((m) => m.type === "error");
        if (errors.length)
          throw new Error(
            "DSpark " +
              name +
              ": " +
              errors.map((e) => e.message + " at " + e.lineNum).join("; "),
          );
        const types =
          extraTypes[name] ||
          layoutTypes[name] ||
          layoutTypes[name.replace(/Block$/, "")];
        const layout = this.device.createBindGroupLayout({
          entries: types.map((type, binding) => ({
            binding,
            visibility: GPUShaderStage.COMPUTE,
            buffer: {
              type:
                type === "u"
                  ? "uniform"
                  : type === "w"
                    ? "storage"
                    : "read-only-storage",
            },
          })),
        });
        this.layouts[name] = layout;
        this.pipelines[name] = await this.device.createComputePipelineAsync({
          layout: this.device.createPipelineLayout({
            bindGroupLayouts: [this.stepLayout, layout],
          }),
          compute: { module, entryPoint: "main" },
          label: "DSpark " + name,
        });
      }),
    );
    const headSize = Number(
      new DataView(await blob.slice(0, 8).arrayBuffer()).getBigUint64(0, true),
    );
    if (headSize > 1024 * 1024)
      throw new Error("Invalid DSpark weights header.");
    const header = JSON.parse(await blob.slice(8, 8 + headSize).text());
    const offset = 8 + headSize;
    const upload = async (names, label, fp32 = false) => {
      const list = [];
      for (const name of names) {
        signal?.throwIfAborted();
        const item = header[name];
        if (!item || item.dtype !== "BF16")
          throw new Error("Missing BF16 DSpark tensor: " + name);
        const bytes = await blob
          .slice(offset + item.data_offsets[0], offset + item.data_offsets[1])
          .arrayBuffer();
        if (fp32) {
          const source = new Uint16Array(bytes),
            out = new Uint32Array(source.length);
          for (let i = 0; i < source.length; i++) out[i] = source[i] << 16;
          list.push(out.buffer);
        } else list.push(bytes);
      }
      const b = this.buffer(
        list.reduce((n, a) => n + Math.ceil(a.byteLength / 4) * 4, 0),
        label,
      );
      let at = 0;
      for (const bytes of list) {
        let data = bytes;
        if (data.byteLength % 4) {
          const padded = new Uint8Array(Math.ceil(data.byteLength / 4) * 4);
          padded.set(new Uint8Array(data));
          data = padded.buffer;
        }
        this.device.queue.writeBuffer(b, at, data);
        at += data.byteLength;
      }
      return b;
    };
    const w = (this.work = {});
    for (const [name, count] of Object.entries({
      features: PREFILL_BATCH * H * 5,
      projected: PREFILL_BATCH * H,
      context: PREFILL_BATCH * H,
      contextKV: PREFILL_BATCH * KV * 2,
      x: PREFILL_BATCH * H,
      norm: PREFILL_BATCH * H,
      temp: PREFILL_BATCH * H,
      zeros: PREFILL_BATCH * H,
      qkv: PREFILL_BATCH * (H + 2 * KV),
      q: PREFILL_BATCH * H,
      k: PREFILL_BATCH * KV,
      v: PREFILL_BATCH * KV,
      attention: PREFILL_BATCH * H,
      upGate: PREFILL_BATCH * F * 2,
      hidden: PREFILL_BATCH * F,
      partial: PREFILL_BATCH * 16 * MAX_CHUNKS * 128,
      md: PREFILL_BATCH * 16 * MAX_CHUNKS * 2,
      tokens: this.maxContext + 9,
      baseLogits: 8 * VOCAB,
      bias: VOCAB,
      logits: VOCAB,
      markov: 256,
      argmax: 256,
    }))
      w[name] = this.buffer(count * 4, "DSpark " + name);
    w.sampler = this.uniform(new Float32Array(4), "draft sampler");
    w.steps = this.buffer(
      (this.maxContext + 9) * 256,
      "DSpark steps",
      GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    );
    w.sampleSteps = this.buffer(
      (this.maxContext + 9) * 256,
      "DSpark sample steps",
      GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    );
    this.stepGroup = this.device.createBindGroup({
      layout: this.stepLayout,
      entries: [{ binding: 0, resource: { buffer: w.steps, size: 16 } }],
    });
    this.sampleStepGroup = this.device.createBindGroup({
      layout: this.stepLayout,
      entries: [{ binding: 0, resource: { buffer: w.sampleSteps, size: 16 } }],
    });
    const [fc, hiddenNorm, finalNorm, w1, w2, confidence, confidenceBias] =
      await Promise.all([
        upload(["fc.weight"], "DSpark projection"),
        upload(["hidden_norm.weight"], "DSpark hidden norm", true),
        upload(["norm.weight"], "DSpark output norm", true),
        upload(["markov_head.markov_w1.weight"], "DSpark Markov embeddings"),
        upload(["markov_head.markov_w2.weight"], "DSpark Markov projection"),
        upload(["confidence_head.proj.weight"], "DSpark confidence", true),
        upload(["confidence_head.proj.bias"], "DSpark confidence bias", true),
      ]);
    this.confidence = { weight: confidence, bias: confidenceBias };
    this.groups = {
      projectBatch: this.group("projectBatch", [w.features, fc, w.projected]),
      projectBlock: this.group("projectBlock", [w.features, fc, w.projected]),
      hiddenNorm: this.group("norm", [
        w.projected,
        w.zeros,
        hiddenNorm,
        w.context,
      ]),
      noise: this.group("noise", [parent.work.tokens, w.tokens]),
      embed: parent.group("embed", [...parent.embedding, w.tokens, w.x]),
      finalNorm: this.group("norm", [w.x, w.temp, finalNorm, w.norm]),
      head: parent.group("headBlock", [w.norm, ...parent.head, w.baseLogits]),
      markovEmbed: this.group("markovEmbed", [w1, w.tokens, w.markov]),
      markov: this.group("markov", [w.markov, w2, w.bias]),
      argmax: this.group("argmax", [w.logits, w.argmax]),
      argmaxFinal: this.group("argmaxFinal", [w.argmax, w.tokens]),
    };
    this.biasGroups = Array.from({ length: 7 }, (_, i) =>
      this.group("addBias", [
        { buffer: w.baseLogits, offset: i * VOCAB * 4, size: VOCAB * 4 },
        w.bias,
        w.logits,
      ]),
    );
    this.captureGroups = DSPARK_LAYERS.map((_, i) =>
      this.group("capture", [
        parent.work.x,
        parent.work.temp,
        w.features,
        this.uniform(new Uint32Array([i, 0, 0, 0]), "capture layer"),
      ]),
    );
    for (let i = 0; i < 5; i++) {
      const prefix = "layers." + i + ".";
      const get = (names, label, float = false) =>
        upload(
          names.map((n) => prefix + n + ".weight"),
          label + " " + i,
          float,
        );
      const [norm, post, qkv, contextKV, o, upGate, down, qNorm, kNorm] =
        await Promise.all([
          get(["input_layernorm"], "DSpark norm", true),
          get(["post_attention_layernorm"], "DSpark post norm", true),
          get(
            ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"],
            "DSpark qkv",
          ),
          get(["self_attn.k_proj", "self_attn.v_proj"], "DSpark context KV"),
          get(["self_attn.o_proj"], "DSpark attention out"),
          get(["mlp.gate_proj", "mlp.up_proj"], "DSpark gate/up"),
          get(["mlp.down_proj"], "DSpark down"),
          get(["self_attn.q_norm"], "DSpark Q norm", true),
          get(["self_attn.k_norm"], "DSpark K norm", true),
        ]);
      const k = this.buffer(
          cacheBytes(this.maxContext, this.settings.cacheK),
          "DSpark key cache",
        ),
        v = this.buffer(
          cacheBytes(this.maxContext, this.settings.cacheK),
          "DSpark value cache",
        );
      const groups = {
        input: this.group("norm", [w.x, i ? w.temp : w.zeros, norm, w.norm]),
        post: this.group("norm", [w.x, w.temp, post, w.norm]),
        qkv: this.group("qkvBlock", [w.norm, qkv, w.qkv]),
        contextBatch: this.group("contextBatch", [
          w.context,
          contextKV,
          w.contextKV,
        ]),
        contextBlock: this.group("contextBlock", [
          w.context,
          contextKV,
          w.contextKV,
        ]),
        noiseRope: this.group("noiseNormRope", [
          w.qkv,
          qNorm,
          kNorm,
          w.q,
          w.k,
          w.v,
        ]),
        contextRope: this.group("contextNormRope", [
          w.contextKV,
          qNorm,
          kNorm,
          w.q,
          w.k,
          w.v,
        ]),
        store: this.group("store", [w.k, w.v, k, v]),
        attention: this.group("attention", [w.q, k, v, w.partial, w.md]),
        combine: this.group("combine", [w.partial, w.md, w.attention]),
        output: this.group("outputBlock", [w.attention, o, w.temp]),
        upGate: this.group("upGateBlock", [w.norm, upGate, w.upGate]),
        silu: this.group("silu", [w.upGate, w.hidden]),
        down: this.group("downBlock", [w.hidden, down, w.temp]),
      };
      this.layers.push({ groups, k, v });
      await this.device.queue.onSubmittedWorkDone();
      parent.check();
    }
    this.sortGroups = {};
    for (const keep of [64, 256]) {
      const pg = parent.sortGroups[keep];
      this.sortGroups[keep] = {
        ...pg,
        first: parent.group("sortFirst" + keep, [
          w.logits,
          parent.work.sortA,
          w.tokens,
          parent.work.penalty,
        ]),
        sample: parent.group("sample", [pg.source, w.tokens, w.sampler]),
      };
    }
    this.ready = true;
  }
  capture(pass, layer, count) {
    const index = DSPARK_LAYERS.indexOf(layer);
    if (index >= 0)
      this.dispatch(pass, "capture", this.captureGroups[index], H / 256, count);
  }
  updateContext(encoder, position, count) {
    if (!count) return;
    this.contextRun = position + count;
    this.prepareStep(position, count);
    const pass = encoder.beginComputePass();
    pass.setBindGroup(0, this.stepGroup, [position * 256]);
    const run = (name, g, x = 1, y = 1, z = 1) =>
        this.dispatch(pass, name, g, x, y, z),
      block = count <= 8;
    run(
      block ? "projectBlock" : "projectBatch",
      block ? this.groups.projectBlock : this.groups.projectBatch,
      H / (block ? 8 : 32),
    );
    run("norm", this.groups.hiddenNorm, count);
    for (const { groups: g } of this.layers) {
      run(
        block ? "contextBlock" : "contextBatch",
        block ? g.contextBlock : g.contextBatch,
        (KV * 2) / (block ? 8 : 32),
      );
      run("contextNormRope", g.contextRope, 2, count);
      run("store", g.store, KV / 32, count, 2);
    }
    pass.end();
  }
  propose(encoder, position) {
    const from = Math.max(0, position - 64);
    encoder.copyBufferToBuffer(
      this.parent.work.tokens,
      from * 4,
      this.work.tokens,
      from * 4,
      (position - from + 1) * 4,
    );
    const count = 7,
      plan = this.prepareStep(position, count),
      pass = encoder.beginComputePass();
    pass.setBindGroup(0, this.stepGroup, [position * 256]);
    const run = (name, g, x = 1, y = 1, z = 1) =>
      this.dispatch(pass, name, g, x, y, z);
    run("noise", this.groups.noise);
    this.parent.dispatch(pass, "embed", this.groups.embed, H / 256, count);
    for (const { groups: g } of this.layers) {
      run("norm", g.input, count);
      run("qkvBlock", g.qkv, (H + KV * 2) / 8);
      run("noiseNormRope", g.noiseRope, 18, count);
      run("store", g.store, KV / 32, count, 2);
      run("attention", g.attention, plan.chunks, 16, count);
      run("combine", g.combine, 16, count);
      run("outputBlock", g.output, H / 8);
      run("norm", g.post, count);
      run("upGateBlock", g.upGate, (F * 2) / 8);
      run("silu", g.silu, F / 256, count);
      run("downBlock", g.down, H / 8);
    }
    run("norm", this.groups.finalNorm, count);
    this.parent.dispatch(pass, "headBlock", this.groups.head, VOCAB / 8);
    for (let i = 0; i < 7; i++) {
      this.device.queue.writeBuffer(
        this.work.sampleSteps,
        (position + i) * 256,
        new Uint32Array([position + i, 1, 0, 1]),
      );
      pass.setBindGroup(0, this.sampleStepGroup, [(position + i) * 256]);
      run("markovEmbed", this.groups.markovEmbed);
      run("markov", this.groups.markov, VOCAB / 8);
      run("addBias", this.biasGroups[i], VOCAB / 256);
      if (this.parent.samplingSettings.temperature > 0) {
        const keep = this.parent.sampleK <= 64 ? 64 : 256,
          g = this.sortGroups[keep];
        this.parent.encodeSample(pass, g.first, g);
        this.parent.dispatch(pass, "saveCandidates", g.save[8 + i], keep / 64);
      } else {
        run("argmax", this.groups.argmax, 128);
        run("argmaxFinal", this.groups.argmaxFinal);
      }
    }
    pass.end();
    encoder.copyBufferToBuffer(
      this.work.tokens,
      (position + 1) * 4,
      this.parent.work.tokens,
      (position + 1) * 4,
      7 * 4,
    );
  }
}