worker.js

// NanoMaestro Web Worker for isolated model inference and parsing
importScripts("https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/ort.min.js");

// ONNX runtime configuration
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 activeTokenizer = "abc";

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

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

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

// Seeding settings
let midiSeedActive = false;
let midiTokens = [];
let midiStartBar = 0;
let isWarmingUp = false;

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

// MIDI to Pitch Name helpers
const abcNotePattern = /^(?:\^\^|__|\^|_|=)?[A-Ga-g][,']*$/;

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 * 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);
    // Fall back to default
    numLayers = 2;
    hiddenSize = 1024;
  }
}

// Model Step Execution
async function stepModel(record = true) {
  if (!session) return "?";
  const output = await session.run({ input: makeTensorId(currentId), h, c });
  h = output.h_out;
  c = output.c_out;
  currentId = sampleFromLogits(output.logits.data);
  const token = itos[currentId] ?? "?";
  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 durationToSecondsFromAbc(duration) {
  let eighths = 1;
  if (duration) {
    if (duration === "/") eighths = 0.5;
    else if (duration.startsWith("/")) eighths = 1 / Number(duration.slice(1));
    else if (duration.includes("/")) {
      const [a, b] = duration.split("/").map(Number);
      eighths = a / b;
    } else {
      eighths = Number(duration);
    }
  }
  const quarterSeconds = 60 / currentBpm;
  return Math.max(0.08, eighths * quarterSeconds * 0.5);
}

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 abcToTonePitch(token) {
  const match = token.match(/^(\^\^|__|\^|_|=)?([A-Ga-g])([,']*)$/);
  if (!match) return null;
  const accidental = match[1] || "";
  const step = match[2];
  const marks = match[3] || "";
  let octave = step === step.toLowerCase() ? 5 : 4;
  for (const mark of marks) octave += mark === "'" ? 1 : -1;
  const accidentalText = accidental === "^^" ? "##" : accidental === "__" ? "bb" : accidental === "^" ? "#" : accidental === "_" ? "b" : "";
  return `${step.toUpperCase()}${accidentalText}${octave}`;
}

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;
}

// ABC Parser
function makeAbcParser() {
  return {
    pending: null,
    chord: null,
    reset() {
      this.pending = null;
      this.chord = null;
    },
    feed(token) {
      if (token === "Q:") {
        this.pending = { type: "tempoHeader" };
        return;
      }
      if (this.pending?.type === "tempoHeader") {
        const bpm = Number(String(token).split("=").pop());
        if (Number.isFinite(bpm) && bpm >= 40 && bpm <= 220) {
          if (!midiSeedActive || isWarmingUp) {
            currentBpm = bpm;
            self.postMessage({ action: "tempo", bpm: bpm });
          }
        }
        this.pending = null;
        return;
      }
      if (token === "[") {
        this.chord = [];
        this.pending = null;
        return;
      }
      if (this.chord) {
        if (token === "]") {
          this.pending = { type: "chord", notes: this.chord.map(abcToTonePitch).filter(Boolean) };
          this.chord = null;
          return;
        }
        if (abcNotePattern.test(token)) this.chord.push(token);
        return;
      }
      if (token === "z") {
        this.pending = { type: "rest" };
        return;
      }
      if (abcNotePattern.test(token)) {
        this.pending = { type: "note", note: abcToTonePitch(token) };
        return;
      }
      if (token.startsWith("DUR:") && this.pending) {
        this.emitPending(token.slice(4));
        return;
      }
      if (this.pending && (token === "|" || token === "<NL>" || token === "<EOP>")) this.emitPending(null);
    },
    emitPending(durationToken) {
      const duration = durationToSecondsFromAbc(durationToken);
      if (this.pending.type === "rest") pushEvent({ type: "rest", duration });
      if (this.pending.type === "note" && this.pending.note) pushEvent({ type: "note", notes: [this.pending.note], duration, advance: duration });
      if (this.pending.type === "chord" && this.pending.notes.length) pushEvent({ type: "note", notes: this.pending.notes.slice(0, 8), duration, advance: duration });
      this.pending = null;
    },
  };
}

// 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 (!midiSeedActive || isWarmingUp) {
            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 activeTokenizer === "giantmidi_event" ? makeEventParser() : makeAbcParser();
}

async function warmPrompt() {
  h = zeroState();
  c = zeroState();
  
  const allTokensStr = midiSeedActive && midiTokens.length > 0 
    ? getTokensUpToBar(midiTokens, midiStartBar) 
    : (activeTokenizer === "giantmidi_event" 
        ? ["BOS", "BPM_120", "GRID_64", "BAR", "POS_0"] 
        : ["X:", "1", "<NL>", "T:", "piece", "<NL>", "M:", "4/4", "<NL>", "L:", "1/8", "<NL>", "Q:", "1/4=120", "<NL>", "K:", "C", "<NL>"]);
  
  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}`);
        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;
        activeTokenizer = data.vocab.tokenizer || "abc";
        
        session = await ort.InferenceSession.create(data.modelBuffer, {
          executionProviders: ["wasm"],
          graphOptimizationLevel: "all",
        });
        
        updateModelDimensions(session);
        self.postMessage({ action: "initialized" });
      } 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;
        midiSeedActive = data.midiSeedActive;
        midiTokens = data.midiTokens;
        midiStartBar = data.midiStartBar;
        
        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;
  }
};