worker.js

// NanoMaestro Web Worker for isolated model inference and parsing
try {
  importScripts("https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/ort.min.js");
} catch (err) {
  console.error("Worker failed to import ONNX Runtime script:", err);
}

// ONNX runtime configuration
if (typeof ort !== "undefined") {
  ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/";
  ort.env.wasm.numThreads = Math.min(4, navigator.hardwareConcurrency || 1);
}

// State variables
let session = null;
let stoi = null;
let itos = null;
let vocabSize = 0;

let h = null;
let c = null;
let currentId = 0;
let activeModelName = "";

// Model parameters
let hiddenSize = 2048;
let numLayers = 2;

// Playback settings
let temperature = 0.85;
let topK = 40;
let currentBpm = 120;
let midiBpmMode = "lock";

// Seeding settings
let playingStyle = "default";
let midiTokens = [];
let midiStartBar = 0;
let isWarmingUp = false;

// Temporal queues for gathering parsed events
let tempQueue = [];
let parser = null;



function tokenId(token) {
  return stoi?.[token] ?? null;
}

function makeTensorId(id) {
  return new ort.Tensor("int64", BigInt64Array.from([BigInt(id)]), [1, 1]);
}

function zeroState() {
  return new ort.Tensor("float32", new Float32Array(numLayers * 1 * hiddenSize), [numLayers, 1, hiddenSize]);
}

function updateModelDimensions(session) {
  try {
    const inputNames = session.inputNames;
    if (inputNames.includes("h") && session.handler && session.handler._model) {
      const inputs = session.handler._model.graph.inputs;
      const hInput = inputs.find(i => i.name === "h");
      if (hInput && hInput.type && hInput.type.tensorType && hInput.type.tensorType.shape) {
        const shape = hInput.type.tensorType.shape.dim;
        const layers = Number(shape[0].dimValue);
        const hidden = Number(shape[2].dimValue);
        if (layers > 0 && hidden > 0) {
          numLayers = layers;
          hiddenSize = hidden;
          console.log(`Worker set model dimensions: layers=${numLayers}, hiddenSize=${hiddenSize}`);
        }
      }
    }
  } catch (e) {
    console.error("Worker failed to detect model dimensions:", e);
  }
}

// Model Step Execution with Memory Management
async function stepModel(record = true) {
  if (!session) return "?";
  
  const inputTensor = makeTensorId(currentId);
  const output = await session.run({ input: inputTensor, h, c });
  
  // Dispose input tensor to prevent memory leak
  inputTensor.dispose();
  
  // Dispose previous recurrent states to prevent memory leaks during continuous generation
  if (h) {
    try { h.dispose(); } catch (e) {}
  }
  if (c) {
    try { c.dispose(); } catch (e) {}
  }
  
  h = output.h_out;
  c = output.c_out;
  
  const logitsData = output.logits.data;
  currentId = sampleFromLogits(logitsData);
  output.logits.dispose();
  
  const token = itos[currentId] ?? "?";
  
  // Reset state on EOS
  if (token === "EOS") {
    if (h) { try { h.dispose(); } catch (e) {} }
    if (c) { try { c.dispose(); } catch (e) {} }
    h = zeroState();
    c = zeroState();
  }
  
  if (record && parser) {
    parser.feed(token);
  }
  return token;
}

function sampleFromLogits(logits) {
  const temp = Math.max(0.05, temperature);
  const k = Math.max(1, topK);
  const scored = [];
  for (let i = 0; i < logits.length; i += 1) {
    const token = itos[i];
    if (token === undefined) continue;
    let score = logits[i] / temp;
    if (token === "<EOP>" || token === "EOS") score -= 1.0;
    scored.push([i, score]);
  }
  scored.sort((a, b) => b[1] - a[1]);
  const picked = scored.slice(0, Math.min(k, scored.length));
  const maxScore = picked[0]?.[1] ?? 0;
  let sum = 0;
  for (const item of picked) {
    item[2] = Math.exp(item[1] - maxScore);
    sum += item[2];
  }
  let r = Math.random() * sum;
  for (const item of picked) {
    r -= item[2];
    if (r <= 0) return item[0];
  }
  return picked[picked.length - 1][0];
}

// Duration converters


function durationToSecondsFromEventSteps(steps, grid) {
  const quarterSeconds = 60 / currentBpm;
  return Math.max(0.035, (Math.max(1, steps) * 4 * quarterSeconds) / Math.max(1, grid));
}

function midiToTonePitch(midi) {
  const names = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"];
  return `${names[((midi % 12) + 12) % 12]}${Math.floor(midi / 12) - 1}`;
}



function pushEvent(event) {
  if (isWarmingUp) return; // Discard prompt history events
  tempQueue.push(event);
}

// Token Slicers
function getTokensUpToBar(tokens, targetBar) {
  let barCount = 0;
  const sliced = [];
  for (const t of tokens) {
    if (t === "BAR") {
      barCount += 1;
      if (barCount > targetBar) {
        break;
      }
    }
    sliced.push(t);
  }
  return sliced;
}



// Event Parser
function makeEventParser() {
  return {
    grid: 64,
    bar: -1,
    pos: 0,
    pendingPosition: null,
    pendingNotes: [],
    pendingNote: null,
    lastQ: 0,
    reset() {
      this.grid = 64;
      this.bar = -1;
      this.pos = 0;
      this.pendingPosition = null;
      this.pendingNotes = [];
      this.pendingNote = null;
      this.lastQ = 0;
    },
    feed(token) {
      if (token.startsWith("BPM_")) {
        const bpm = Number(token.slice(4));
        if (Number.isFinite(bpm) && bpm >= 40 && bpm <= 220) {
          if (playingStyle !== "midi" || isWarmingUp || midiBpmMode === "model") {
            currentBpm = bpm;
            self.postMessage({ action: "tempo", bpm: bpm });
          }
        }
        return;
      }
      if (token.startsWith("GRID_")) {
        const grid = Number(token.slice(5));
        if (Number.isFinite(grid) && grid > 0) this.grid = grid;
        return;
      }
      if (token === "BAR") {
        this.flushTo(this.absoluteQFor(this.bar + 1, 0));
        this.bar += 1;
        this.pos = 0;
        return;
      }
      if (token.startsWith("POS_")) {
        const pos = Number(token.slice(4));
        if (!Number.isFinite(pos)) return;
        this.flushTo(this.absoluteQFor(this.bar, pos));
        this.pos = pos;
        return;
      }
      if (token.startsWith("NOTE_")) {
        const midi = Number(token.slice(5));
        if (Number.isFinite(midi) && midi >= 0 && midi <= 127) {
          this.pendingNote = { midi, durationSteps: 1, velocity: 0.72 };
        }
        return;
      }
      if (token.startsWith("DUR_") && this.pendingNote) {
        const steps = Number(token.slice(4));
        if (Number.isFinite(steps)) this.pendingNote.durationSteps = Math.max(1, steps);
        return;
      }
      if (token.startsWith("VEL_") && this.pendingNote) {
        const bucket = Number(token.slice(4));
        if (Number.isFinite(bucket)) this.pendingNote.velocity = Math.max(0.2, Math.min(0.95, bucket / 8));
        const q = this.absoluteQFor(this.bar, this.pos);
        this.pendingPosition ??= q;
        this.pendingNotes.push(this.pendingNote);
        this.pendingNote = null;
      }
    },
    absoluteQFor(bar, pos) {
      return Math.max(0, bar) * 4 + (Math.max(0, pos) * 4) / Math.max(1, this.grid);
    },
    flushTo(nextQ) {
      if (this.pendingNote) {
        const q = this.absoluteQFor(this.bar, this.pos);
        this.pendingPosition ??= q;
        this.pendingNotes.push(this.pendingNote);
        this.pendingNote = null;
      }
      if (this.pendingNotes.length && this.pendingPosition !== null) {
        const gap = Math.max(0, this.pendingPosition - this.lastQ);
        if (gap > 0) pushEvent({ type: "rest", duration: this.quartersToSeconds(gap) });
        pushEvent({
          type: "note",
          notes: this.pendingNotes.map((note) => midiToTonePitch(note.midi)),
          perNoteDurations: this.pendingNotes.map((note) => durationToSecondsFromEventSteps(note.durationSteps, this.grid)),
          velocities: this.pendingNotes.map((note) => note.velocity),
          duration: 0,
          advance: 0,
        });
        this.lastQ = this.pendingPosition;
      }
      const finalGap = Math.max(0, nextQ - this.lastQ);
      if (finalGap > 0) pushEvent({ type: "rest", duration: this.quartersToSeconds(finalGap) });
      this.lastQ = Math.max(this.lastQ, nextQ);
      this.pendingPosition = null;
      this.pendingNotes = [];
    },
    quartersToSeconds(quarters) {
      return (quarters * 60) / currentBpm;
    },
  };
}

function makeParser() {
  return makeEventParser();
}

async function warmPrompt() {
  h = zeroState();
  c = zeroState();
  
  let allTokensStr;
  if (playingStyle === "midi" && midiTokens.length > 0) {
    allTokensStr = getTokensUpToBar(midiTokens, midiStartBar);
  } else if (playingStyle === "bach") {
    allTokensStr = [
      "BOS", "BPM_100", "GRID_64",
      "BAR",
      "POS_0", "NOTE_81", "DUR_4", "VEL_6",
      "POS_4", "NOTE_79", "DUR_4", "VEL_6",
      "POS_8", "NOTE_77", "DUR_4", "VEL_6",
      "POS_12", "NOTE_76", "DUR_4", "VEL_6",
      "POS_16", "NOTE_74", "DUR_4", "VEL_6",
      "POS_20", "NOTE_73", "DUR_4", "VEL_6",
      "POS_24", "NOTE_74", "DUR_16", "VEL_6",
      "BAR"
    ];
  } else if (playingStyle === "mozart") {
    allTokensStr = [
      "BOS", "BPM_132", "GRID_64",
      "BAR",
      // Bright G-major chamber-style opening inspired by Mozart's serenades.
      "POS_0", "NOTE_55", "DUR_8", "VEL_6",
      "NOTE_67", "DUR_8", "VEL_6",
      "POS_8", "NOTE_74", "DUR_4", "VEL_7",
      "POS_12", "NOTE_71", "DUR_4", "VEL_6",
      "POS_16", "NOTE_79", "DUR_8", "VEL_7",
      "POS_24", "NOTE_78", "DUR_4", "VEL_6",
      "POS_28", "NOTE_76", "DUR_4", "VEL_6",
      "POS_32", "NOTE_74", "DUR_8", "VEL_7",
      "POS_40", "NOTE_72", "DUR_4", "VEL_6",
      "POS_44", "NOTE_71", "DUR_4", "VEL_6",
      "POS_48", "NOTE_69", "DUR_8", "VEL_6",
      "POS_56", "NOTE_67", "DUR_8", "VEL_7",
      "BAR",
      "POS_0", "NOTE_60", "DUR_16", "VEL_5",
      "NOTE_64", "DUR_16", "VEL_5",
      "NOTE_69", "DUR_16", "VEL_5",
      "POS_16", "NOTE_72", "DUR_4", "VEL_6",
      "POS_20", "NOTE_76", "DUR_4", "VEL_6",
      "POS_24", "NOTE_81", "DUR_8", "VEL_7",
      "POS_32", "NOTE_79", "DUR_4", "VEL_6",
      "POS_36", "NOTE_78", "DUR_4", "VEL_6",
      "POS_40", "NOTE_76", "DUR_4", "VEL_6",
      "POS_44", "NOTE_74", "DUR_4", "VEL_6",
      "POS_48", "NOTE_71", "DUR_8", "VEL_6",
      "POS_56", "NOTE_67", "DUR_8", "VEL_7",
      "BAR"
    ];
  } else if (playingStyle === "chopin") {
    allTokensStr = [
      "BOS", "BPM_68", "GRID_64",
      "BAR",
      // Original C-minor nocturne texture: rolling bass beneath a lyrical upper line.
      "POS_0", "NOTE_36", "DUR_16", "VEL_4", "NOTE_72", "DUR_12", "VEL_6",
      "POS_8", "NOTE_43", "DUR_8", "VEL_4",
      "POS_12", "NOTE_75", "DUR_8", "VEL_6",
      "POS_16", "NOTE_48", "DUR_8", "VEL_4",
      "POS_24", "NOTE_51", "DUR_8", "VEL_4", "NOTE_77", "DUR_12", "VEL_6",
      "POS_32", "NOTE_39", "DUR_16", "VEL_4", "NOTE_79", "DUR_16", "VEL_7",
      "POS_40", "NOTE_46", "DUR_8", "VEL_4",
      "POS_48", "NOTE_51", "DUR_8", "VEL_4", "NOTE_77", "DUR_8", "VEL_6",
      "POS_56", "NOTE_55", "DUR_8", "VEL_4", "NOTE_75", "DUR_8", "VEL_5",
      "BAR",
      "POS_0", "NOTE_41", "DUR_16", "VEL_4", "NOTE_72", "DUR_16", "VEL_6",
      "POS_8", "NOTE_48", "DUR_8", "VEL_4",
      "POS_16", "NOTE_53", "DUR_8", "VEL_4", "NOTE_74", "DUR_8", "VEL_6",
      "POS_24", "NOTE_56", "DUR_8", "VEL_4", "NOTE_75", "DUR_8", "VEL_6",
      "POS_32", "NOTE_43", "DUR_16", "VEL_4", "NOTE_79", "DUR_12", "VEL_7",
      "POS_40", "NOTE_50", "DUR_8", "VEL_4",
      "POS_44", "NOTE_77", "DUR_8", "VEL_6",
      "POS_48", "NOTE_55", "DUR_8", "VEL_4",
      "POS_56", "NOTE_60", "DUR_8", "VEL_4", "NOTE_75", "DUR_8", "VEL_5",
      "BAR"
    ];
  } else if (playingStyle === "debussy") {
    allTokensStr = [
      "BOS", "BPM_74", "GRID_64",
      "BAR",
      // Original impressionist sketch: open fifths and slowly shifting parallel colors.
      "POS_0", "NOTE_38", "DUR_32", "VEL_3", "NOTE_57", "DUR_16", "VEL_4", "NOTE_62", "DUR_16", "VEL_4", "NOTE_69", "DUR_16", "VEL_5",
      "POS_16", "NOTE_59", "DUR_16", "VEL_4", "NOTE_64", "DUR_16", "VEL_4", "NOTE_71", "DUR_16", "VEL_5",
      "POS_32", "NOTE_45", "DUR_32", "VEL_3", "NOTE_61", "DUR_16", "VEL_4", "NOTE_66", "DUR_16", "VEL_4", "NOTE_73", "DUR_16", "VEL_5",
      "POS_48", "NOTE_64", "DUR_16", "VEL_4", "NOTE_69", "DUR_16", "VEL_4", "NOTE_76", "DUR_16", "VEL_5",
      "BAR",
      "POS_0", "NOTE_40", "DUR_32", "VEL_3", "NOTE_59", "DUR_16", "VEL_4", "NOTE_64", "DUR_16", "VEL_4", "NOTE_71", "DUR_16", "VEL_5",
      "POS_16", "NOTE_61", "DUR_16", "VEL_4", "NOTE_66", "DUR_16", "VEL_4", "NOTE_73", "DUR_16", "VEL_5",
      "POS_32", "NOTE_43", "DUR_32", "VEL_3", "NOTE_62", "DUR_16", "VEL_4", "NOTE_67", "DUR_16", "VEL_4", "NOTE_74", "DUR_16", "VEL_5",
      "POS_48", "NOTE_66", "DUR_16", "VEL_4", "NOTE_71", "DUR_16", "VEL_4", "NOTE_78", "DUR_16", "VEL_5",
      "BAR"
    ];
  } else if (playingStyle === "beethoven") {
    allTokensStr = [
      "BOS", "BPM_126", "GRID_64",
      "BAR",
      // Original dramatic sonata gesture: compact motif, octave bass, accented response.
      "POS_0", "NOTE_36", "DUR_8", "VEL_7", "NOTE_48", "DUR_8", "VEL_7", "NOTE_67", "DUR_4", "VEL_7",
      "POS_4", "NOTE_70", "DUR_4", "VEL_7",
      "POS_8", "NOTE_68", "DUR_8", "VEL_8",
      "POS_16", "NOTE_43", "DUR_8", "VEL_6", "NOTE_55", "DUR_8", "VEL_6", "NOTE_65", "DUR_4", "VEL_7",
      "POS_20", "NOTE_68", "DUR_4", "VEL_7",
      "POS_24", "NOTE_67", "DUR_8", "VEL_8",
      "POS_32", "NOTE_39", "DUR_8", "VEL_7", "NOTE_51", "DUR_8", "VEL_7", "NOTE_63", "DUR_4", "VEL_7",
      "POS_36", "NOTE_67", "DUR_4", "VEL_7",
      "POS_40", "NOTE_72", "DUR_8", "VEL_8",
      "POS_48", "NOTE_44", "DUR_8", "VEL_6", "NOTE_56", "DUR_8", "VEL_6", "NOTE_70", "DUR_4", "VEL_7",
      "POS_52", "NOTE_68", "DUR_4", "VEL_7",
      "POS_56", "NOTE_67", "DUR_8", "VEL_8",
      "BAR",
      "POS_0", "NOTE_41", "DUR_16", "VEL_7", "NOTE_53", "DUR_16", "VEL_7", "NOTE_65", "DUR_8", "VEL_7",
      "POS_8", "NOTE_68", "DUR_4", "VEL_7",
      "POS_12", "NOTE_72", "DUR_4", "VEL_8",
      "POS_16", "NOTE_77", "DUR_8", "VEL_8",
      "POS_24", "NOTE_75", "DUR_8", "VEL_7",
      "POS_32", "NOTE_43", "DUR_16", "VEL_7", "NOTE_55", "DUR_16", "VEL_7", "NOTE_74", "DUR_8", "VEL_7",
      "POS_40", "NOTE_72", "DUR_8", "VEL_7",
      "POS_48", "NOTE_67", "DUR_16", "VEL_8", "NOTE_72", "DUR_16", "VEL_8", "NOTE_75", "DUR_16", "VEL_8",
      "BAR"
    ];
  } else if (playingStyle === "satie") {
    allTokensStr = [
      "BOS", "BPM_62", "GRID_64",
      "BAR",
      // Original sparse meditation: low bass, suspended chords, unhurried melody.
      "POS_0", "NOTE_38", "DUR_32", "VEL_3", "NOTE_69", "DUR_20", "VEL_5",
      "POS_16", "NOTE_57", "DUR_24", "VEL_3", "NOTE_62", "DUR_24", "VEL_3", "NOTE_65", "DUR_24", "VEL_3",
      "POS_32", "NOTE_45", "DUR_32", "VEL_3", "NOTE_72", "DUR_16", "VEL_5",
      "POS_48", "NOTE_55", "DUR_24", "VEL_3", "NOTE_60", "DUR_24", "VEL_3", "NOTE_64", "DUR_24", "VEL_3",
      "BAR",
      "POS_0", "NOTE_40", "DUR_32", "VEL_3", "NOTE_71", "DUR_20", "VEL_5",
      "POS_16", "NOTE_59", "DUR_24", "VEL_3", "NOTE_64", "DUR_24", "VEL_3", "NOTE_67", "DUR_24", "VEL_3",
      "POS_32", "NOTE_43", "DUR_32", "VEL_3", "NOTE_74", "DUR_16", "VEL_5",
      "POS_48", "NOTE_57", "DUR_24", "VEL_3", "NOTE_62", "DUR_24", "VEL_3", "NOTE_66", "DUR_24", "VEL_3",
      "BAR"
    ];
  } else {
    allTokensStr = ["BOS", "BPM_120", "GRID_64", "BAR", "POS_0"];
  }
  
  isWarmingUp = true;
  
  // 1. Feed all prompt tokens to parser to advance its clock
  if (parser) {
    for (const token of allTokensStr) {
      parser.feed(token);
    }
  }
  
  // 2. Slice model warm-up tokens to last 256 (matching model's context window)
  const modelTokensStr = allTokensStr.slice(-256);
  const ids = modelTokensStr.map(tokenId).filter((id) => id !== null);
  
  for (const id of ids) {
    currentId = id;
    await stepModel(false);
  }
  
  isWarmingUp = false;
}

async function pumpTokens(targetCount = 96) {
  tempQueue = [];
  let steps = 0;
  // Keep stepping the model until we collect enough parsed events or hit loop bounds
  while (tempQueue.length < targetCount && steps < 300) {
    await stepModel(true);
    steps += 1;
  }
  return {
    events: tempQueue,
    lastToken: itos[currentId] ?? "?"
  };
}

// Message Router
self.onmessage = async function (e) {
  const data = e.data;
  
  switch (data.action) {
    case "init":
      try {
        console.log(`Worker loading model: ${data.activeModelName}, requested EP: ${data.ep || "wasm"}`);
        activeModelName = data.activeModelName;
        stoi = data.vocab.stoi;
        itos = Object.fromEntries(Object.entries(data.vocab.itos).map(([key, value]) => [Number(key), value]));
        vocabSize = data.vocab.vocab_size;
        
        if (data.modelConfig) {
          numLayers = data.modelConfig.numLayers || 2;
          hiddenSize = data.modelConfig.hiddenSize || 2048;
        }
        
        if (session) {
          try { session.dispose(); } catch (e) {}
          session = null;
        }
        
        let epUsed = data.ep || "wasm";
        if (data.ep === "webgpu") {
          try {
            console.log("Attempting WebGPU execution provider...");
            session = await ort.InferenceSession.create(data.modelBuffer, {
              executionProviders: ["webgpu"],
            });
            console.log("WebGPU session created successfully!");
            epUsed = "webgpu";
          } catch (webgpuErr) {
            console.warn("WebGPU EP failed or unsupported by browser, falling back to CPU WASM:", webgpuErr);
            session = await ort.InferenceSession.create(data.modelBuffer, {
              executionProviders: ["wasm"],
              graphOptimizationLevel: "all",
            });
            epUsed = "wasm (CPU fallback)";
          }
        } else {
          session = await ort.InferenceSession.create(data.modelBuffer, {
            executionProviders: ["wasm"],
            graphOptimizationLevel: "all",
          });
          epUsed = "wasm";
        }
        
        updateModelDimensions(session);
        
        if (h) { try { h.dispose(); } catch (e) {} h = null; }
        if (c) { try { c.dispose(); } catch (e) {} c = null; }
        
        self.postMessage({
          action: "initialized",
          activeModelName: activeModelName,
          vocabSize: vocabSize,
          epUsed: epUsed
        });
      } catch (err) {
        console.error("Worker initialization failed:", err);
        self.postMessage({ action: "error", message: `Init failed: ${err.message}` });
      }
      break;
      
    case "start":
      try {
        temperature = data.temperature;
        topK = data.topK;
        currentBpm = data.bpm;
        playingStyle = data.playingStyle || "default";
        midiTokens = data.midiTokens;
        midiStartBar = data.midiStartBar;
        midiBpmMode = data.midiBpmMode || "lock";
        
        parser = makeParser();
        parser.reset();
        
        await warmPrompt();
        
        const initialBatch = await pumpTokens(96);
        self.postMessage({
          action: "started",
          events: initialBatch.events,
          lastToken: initialBatch.lastToken
        });
      } catch (err) {
        console.error("Worker start failed:", err);
        self.postMessage({ action: "error", message: `Start failed: ${err.message}` });
      }
      break;
      
    case "pump":
      try {
        temperature = data.temperature;
        topK = data.topK;
        const batch = await pumpTokens(96);
        self.postMessage({
          action: "events",
          events: batch.events,
          lastToken: batch.lastToken
        });
      } catch (err) {
        console.error("Worker pump failed:", err);
        self.postMessage({ action: "error", message: `Pump failed: ${err.message}` });
      }
      break;
      
    case "stop":
      // Halts ongoing actions, resets states if needed
      break;
  }
};