let hiddenSize = 1024;
let numLayers = 2;
const TARGET_QUEUE_EVENTS = 96;
const MIN_QUEUE_EVENTS = 48;
const LOOKAHEAD_MS = 120;
const SCHEDULE_AHEAD_SECONDS = 4.0;
const STEPS_PER_PUMP = 128;

// DOM Selectors
const appStatusEl = document.querySelector("#appStatus");
const statusDotEl = document.querySelector("#statusDot");
const statusTextEl = document.querySelector("#statusText");

const modelFileInput = document.querySelector("#modelFile");
const vocabFileInput = document.querySelector("#vocabFile");
const startBtn = document.querySelector("#startBtn");
const stopBtn = document.querySelector("#stopBtn");
const newSongBtn = document.querySelector("#newSongBtn");
const recordBtn = document.querySelector("#recordBtn");
const recordIcon = document.querySelector("#recordIcon");
const recordBtnText = document.querySelector("#recordBtnText");

const tempInput = document.querySelector("#temperature");
const tempValue = document.querySelector("#temperatureValue");
const topKInput = document.querySelector("#topK");
const topKValue = document.querySelector("#topKValue");

const volumeInput = document.querySelector("#masterVolume");
const volumeValue = document.querySelector("#volumeValue");
const reverbInput = document.querySelector("#reverbWet");
const reverbValue = document.querySelector("#reverbValue");
const delayInput = document.querySelector("#delayWet");
const delayValue = document.querySelector("#delayValue");
const releaseInput = document.querySelector("#synthRelease");
const releaseValue = document.querySelector("#releaseValue");

const lastTokenEl = document.querySelector("#lastToken");
const eventCountEl = document.querySelector("#eventCount");
const queueCountEl = document.querySelector("#queueCount");
const queueFillEl = document.querySelector("#queueFill");
const tempoEl = document.querySelector("#tempo");
const noteCanvas = document.querySelector("#noteCanvas");
const canvasCtx = noteCanvas.getContext("2d");

// Inference & Audio Variables
let activeModelName = "";
let isModelReady = false;
let running = false;
let pumping = false;
let schedulerId = 0;

// Web Worker instance & routing
const worker = new Worker("./worker.js");
let modelInitResolve = null;
let modelInitReject = null;
let startResolve = null;
let startReject = null;

worker.onmessage = function (e) {
  const data = e.data;
  switch (data.action) {
    case "initialized":
      isModelReady = true;
      if (modelInitResolve) modelInitResolve();
      break;
    case "started":
      if (data.events) {
        data.events.forEach(event => pushEvent(event));
      }
      if (data.lastToken) {
        updateReadouts(data.lastToken);
      }
      if (startResolve) startResolve();
      break;
    case "events":
      if (data.events) {
        data.events.forEach(event => pushEvent(event));
      }
      if (data.lastToken) {
        updateReadouts(data.lastToken);
      }
      pumping = false;
      break;
    case "tempo":
      if (tempoEl) tempoEl.textContent = String(Math.round(data.bpm));
      break;
    case "error":
      console.error("Worker error:", data.message);
      setStatus(data.message, "idle");
      stop();
      if (modelInitReject) modelInitReject(new Error(data.message));
      if (startReject) startReject(new Error(data.message));
      break;
  }
};

worker.onerror = function (err) {
  console.error("Web Worker uncaught error:", err);
  const errMsg = err.message || "Web Worker script runtime error";
  setStatus(`Worker error: ${errMsg}`, "idle");
  stop();
  if (modelInitReject) modelInitReject(new Error(errMsg));
  if (startReject) startReject(new Error(errMsg));
};
let nextNoteTime = 0;
let generatedEvents = 0;
let queue = [];
let visualNotes = [];
let synth;
let reverb;
let delay;
let visualFrame = 0;
let canvasWidth = 0;
let canvasHeight = 0;
let parser;
let lastAudioTime = 0;
let lastPerfTime = 0;
const VISUAL_DELAY = 0.06; // 60ms audio latency compensation

// MIDI Seeding Elements & Variables
const seedSourceGroup = document.querySelector("#seedSourceGroup");
const toggleButtons = seedSourceGroup ? seedSourceGroup.querySelectorAll(".toggle-btn") : [];
const midiSeedControls = document.querySelector("#midiSeedControls");
const midiFileInput = document.querySelector("#midiFileInput");
const midiFileInfo = document.querySelector("#midiFileInfo");
const midiFileNameEl = document.querySelector("#midiFileName");
const midiFileStatsEl = document.querySelector("#midiFileStats");
const midiStartBarInput = document.querySelector("#midiStartBar");
const midiStartBarValue = document.querySelector("#midiStartBarValue");
const midiStartBarTimeEl = document.querySelector("#midiStartBarTime");
const midiBpmModeGroup = document.querySelector("#midiBpmModeGroup");
const midiBpmModeButtons = midiBpmModeGroup ? midiBpmModeGroup.querySelectorAll(".toggle-btn") : [];

let activePlayingStyle = "beethoven";
let midiBpmMode = "lock";
let midiTokens = [];
let midiStartBar = 0;
let midiFileName = "";
let midiBpm = 120;
let midiTotalBars = 0;
let isWarmingUp = false;
let isRecording = false;
let recordingStartTime = 0;
let recordedNotes = [];

// Binary MIDI Parser
class MidiParser {
  constructor(arrayBuffer) {
    this.view = new DataView(arrayBuffer);
    this.pos = 0;
  }
  
  readUint8() {
    const v = this.view.getUint8(this.pos);
    this.pos += 1;
    return v;
  }
  
  readUint16() {
    const v = this.view.getUint16(this.pos);
    this.pos += 2;
    return v;
  }
  
  readUint32() {
    const v = this.view.getUint32(this.pos);
    this.pos += 4;
    return v;
  }
  
  readBytes(len) {
    const buf = new Uint8Array(this.view.buffer, this.pos + this.view.byteOffset, len);
    this.pos += len;
    return buf;
  }
  
  readVarInt() {
    let value = 0;
    while (true) {
      const b = this.readUint8();
      value = (value << 7) | (b & 0x7F);
      if (!(b & 0x80)) break;
    }
    return value;
  }
  
  parse() {
    const mthd = String.fromCharCode(...this.readBytes(4));
    if (mthd !== "MThd") throw new Error("Not a valid MIDI file (missing MThd)");
    
    const headerLength = this.readUint32();
    const format = this.readUint16();
    const numTracks = this.readUint16();
    const division = this.readUint16();
    
    if (headerLength > 6) {
      this.readBytes(headerLength - 6);
    }
    
    let bpm = 120;
    const notes = [];
    const activeNotes = new Map();
    
    for (let t = 0; t < numTracks; t++) {
      const mtrk = String.fromCharCode(...this.readBytes(4));
      if (mtrk !== "MTrk") {
        const len = this.readUint32();
        this.readBytes(len);
        continue;
      }
      
      const trackLength = this.readUint32();
      const endPos = this.pos + trackLength;
      
      let tick = 0;
      let runningStatus = 0;
      
      while (this.pos < endPos) {
        const deltaTime = this.readVarInt();
        tick += deltaTime;
        
        let status = this.readUint8();
        if (status < 0x80) {
          this.pos -= 1;
          status = runningStatus;
        } else {
          runningStatus = status;
        }
        
        const eventType = status & 0xF0;
        const channel = status & 0x0F;
        
        if (eventType === 0x90) {
          const pitch = this.readUint8();
          const velocity = this.readUint8();
          const key = `${pitch}_${channel}`;
          if (velocity > 0) {
            if (!activeNotes.has(key)) {
              activeNotes.set(key, { startTick: tick, velocity });
            }
          } else {
            const active = activeNotes.get(key);
            if (active) {
              notes.push({
                pitch,
                startTick: active.startTick,
                endTick: tick,
                velocity: active.velocity
              });
              activeNotes.delete(key);
            }
          }
        } else if (eventType === 0x80) {
          const pitch = this.readUint8();
          const velocity = this.readUint8();
          const key = `${pitch}_${channel}`;
          const active = activeNotes.get(key);
          if (active) {
            notes.push({
              pitch,
              startTick: active.startTick,
              endTick: tick,
              velocity: active.velocity
            });
            activeNotes.delete(key);
          }
        } else if (status === 0xFF) {
          const metaType = this.readUint8();
          const len = this.readVarInt();
          const data = this.readBytes(len);
          
          if (metaType === 0x51 && len === 3) {
            const tempo = (data[0] << 16) | (data[1] << 8) | data[2];
            bpm = Math.round(60000000 / tempo);
          }
        } else if (eventType === 0xA0 || eventType === 0xB0 || eventType === 0xE0) {
          this.readBytes(2);
        } else if (eventType === 0xC0 || eventType === 0xD0) {
          this.readBytes(1);
        } else if (status === 0xF0 || status === 0xF7) {
          const len = this.readVarInt();
          this.readBytes(len);
        }
      }
    }
    
    for (const [key, active] of activeNotes.entries()) {
      const [pitch, channel] = key.split("_").map(Number);
      notes.push({
        pitch,
        startTick: active.startTick,
        endTick: active.startTick + division,
        velocity: active.velocity
      });
    }
    
    return { bpm, division, notes };
  }
}

// Tokenize MIDI to Model-Compatible Event Tokens
function tokenizeMidi(parsedMidi, grid = 64) {
  const { bpm, division, notes } = parsedMidi;
  const stepsPerQuarter = grid / 4;
  
  const quantizedNotes = notes.map(note => {
    const start_step = Math.round((note.startTick / division) * stepsPerQuarter);
    const end_step = Math.round((note.endTick / division) * stepsPerQuarter);
    const duration_steps = Math.max(1, end_step - start_step);
    const velocity_bucket = Math.max(1, Math.min(8, Math.ceil(note.velocity / 16)));
    
    return {
      pitch: note.pitch,
      start_step,
      duration_steps,
      velocity_bucket
    };
  });
  
  quantizedNotes.sort((a, b) => {
    if (a.start_step !== b.start_step) return a.start_step - b.start_step;
    return a.pitch - b.pitch;
  });
  
  const tokens = ["BOS", `BPM_${Math.round(bpm)}`, `GRID_${grid}`];
  let currentBar = -1;
  
  const notesByStep = new Map();
  for (const note of quantizedNotes) {
    if (!notesByStep.has(note.start_step)) {
      notesByStep.set(note.start_step, []);
    }
    notesByStep.get(note.start_step).push(note);
  }
  
  const steps = Array.from(notesByStep.keys()).sort((a, b) => a - b);
  for (const step of steps) {
    const note_bar = Math.floor(step / grid);
    const note_pos = step % grid;
    
    while (currentBar < note_bar) {
      tokens.push("BAR");
      currentBar += 1;
    }
    
    tokens.push(`POS_${note_pos}`);
    
    const stepNotes = notesByStep.get(step);
    for (const note of stepNotes) {
      tokens.push(`NOTE_${note.pitch}`);
      tokens.push(`DUR_${Math.min(256, note.duration_steps)}`);
      tokens.push(`VEL_${note.velocity_bucket}`);
    }
  }
  
  return { tokens, maxBar: currentBar };
}

// Slice prompt tokens up to target bar (exclusive start of target bar)
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;
}

// Find the first bar that has at least 256 tokens preceding it
function findMinStartBar(tokens) {
  let barCount = 0;
  let tokenCount = 0;
  
  for (let i = 0; i < tokens.length; i++) {
    const t = tokens[i];
    if (t === "BAR") {
      if (tokenCount >= 256) {
        return barCount;
      }
      barCount += 1;
    }
    tokenCount += 1;
  }
  return 0;
}

// Analyse tokens to get per-bar metadata (token counts, note densities)
function analyzeMidiTokens(tokens) {
  const barsData = [];
  let currentBar = -1;
  let barTokensCount = 0;
  let barNotesCount = 0;

  for (const t of tokens) {
    if (t === "BAR") {
      if (currentBar >= 0) {
        barsData.push({
          barNum: currentBar,
          tokenCount: barTokensCount,
          noteCount: barNotesCount
        });
      }
      currentBar += 1;
      barTokensCount = 1;
      barNotesCount = 0;
    } else {
      barTokensCount += 1;
      if (t.startsWith("NOTE_")) {
        barNotesCount += 1;
      }
    }
  }

  if (currentBar >= 0) {
    barsData.push({
      barNum: currentBar,
      tokenCount: barTokensCount,
      noteCount: barNotesCount
    });
  }

  return barsData;
}

// Calculate which bars fall in the 256-token context window
function getWarmUpInfo(tokens, targetBar) {
  let barCount = 0;
  const tokensUpToBar = [];
  
  for (const t of tokens) {
    if (t === "BAR") {
      if (barCount === targetBar) {
        break;
      }
      barCount += 1;
    }
    tokensUpToBar.push(t);
  }
  
  const totalTokens = tokensUpToBar.length;
  const warmUpCount = Math.min(256, totalTokens);
  const startIndex = totalTokens - warmUpCount;
  
  let currentBarOfStart = -1;
  for (let i = 0; i < startIndex; i++) {
    if (tokens[i] === "BAR") {
      currentBarOfStart += 1;
    }
  }
  
  return {
    startBar: Math.max(0, currentBarOfStart),
    endBar: targetBar - 1,
    tokenCount: warmUpCount,
    totalTokens: totalTokens
  };
}

// Render the visual timeline bars
function renderMidiTimeline(barsData) {
  const midiTimeline = document.querySelector("#midiTimeline");
  if (!midiTimeline) return;
  midiTimeline.innerHTML = "";
  
  if (!barsData || barsData.length === 0) {
    midiTimeline.style.display = "none";
    return;
  }
  
  midiTimeline.style.display = "flex";
  
  const maxNotes = Math.max(...barsData.map(b => b.noteCount), 1);
  const minBar = findMinStartBar(midiTokens);
  
  barsData.forEach((bar, index) => {
    const col = document.createElement("div");
    col.className = "timeline-bar-col";
    col.dataset.barNum = String(index);
    col.dataset.barLabel = `${index + 1}`;
    
    // Label every Nth bar to prevent crowding
    const totalBars = barsData.length;
    const labelInterval = totalBars > 64 ? 16 : (totalBars > 32 ? 8 : 4);
    if (index % labelInterval === 0) {
      col.classList.add("labeled");
    }
    
    const densityVal = (bar.noteCount / maxNotes) * 100;
    
    const fill = document.createElement("div");
    fill.className = "density-fill";
    fill.style.height = `${Math.max(10, densityVal)}%`;
    
    col.appendChild(fill);
    
    if (index < minBar) {
      col.classList.add("disabled");
    } else {
      col.addEventListener("click", () => {
        midiStartBarInput.value = String(index);
        midiStartBar = index;
        midiStartBarValue.textContent = String(index);
        updateMidiTimelineHighlights();
        updateMidiStartBarText();
      });
    }
    
    midiTimeline.appendChild(col);
  });
  
  updateMidiTimelineHighlights();
}

// Refresh visual highlighting on the timeline
function updateMidiTimelineHighlights() {
  const midiTimeline = document.querySelector("#midiTimeline");
  if (!midiTimeline || !midiTokens || midiTokens.length === 0) return;
  
  const cols = midiTimeline.querySelectorAll(".timeline-bar-col");
  const info = getWarmUpInfo(midiTokens, midiStartBar);
  
  cols.forEach((col, index) => {
    col.classList.remove("in-warmup", "is-selected");
    
    if (index === midiStartBar) {
      col.classList.add("is-selected");
      // Scroll active bar into view if timeline overflows
      if (typeof col.scrollIntoViewIfNeeded === "function") {
        col.scrollIntoViewIfNeeded({ behavior: "smooth", block: "nearest" });
      } else {
        col.scrollIntoView({ behavior: "smooth", block: "nearest", inline: "nearest" });
      }
    } else if (midiStartBar > 0 && index >= info.startBar && index <= info.endBar) {
      col.classList.add("in-warmup");
    }
  });
}

// Update text detail about context window priming
function updateMidiStartBarText() {
  if (!midiStartBarTimeEl) return;
  if (midiStartBar === 0) {
    midiStartBarTimeEl.textContent = "Will prime model with bars 0 to -1 (empty prompt)";
  } else {
    const info = getWarmUpInfo(midiTokens, midiStartBar);
    if (info.startBar === 0) {
      midiStartBarTimeEl.textContent = `Bars 0 to ${info.endBar} (${info.tokenCount} tokens) will prime the model context window.`;
    } else {
      midiStartBarTimeEl.textContent = `Bars ${info.startBar} to ${info.endBar} (${info.tokenCount} tokens) will prime the model context window (older bars truncated).`;
    }
  }
}

const pianoSamples = {
  A0: "A0.mp3",
  C1: "C1.mp3",
  "D#1": "Ds1.mp3",
  "F#1": "Fs1.mp3",
  A1: "A1.mp3",
  C2: "C2.mp3",
  "D#2": "Ds2.mp3",
  "F#2": "Fs2.mp3",
  A2: "A2.mp3",
  C3: "C3.mp3",
  "D#3": "Ds3.mp3",
  "F#3": "Fs3.mp3",
  A3: "A3.mp3",
  C4: "C4.mp3",
  "D#4": "Ds4.mp3",
  "F#4": "Fs4.mp3",
  A4: "A4.mp3",
  C5: "C5.mp3",
  "D#5": "Ds5.mp3",
  "F#5": "Fs5.mp3",
  A5: "A5.mp3",
  C6: "C6.mp3",
  "D#6": "Ds6.mp3",
  "F#6": "Fs6.mp3",
  A6: "A6.mp3",
  C7: "C7.mp3",
  "D#7": "Ds7.mp3",
  "F#7": "Fs7.mp3",
  A7: "A7.mp3",
  C8: "C8.mp3",
};

let currentInstrument = "piano";

const instrumentSamples = {
  xylophone: {
    'C8': 'C8.mp3',
    'G4': 'G4.mp3',
    'G5': 'G5.mp3',
    'G6': 'G6.mp3',
    'G7': 'G7.mp3',
    'C5': 'C5.mp3',
    'C6': 'C6.mp3',
    'C7': 'C7.mp3'
  },
  violin: {
    'A3': 'A3.mp3',
    'A4': 'A4.mp3',
    'A5': 'A5.mp3',
    'A6': 'A6.mp3',
    'C4': 'C4.mp3',
    'C5': 'C5.mp3',
    'C6': 'C6.mp3',
    'C7': 'C7.mp3',
    'E4': 'E4.mp3',
    'E5': 'E5.mp3',
    'E6': 'E6.mp3',
    'G4': 'G4.mp3',
    'G5': 'G5.mp3',
    'G6': 'G6.mp3'
  },
  trumpet: {
    'C6': 'C6.mp3',
    'D5': 'D5.mp3',
    'D#4': 'Ds4.mp3',
    'F3': 'F3.mp3',
    'F4': 'F4.mp3',
    'F5': 'F5.mp3',
    'G4': 'G4.mp3',
    'A3': 'A3.mp3',
    'A5': 'A5.mp3',
    'A#4': 'As4.mp3',
    'C4': 'C4.mp3'
  },
  saxophone: {
    'D#5': 'Ds5.mp3',
    'E3': 'E3.mp3',
    'E4': 'E4.mp3',
    'E5': 'E5.mp3',
    'F3': 'F3.mp3',
    'F4': 'F4.mp3',
    'F5': 'F5.mp3',
    'F#3': 'Fs3.mp3',
    'F#4': 'Fs4.mp3',
    'F#5': 'Fs5.mp3',
    'G3': 'G3.mp3',
    'G4': 'G4.mp3',
    'G5': 'G5.mp3',
    'G#3': 'Gs3.mp3',
    'G#4': 'Gs4.mp3',
    'G#5': 'Gs5.mp3',
    'A4': 'A4.mp3',
    'A5': 'A5.mp3',
    'A#3': 'As3.mp3',
    'A#4': 'As4.mp3',
    'B3': 'B3.mp3',
    'B4': 'B4.mp3',
    'C4': 'C4.mp3',
    'C5': 'C5.mp3',
    'C#3': 'Cs3.mp3',
    'C#4': 'Cs4.mp3',
    'C#5': 'Cs5.mp3',
    'D3': 'D3.mp3',
    'D4': 'D4.mp3',
    'D5': 'D5.mp3',
    'D#3': 'Ds3.mp3',
    'D#4': 'Ds4.mp3'
  },
  organ: {
    'C3': 'C3.mp3',
    'C4': 'C4.mp3',
    'C5': 'C5.mp3',
    'C6': 'C6.mp3',
    'D#1': 'Ds1.mp3',
    'D#2': 'Ds2.mp3',
    'D#3': 'Ds3.mp3',
    'D#4': 'Ds4.mp3',
    'D#5': 'Ds5.mp3',
    'F#1': 'Fs1.mp3',
    'F#2': 'Fs2.mp3',
    'F#3': 'Fs3.mp3',
    'F#4': 'Fs4.mp3',
    'F#5': 'Fs5.mp3',
    'A1': 'A1.mp3',
    'A2': 'A2.mp3',
    'A3': 'A3.mp3',
    'A4': 'A4.mp3',
    'A5': 'A5.mp3',
    'C1': 'C1.mp3',
    'C2': 'C2.mp3'
  },
  harp: {
    'C5': 'C5.mp3',
    'D2': 'D2.mp3',
    'D4': 'D4.mp3',
    'D6': 'D6.mp3',
    'D7': 'D7.mp3',
    'E1': 'E1.mp3',
    'E3': 'E3.mp3',
    'E5': 'E5.mp3',
    'F2': 'F2.mp3',
    'F4': 'F4.mp3',
    'F6': 'F6.mp3',
    'F7': 'F7.mp3',
    'G1': 'G1.mp3',
    'G3': 'G3.mp3',
    'G5': 'G5.mp3',
    'A2': 'A2.mp3',
    'A4': 'A4.mp3',
    'A6': 'A6.mp3',
    'B1': 'B1.mp3',
    'B3': 'B3.mp3',
    'B5': 'B5.mp3',
    'B6': 'B6.mp3',
    'C3': 'C3.mp3'
  },
  'guitar-acoustic': {
    'F4': 'F4.mp3',
    'F#2': 'Fs2.mp3',
    'F#3': 'Fs3.mp3',
    'F#4': 'Fs4.mp3',
    'G2': 'G2.mp3',
    'G3': 'G3.mp3',
    'G4': 'G4.mp3',
    'G#2': 'Gs2.mp3',
    'G#3': 'Gs3.mp3',
    'G#4': 'Gs4.mp3',
    'A2': 'A2.mp3',
    'A3': 'A3.mp3',
    'A4': 'A4.mp3',
    'A#2': 'As2.mp3',
    'A#3': 'As3.mp3',
    'A#4': 'As4.mp3',
    'B2': 'B2.mp3',
    'B3': 'B3.mp3',
    'B4': 'B4.mp3',
    'C3': 'C3.mp3',
    'C4': 'C4.mp3',
    'C5': 'C5.mp3',
    'C#3': 'Cs3.mp3',
    'C#4': 'Cs4.mp3',
    'C#5': 'Cs5.mp3',
    'D2': 'D2.mp3',
    'D3': 'D3.mp3',
    'D4': 'D4.mp3',
    'D5': 'D5.mp3',
    'D#2': 'Ds2.mp3',
    'D#3': 'Ds3.mp3',
    'D#4': 'Ds3.mp3',
    'E2': 'E2.mp3',
    'E3': 'E3.mp3',
    'E4': 'E4.mp3',
    'F2': 'F2.mp3',
    'F3': 'F3.mp3'
  },
  'guitar-electric': {
    'D#3': 'Ds3.mp3',
    'D#4': 'Ds4.mp3',
    'D#5': 'Ds5.mp3',
    'E2': 'E2.mp3',
    'F#2': 'Fs2.mp3',
    'F#3': 'Fs3.mp3',
    'F#4': 'Fs4.mp3',
    'F#5': 'Fs5.mp3',
    'A2': 'A2.mp3',
    'A3': 'A3.mp3',
    'A4': 'A4.mp3',
    'A5': 'A5.mp3',
    'C3': 'C3.mp3',
    'C4': 'C4.mp3',
    'C5': 'C5.mp3',
    'C6': 'C6.mp3',
    'C#2': 'Cs2.mp3'
  },
  'bass-electric': {
    'A#1': 'As1.mp3',
    'A#2': 'As2.mp3',
    'A#3': 'As3.mp3',
    'A#4': 'As4.mp3',
    'C#1': 'Cs1.mp3',
    'C#2': 'Cs2.mp3',
    'C#3': 'Cs3.mp3',
    'C#4': 'Cs4.mp3',
    'E1': 'E1.mp3',
    'E2': 'E2.mp3',
    'E3': 'E3.mp3',
    'E4': 'E4.mp3',
    'G1': 'G1.mp3',
    'G2': 'G2.mp3',
    'G3': 'G3.mp3',
    'G4': 'G4.mp3'
  }
};

// Status Handler
function setStatus(message, state = "idle") {
  if (appStatusEl) appStatusEl.textContent = message;
  if (statusTextEl) statusTextEl.textContent = state.toUpperCase();
  
  if (statusDotEl) {
    statusDotEl.className = "status-dot";
    if (state === "playing") {
      statusDotEl.classList.add("playing");
    } else if (state === "buffering" || state === "loading") {
      statusDotEl.classList.add("buffering");
    }
  }
}

// UI readouts update
function updateReadouts(lastToken = null) {
  if (lastToken !== null && lastTokenEl) lastTokenEl.textContent = lastToken;
  if (eventCountEl) eventCountEl.textContent = String(generatedEvents);
  if (queueCountEl) queueCountEl.textContent = String(queue.length);
  if (queueFillEl) {
    queueFillEl.style.width = `${Math.min(100, (queue.length / TARGET_QUEUE_EVENTS) * 100)}%`;
  }
  if (tempValue) tempValue.textContent = Number(tempInput.value).toFixed(2);
  if (topKValue) topKValue.textContent = topKInput.value;
}

function midiFromPitchName(note) {
  const match = String(note).match(/^([A-G])([#b]{0,2})(-?\d+)$/);
  if (!match) return null;
  const semitone = { C: 0, D: 2, E: 4, F: 5, G: 7, A: 9, B: 11 }[match[1]];
  const accidental = match[2].split("").reduce((sum, char) => sum + (char === "#" ? 1 : -1), 0);
  return (Number(match[3]) + 1) * 12 + semitone + accidental;
}

function recordVisualEvent(event, startTime) {
  if (event.type !== "note") return;
  event.notes.forEach((note, index) => {
    const midi = midiFromPitchName(note);
    if (midi === null) return;
    const duration = event.perNoteDurations 
      ? event.perNoteDurations[index] 
      : event.duration;
    visualNotes.push({ 
      midi, 
      note, 
      start: startTime, 
      duration: Math.max(0.08, duration) 
    });
  });
  const now = Tone.now();
  visualNotes = visualNotes.filter((item) => item.start + item.duration > now - 1.5);
}

function pushEvent(event) {
  if (event.duration < 0 || queue.length > TARGET_QUEUE_EVENTS * 3) return;
  queue.push(event);
  generatedEvents += event.type === "note" ? 1 : 0;
  updateReadouts();
}

function pumpTokens() {
  if (!running || pumping || queue.length >= TARGET_QUEUE_EVENTS) return;
  pumping = true;
  worker.postMessage({
    action: "pump",
    temperature: Math.max(0.05, Number(tempInput.value)),
    topK: Math.max(1, Number(topKInput.value))
  });
}

function scheduleAudio() {
  if (!running) return;
  const now = Tone.now();
  if (nextNoteTime < now + 0.08) nextNoteTime = now + 0.08;
  while (queue.length && nextNoteTime < now + SCHEDULE_AHEAD_SECONDS) {
    const event = queue.shift();
    if (event.type === "note") {
      if (event.perNoteDurations) {
        event.notes.forEach((note, index) => {
          synth.triggerAttackRelease(note, Math.min(3.8, event.perNoteDurations[index]), nextNoteTime, event.velocities?.[index] ?? 0.72);
        });
      } else {
        synth.triggerAttackRelease(event.notes, Math.min(3.2, event.duration * 0.92), nextNoteTime, 0.72);
      }
      
      if (isRecording) {
        const eventTime = nextNoteTime;
        if (event.perNoteDurations) {
          event.notes.forEach((note, index) => {
            const midi = midiFromPitchName(note);
            if (midi !== null) {
              recordedNotes.push({
                midi,
                velocity: event.velocities?.[index] ?? 0.72,
                time: eventTime - recordingStartTime,
                duration: event.perNoteDurations[index]
              });
            }
          });
        } else {
          event.notes.forEach((note) => {
            const midi = midiFromPitchName(note);
            if (midi !== null) {
              recordedNotes.push({
                midi,
                velocity: 0.72,
                time: eventTime - recordingStartTime,
                duration: event.duration
              });
            }
          });
        }
      }
      
      recordVisualEvent({ ...event, duration: event.duration || Math.max(...(event.perNoteDurations ?? [0.2])) }, nextNoteTime);
    }
    nextNoteTime += event.advance ?? event.duration;
  }
  if (queue.length < MIN_QUEUE_EVENTS) void pumpTokens();
  updateReadouts();
  schedulerId = window.setTimeout(scheduleAudio, LOOKAHEAD_MS);
}

// Audio Node Setters
function setVolume(pct) {
  if (!Tone) return;
  if (pct === 0) {
    Tone.getDestination().volume.value = -999;
  } else {
    // Map 0-100 percentage to decibels
    Tone.getDestination().volume.value = 20 * Math.log10(pct / 100);
  }
  if (volumeValue) volumeValue.textContent = `${pct}%`;
}

function setRelease(val) {
  if (synth) {
    synth.release = Number(val);
  }
  if (releaseValue) releaseValue.textContent = `${Number(val).toFixed(1)}s`;
}

async function createSampler(instrumentName) {
  if (synth) {
    try {
      synth.dispose();
    } catch (e) {}
    synth = null;
  }
  
  const initialRel = Number(releaseInput.value);
  
  let urls, baseUrl;
  if (instrumentName === "piano") {
    urls = pianoSamples;
    baseUrl = "https://tonejs.github.io/audio/salamander/";
  } else {
    urls = instrumentSamples[instrumentName];
    baseUrl = `https://nbrosowsky.github.io/tonejs-instruments/samples/${instrumentName}/`;
  }
  
  const originalStatus = appStatusEl ? appStatusEl.textContent : "";
  const originalState = statusTextEl ? statusTextEl.textContent : "IDLE";
  
  setStatus(`Loading ${instrumentName} audio samples...`, "loading");
  
  synth = new Tone.Sampler({
    urls: urls,
    baseUrl: baseUrl,
    attack: 0,
    release: initialRel,
    curve: "exponential",
  }).connect(delay);
  
  await Tone.loaded();
  setStatus(originalStatus.includes("Loading") ? "Ready to perform" : originalStatus, originalState.toLowerCase());
}

async function ensureInstrument() {
  if (synth) return;
  
  const initialVol = Number(volumeInput.value);
  const initialRev = Number(reverbInput.value);
  const initialDel = Number(delayInput.value);
  const initialRel = Number(releaseInput.value);

  reverb = new Tone.Reverb({ decay: 2.8, wet: initialRev / 100 }).toDestination();
  delay = new Tone.FeedbackDelay({ delayTime: "8n", feedback: 0.18, wet: initialDel / 100 }).connect(reverb);
  
  await createSampler(currentInstrument);
  
  setVolume(initialVol);
  
  // Set UI readouts
  if (reverbValue) reverbValue.textContent = `${initialRev}%`;
  if (delayValue) delayValue.textContent = `${initialDel}%`;
  if (releaseValue) releaseValue.textContent = `${initialRel.toFixed(1)}s`;
}

async function start() {
  if (running || !isModelReady) return;
  startBtn.disabled = true;
  setStatus("Starting audio...", "loading");
  
  await Tone.start();
  await ensureInstrument();
  
  queue = [];
  generatedEvents = 0;
  visualNotes = [];
  if (tempoEl) tempoEl.textContent = (activePlayingStyle === "midi") ? String(Math.round(midiBpm)) : "120";
  
  setStatus("Warming AI model prompt...", "buffering");
  
  await new Promise((resolve, reject) => {
    startResolve = resolve;
    startReject = reject;
    worker.postMessage({
      action: "start",
      temperature: Math.max(0.05, Number(tempInput.value)),
      topK: Math.max(1, Number(topKInput.value)),
      bpm: Number(tempoEl.textContent || 120),
      playingStyle: activePlayingStyle,
      midiTokens: (activePlayingStyle === "midi") ? midiTokens : [],
      midiStartBar: midiStartBar,
      midiBpmMode: midiBpmMode
    });
  });
  
  running = true;
  stopBtn.disabled = false;
  if (newSongBtn) newSongBtn.disabled = false;
  if (recordBtn) recordBtn.disabled = false;
  if (modelFileInput) modelFileInput.disabled = true;
  if (vocabFileInput) vocabFileInput.disabled = true;
  toggleButtons.forEach(btn => btn.disabled = true);
  midiBpmModeButtons.forEach(btn => btn.disabled = true);
  midiFileInput.disabled = true;
  midiStartBarInput.disabled = true;
  
  setStatus(`Playing ${activeModelName}`, "playing");
  
  nextNoteTime = Tone.now() + 0.12;
  startVisualizer();
  scheduleAudio();
}

function stop() {
  if (isRecording) {
    stopRecording(true);
  }
  running = false;
  window.clearTimeout(schedulerId);
  schedulerId = 0;
  
  worker.postMessage({ action: "stop" });
  
  startBtn.disabled = !isModelReady;
  stopBtn.disabled = true;
  if (newSongBtn) newSongBtn.disabled = true;
  if (recordBtn) recordBtn.disabled = true;
  if (modelFileInput) modelFileInput.disabled = false;
  if (vocabFileInput) vocabFileInput.disabled = false;
  toggleButtons.forEach(btn => btn.disabled = false);
  midiBpmModeButtons.forEach(btn => btn.disabled = false);
  midiFileInput.disabled = false;
  midiStartBarInput.disabled = false;
  if (synth) {
    try { synth.dispose(); } catch (e) { console.error(e); }
    synth = null;
  }
  if (reverb) {
    try { reverb.dispose(); } catch (e) { console.error(e); }
    reverb = null;
  }
  if (delay) {
    try { delay.dispose(); } catch (e) { console.error(e); }
    delay = null;
  }
  if (visualFrame) cancelAnimationFrame(visualFrame);
  visualFrame = 0;
  setStatus(isModelReady ? `Ready: ${activeModelName}` : "Stopped", "idle");
}

// MIDI Recording functions
function startRecording() {
  if (isRecording || !running) return;
  isRecording = true;
  recordedNotes = [];
  recordingStartTime = Tone.now();
  
  if (recordBtn) {
    recordBtn.classList.add("btn-recording");
    if (recordIcon) recordIcon.classList.add("pulse-record");
    if (recordIcon) {
      recordIcon.innerHTML = `<rect x="6" y="6" width="12" height="12" fill="#d9534f"/>`;
    }
    if (recordBtnText) recordBtnText.textContent = "Stop & Download";
  }
}

function stopRecording(download = true) {
  if (!isRecording) return;
  isRecording = false;
  
  if (recordBtn) {
    recordBtn.classList.remove("btn-recording");
    if (recordIcon) recordIcon.classList.remove("pulse-record");
    if (recordIcon) {
      recordIcon.innerHTML = `<circle cx="12" cy="12" r="8" fill="#d9534f"/>`;
    }
    if (recordBtnText) recordBtnText.textContent = "Record";
  }
  
  if (download && recordedNotes.length > 0) {
    const bpm = Number(tempoEl?.textContent || 120);
    const fileBytes = compileMidiFile(recordedNotes, bpm);
    downloadMidi(fileBytes);
  }
  recordedNotes = [];
}

function encodeVLQ(value) {
  const bytes = [];
  let buffer = value & 0x7F;
  value = value >> 7;
  while (value > 0) {
    bytes.push((value & 0x7F) | 0x80);
    value = value >> 7;
  }
  bytes.reverse();
  bytes.push(buffer);
  return bytes;
}

function compileMidiFile(recordedNotes, bpm) {
  const division = 480; // Standard ticks per quarter note
  const trackBytes = [];
  
  // 1. Add tempo meta event (microsec per quarter note) at tick 0
  const tempoMicro = Math.round(60000000 / bpm);
  trackBytes.push(
    0, // delta-time = 0
    0xFF, 0x51, 0x03, // Meta Tempo, length = 3
    (tempoMicro >> 16) & 0xFF,
    (tempoMicro >> 8) & 0xFF,
    tempoMicro & 0xFF
  );
  
  // 2. Add note events to a flat list in absolute ticks
  const midiEvents = [];
  for (const note of recordedNotes) {
    const onTicks = Math.round(note.time * (bpm * division / 60));
    const offTicks = Math.round((note.time + note.duration) * (bpm * division / 60));
    
    midiEvents.push({
      ticks: onTicks,
      type: 0x90,
      note: note.midi,
      velocity: Math.round(note.velocity * 127)
    });
    
    midiEvents.push({
      ticks: offTicks,
      type: 0x80,
      note: note.midi,
      velocity: 0
    });
  }
  
  // 3. Sort events chronologically (with note-off before note-on for identical ticks)
  midiEvents.sort((a, b) => {
    if (a.ticks !== b.ticks) return a.ticks - b.ticks;
    return a.type - b.type;
  });
  
  // 4. Generate delta-times and append events
  let lastTicks = 0;
  for (const event of midiEvents) {
    const delta = event.ticks - lastTicks;
    lastTicks = event.ticks;
    
    trackBytes.push(...encodeVLQ(delta));
    trackBytes.push(event.type, event.note, event.velocity);
  }
  
  // 5. Add End of Track event
  trackBytes.push(0, 0xFF, 0x2F, 0x00);
  
  // 6. Build File Bytes
  const headerBytes = [
    0x4d, 0x54, 0x68, 0x64, // "MThd"
    0, 0, 0, 6,             // Length
    0, 0,                   // Format 0
    0, 1,                   // Number of tracks
    (division >> 8) & 0xFF,
    division & 0xFF
  ];
  
  const trackLen = trackBytes.length;
  const trackHeaderBytes = [
    0x4d, 0x54, 0x72, 0x6b, // "MTrk"
    (trackLen >> 24) & 0xFF,
    (trackLen >> 16) & 0xFF,
    (trackLen >> 8) & 0xFF,
    trackLen & 0xFF
  ];
  
  const fileBytes = new Uint8Array(headerBytes.length + trackHeaderBytes.length + trackBytes.length);
  fileBytes.set(headerBytes, 0);
  fileBytes.set(trackHeaderBytes, headerBytes.length);
  fileBytes.set(trackBytes, headerBytes.length + trackHeaderBytes.length);
  
  return fileBytes;
}

function downloadMidi(fileBytes, filename = "nanomaestro_performance.mid") {
  const blob = new Blob([fileBytes], { type: "audio/midi" });
  const url = URL.createObjectURL(blob);
  const a = document.createElement("a");
  a.href = url;
  a.download = filename;
  document.body.appendChild(a);
  a.click();
  document.body.removeChild(a);
  URL.revokeObjectURL(url);
}

async function loadCustomUserFiles() {
  stop();
  isModelReady = false;
  startBtn.disabled = true;
  if (!modelFileInput || !vocabFileInput) {
    setStatus("Select model_int8.onnx and matching vocab.json in Advanced settings", "idle");
    return;
  }
  const modelFile = modelFileInput.files?.[0];
  const vocabFile = vocabFileInput.files?.[0];
  if (!modelFile || !vocabFile) {
    activeModelName = "";
    setStatus("Select model_int8.onnx and matching vocab.json in Advanced settings", "idle");
    return;
  }
  activeModelName = modelFile.name;
  setStatus(`Reading ${vocabFile.name}...`, "loading");
  const vocab = JSON.parse(await vocabFile.text());
  
  setStatus(`Loading ${modelFile.name}...`, "loading");
  const modelBuffer = await modelFile.arrayBuffer();
  
  await new Promise((resolve, reject) => {
    modelInitResolve = resolve;
    modelInitReject = reject;
    worker.postMessage({
      action: "init",
      activeModelName: modelFile.name,
      vocab: vocab,
      modelBuffer: modelBuffer
    }, [modelBuffer]);
  });
  
  setStatus(`Ready: ${activeModelName}`, "idle");
  startBtn.disabled = false;
}

function resizeCanvas() {
  const rect = noteCanvas.getBoundingClientRect();
  const dpr = window.devicePixelRatio || 1;
  const width = Math.max(1, Math.floor(rect.width * dpr));
  const height = Math.max(1, Math.floor(rect.height * dpr));
  if (noteCanvas.width !== width || noteCanvas.height !== height) {
    noteCanvas.width = width;
    noteCanvas.height = height;
  }
  canvasCtx.setTransform(dpr, 0, 0, dpr, 0, 0);
  canvasWidth = rect.width;
  canvasHeight = rect.height;
}

// Maps MIDI note number to its key position on the horizontal keyboard
function getKeyPosition(midi, width) {
  const minMidi = 36;
  const maxMidi = 96;
  const numWhiteKeys = 36;
  const whiteKeyW = width / numWhiteKeys;
  const noteInOctave = (midi - minMidi) % 12;
  const octave = Math.floor((midi - minMidi) / 12);
  
  const whiteKeyOffsets = [0, null, 1, null, 2, 3, null, 4, null, 5, null, 6];
  const isBlack = whiteKeyOffsets[noteInOctave] === null;
  
  if (!isBlack) {
    const whiteIdx = octave * 7 + whiteKeyOffsets[noteInOctave];
    const x = whiteIdx * whiteKeyW;
    return { x, w: whiteKeyW, isBlack: false, center: x + whiteKeyW / 2 };
  } else {
    const blackKeyBorderOffsets = {
      1: 1, // C#
      3: 2, // D#
      6: 4, // F#
      8: 5, // G#
      10: 6 // A#
    };
    const borderIdx = octave * 7 + blackKeyBorderOffsets[noteInOctave];
    const borderX = borderIdx * whiteKeyW;
    const w = whiteKeyW * 0.62; // Black key width
    const x = borderX - w / 2;
    return { x, w, isBlack: true, center: borderX };
  }
}

function drawVisualizer() {
  const width = canvasWidth;
  const height = canvasHeight;
  
  // Interpolate smooth audio time from high-precision performance clock
  const perfNow = performance.now();
  const rawAudioTime = Tone.now();
  if (rawAudioTime !== lastAudioTime) {
    lastAudioTime = rawAudioTime;
    lastPerfTime = perfNow;
  }
  const smoothNow = lastAudioTime + (perfNow - lastPerfTime) / 1000;
  const visualNow = smoothNow - VISUAL_DELAY;
  
  const minMidi = 36;
  const maxMidi = 96;
  const secondsVisible = 4.5;
  const keyboardH = width < 640 ? 80 : 120;
  const playLineY = height - keyboardH;
  const speed = playLineY / secondsVisible;
  
  canvasCtx.clearRect(0, 0, width, height);
  
  // 1. Draw background lanes for black keys
  for (let m = minMidi; m <= maxMidi; m++) {
    const keyInfo = getKeyPosition(m, width);
    if (keyInfo.isBlack) {
      canvasCtx.fillStyle = "rgba(255, 255, 255, 0.015)";
      canvasCtx.fillRect(keyInfo.x, 0, keyInfo.w, playLineY);
    }
  }
  
  // 2. Draw horizontal scrolling beat and bar grid lines
  const currentBpm = Number(tempoEl?.textContent || 120);
  const secondsPerBeat = 60 / currentBpm;
  
  const startBeat = Math.floor((visualNow - 1.0) / secondsPerBeat);
  const endBeat = Math.ceil((visualNow + secondsVisible) / secondsPerBeat);
  
  canvasCtx.font = "8px ui-sans-serif, system-ui, sans-serif";
  for (let b = startBeat; b <= endBeat; b++) {
    if (b < 0) continue;
    const beatTime = b * secondsPerBeat;
    const y = playLineY - (beatTime - visualNow) * speed;
    if (y < 0 || y > playLineY) continue;
    
    const isBar = b % 4 === 0;
    if (isBar) {
      canvasCtx.strokeStyle = "rgba(255, 255, 255, 0.045)";
      canvasCtx.lineWidth = 1;
      
      // Draw Bar text indicator
      canvasCtx.fillStyle = "rgba(255, 255, 255, 0.2)";
      canvasCtx.fillText(`BAR ${b / 4 + 1}`, 10, y - 4);
    } else {
      canvasCtx.strokeStyle = "rgba(255, 255, 255, 0.015)";
      canvasCtx.lineWidth = 0.5;
    }
    
    canvasCtx.beginPath();
    canvasCtx.moveTo(0, y);
    canvasCtx.lineTo(width, y);
    canvasCtx.stroke();
  }
  
  // Track active pitches for the keyboard representation
  const activePitches = new Set();
  
  // 3. Draw note capsules falling down
  for (const item of visualNotes) {
    const end = item.start + item.duration;
    if (end < visualNow - 0.5 || item.start > visualNow + secondsVisible) continue;
    
    const active = item.start <= visualNow && end >= visualNow;
    const midiClamped = Math.max(minMidi, Math.min(maxMidi, item.midi));
    if (active) activePitches.add(midiClamped);
    
    const keyInfo = getKeyPosition(midiClamped, width);
    const noteW = keyInfo.isBlack ? keyInfo.w * 0.85 : keyInfo.w * 0.75;
    const noteX = keyInfo.center - noteW / 2;
    
    const bottomY = playLineY - (item.start - visualNow) * speed;
    const topY = playLineY - (end - visualNow) * speed;
    
    const drawTop = Math.max(-100, topY);
    const drawBottom = Math.max(-100, bottomY);
    const noteH = drawBottom - drawTop;
    
    if (noteH > 0) {
      if (active) {
        canvasCtx.fillStyle = "#ffffff";
        canvasCtx.shadowBlur = 8;
        canvasCtx.shadowColor = "rgba(255, 255, 255, 0.5)";
        canvasCtx.globalAlpha = 0.95;
      } else {
        canvasCtx.fillStyle = "#3d4853"; // Muted slate gray note
        canvasCtx.shadowBlur = 0;
        canvasCtx.globalAlpha = 0.55;
      }
      
      canvasCtx.beginPath();
      if (typeof canvasCtx.roundRect === "function") {
        canvasCtx.roundRect(noteX, drawTop, noteW, noteH, Math.min(noteW / 2, 4));
      } else {
        canvasCtx.rect(noteX, drawTop, noteW, noteH);
      }
      canvasCtx.fill();
    }
  }
  
  // Reset shadow settings
  canvasCtx.shadowBlur = 0;
  canvasCtx.globalAlpha = 1.0;
  
  // 4. Draw Piano Keyboard at the bottom
  // Draw white keys first
  const numWhiteKeys = 36;
  const whiteKeyW = width / numWhiteKeys;
  for (let i = 0; i < numWhiteKeys; i++) {
    const octave = Math.floor(i / 7);
    const offsetIdx = i % 7;
    const whiteToNote = [0, 2, 4, 5, 7, 9, 11];
    const midi = 36 + octave * 12 + whiteToNote[offsetIdx];
    
    const active = activePitches.has(midi);
    const x = i * whiteKeyW;
    
    canvasCtx.fillStyle = active ? "#ffffff" : "#1a1f26";
    
    canvasCtx.beginPath();
    if (typeof canvasCtx.roundRect === "function") {
      canvasCtx.roundRect(x + 1, playLineY, whiteKeyW - 2, keyboardH, [0, 0, 3, 3]);
    } else {
      canvasCtx.rect(x + 1, playLineY, whiteKeyW - 2, keyboardH);
    }
    canvasCtx.fill();
    
    // Key divider shadow
    canvasCtx.fillStyle = "rgba(0, 0, 0, 0.35)";
    canvasCtx.fillRect(x, playLineY, 1, keyboardH);
  }
  
  // Draw black keys on top
  const blackKeyHeight = keyboardH * 0.62;
  for (let m = minMidi; m <= maxMidi; m++) {
    const noteInOctave = (m - minMidi) % 12;
    const isBlack = [1, 3, 6, 8, 10].includes(noteInOctave);
    if (!isBlack) continue;
    
    const active = activePitches.has(m);
    const keyInfo = getKeyPosition(m, width);
    
    canvasCtx.fillStyle = active ? "#6f9c78" : "#090a0d";
    
    canvasCtx.beginPath();
    if (typeof canvasCtx.roundRect === "function") {
      canvasCtx.roundRect(keyInfo.x, playLineY, keyInfo.w, blackKeyHeight, [0, 0, 2, 2]);
    } else {
      canvasCtx.rect(keyInfo.x, playLineY, keyInfo.w, blackKeyHeight);
    }
    canvasCtx.fill();
    
    // 3D highlight
    canvasCtx.fillStyle = active ? "rgba(255, 255, 255, 0.15)" : "rgba(255, 255, 255, 0.05)";
    canvasCtx.fillRect(keyInfo.x, playLineY, keyInfo.w, 2);
  }
  
  // Play line divider
  canvasCtx.fillStyle = "rgba(255, 255, 255, 0.1)";
  canvasCtx.fillRect(0, playLineY - 1, width, 1);
  
  // Shadow below the play line onto the keyboard keys
  const playLineShadow = canvasCtx.createLinearGradient(0, playLineY, 0, playLineY + 6);
  playLineShadow.addColorStop(0, "rgba(0, 0, 0, 0.45)");
  playLineShadow.addColorStop(1, "rgba(0, 0, 0, 0)");
  canvasCtx.fillStyle = playLineShadow;
  canvasCtx.fillRect(0, playLineY, width, 6);
  
  // 5. Draw HUD label texts inside visualizer
  canvasCtx.fillStyle = "rgba(255, 255, 255, 0.15)";
  canvasCtx.font = "9px ui-sans-serif, system-ui, sans-serif";
  canvasCtx.fillText("BASS", 12, playLineY - 10);
  canvasCtx.fillText("TREBLE", width - 48, playLineY - 10);
  
  visualNotes = visualNotes.filter((item) => item.start + item.duration > visualNow - 1.5);
  if (running) visualFrame = requestAnimationFrame(drawVisualizer);
}

function startVisualizer() {
  if (visualFrame) cancelAnimationFrame(visualFrame);
  visualFrame = requestAnimationFrame(drawVisualizer);
}

// Dual-Layer Model Caching Engine (IndexedDB + Cache Storage API)
let cacheDBPromise = null;

function getCacheDB() {
  if (!cacheDBPromise) {
    cacheDBPromise = new Promise((resolve) => {
      if (!window.indexedDB) return resolve(null);
      const req = indexedDB.open("nanomaestro_cache_v2", 1);
      req.onupgradeneeded = (e) => {
        const db = e.target.result;
        if (!db.objectStoreNames.contains("models")) {
          db.createObjectStore("models");
        }
      };
      req.onsuccess = (e) => resolve(e.target.result);
      req.onerror = (e) => {
        console.warn("IndexedDB open error:", e);
        resolve(null);
      };
    });
  }
  return cacheDBPromise;
}

async function getCachedData(key) {
  try {
    const db = await getCacheDB();
    if (!db) return null;
    return new Promise((resolve) => {
      const tx = db.transaction("models", "readonly");
      const store = tx.objectStore("models");
      const req = store.get(key);
      req.onsuccess = () => resolve(req.result || null);
      req.onerror = () => resolve(null);
    });
  } catch (e) {
    return null;
  }
}

async function setCachedData(key, data) {
  try {
    const db = await getCacheDB();
    if (!db) return;
    const tx = db.transaction("models", "readwrite");
    const store = tx.objectStore("models");
    store.put(data, key);
  } catch (e) {
    // Fail silently if browser quota limit is reached
  }
}

async function getFromCacheStorage(cacheKey) {
  try {
    if (!("caches" in window)) return null;
    const cache = await caches.open("nanomaestro_models_v2");
    const res = await cache.match(cacheKey);
    if (res && res.ok) {
      return res;
    }
  } catch (e) {
    // Ignore cache storage errors
  }
  return null;
}

async function setInCacheStorage(cacheKey, response) {
  try {
    if (!("caches" in window)) return;
    const cache = await caches.open("nanomaestro_models_v2");
    await cache.put(cacheKey, response);
  } catch (e) {
    // Ignore cache storage errors
  }
}

// Stream reader helper with progress callback
async function fetchStreamWithProgress(response, onProgress) {
  const contentLength = response.headers.get("content-length");
  const total = contentLength ? parseInt(contentLength, 10) : 0;
  
  let loaded = 0;
  if (response.body && typeof response.body.getReader === "function") {
    const reader = response.body.getReader();
    const chunks = [];
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
      loaded += value.byteLength;
      if (onProgress) onProgress(loaded, total);
    }
    return new Blob(chunks);
  } else {
    const blob = await response.blob();
    if (onProgress) onProgress(blob.size, blob.size);
    return blob;
  }
}

// Dual-layer cache fetch helper with stable asset keys
async function fetchAssetWithCache(cacheKey, localUrl, remoteUrl, onProgress, isJson = false) {
  // 1. Check IndexedDB FIRST
  try {
    const cachedData = await getCachedData(cacheKey);
    if (cachedData) {
      console.log(`[Cache HIT - IDB] Loaded ${cacheKey}`);
      if (onProgress) {
        const size = isJson ? 1 : (cachedData.size || 1);
        onProgress(size, size);
      }
      return cachedData;
    }
  } catch (e) {
    console.warn("IndexedDB read error:", e);
  }

  // 2. Check Cache Storage SECOND
  try {
    const cachedResponse = await getFromCacheStorage(cacheKey);
    if (cachedResponse) {
      console.log(`[Cache HIT - Cache API] Loaded ${cacheKey}`);
      let data = isJson ? await cachedResponse.json() : await cachedResponse.blob();
      if (onProgress) {
        const size = isJson ? 1 : (data.size || 1);
        onProgress(size, size);
      }
      void setCachedData(cacheKey, data);
      return data;
    }
  } catch (e) {
    console.warn("Cache storage read error:", e);
  }

  // 3. Network Fetch: Try localUrl first, fallback to remoteUrl
  console.log(`[Cache MISS] Fetching ${cacheKey} from network...`);
  let data;
  let responseToCache = null;

  try {
    const response = await fetch(localUrl);
    if (!response.ok) throw new Error(`HTTP ${response.status} fetching ${localUrl}`);
    responseToCache = response.clone();
    if (isJson) {
      data = await response.json();
      if (onProgress) onProgress(1, 1);
    } else {
      data = await fetchStreamWithProgress(response, onProgress);
    }
  } catch (localErr) {
    console.warn(`Local fetch failed for ${localUrl}, fetching from Hugging Face: ${remoteUrl}`);
    const response = await fetch(remoteUrl);
    if (!response.ok) throw new Error(`HTTP ${response.status} fetching ${remoteUrl}`);
    responseToCache = response.clone();
    if (isJson) {
      data = await response.json();
      if (onProgress) onProgress(1, 1);
    } else {
      data = await fetchStreamWithProgress(response, onProgress);
    }
  }

  // 4. Save downloaded asset into BOTH caches!
  void setCachedData(cacheKey, data);
  if (responseToCache) {
    void setInCacheStorage(cacheKey, responseToCache);
  }

  return data;
}

const MODEL_PRESETS = {
  pro: {
    key: "pro",
    name: "NanoMaestro Pro (54M)",
    modelUrl: "./NanoMaestro-Pro/model_int8.onnx",
    remoteModelUrl: "https://huggingface.co/utkucoban/NanoMaestro-Realtime/resolve/main/NanoMaestro-Pro/model_int8.onnx",
    vocabUrl: "./NanoMaestro-Pro/vocab.json",
    remoteVocabUrl: "https://huggingface.co/utkucoban/NanoMaestro-Realtime/resolve/main/NanoMaestro-Pro/vocab.json",
    ep: "wasm",
    config: { numLayers: 2, hiddenSize: 2048 },
    desc: "NanoMaestro Pro is a 50 MB music-generation model with 54 million parameters, designed to run continuously in real time on almost any consumer CPU."
  },
  light: {
    key: "light",
    name: "NanoMaestro Light (13M)",
    modelUrl: "./NanoMaestro-Light/model_int8.onnx",
    remoteModelUrl: "https://huggingface.co/utkucoban/NanoMaestro-Realtime/resolve/main/NanoMaestro-Light/model_int8.onnx",
    vocabUrl: "./NanoMaestro-Light/vocab.json",
    remoteVocabUrl: "https://huggingface.co/utkucoban/NanoMaestro-Realtime/resolve/main/NanoMaestro-Light/vocab.json",
    ep: "wasm",
    config: { numLayers: 2, hiddenSize: 1024 },
    desc: "NanoMaestro Light is a tiny, 13 MB music-generation model with 13 million parameters, designed to run continuously in real time on almost any consumer CPU."
  }
};

let currentModelKey = "pro";

function syncModelSelectors(key) {
  const topSelect = document.querySelector("#topModelSelect");
  const settingsSelect = document.querySelector("#modelSelect");
  if (topSelect && topSelect.value !== key) topSelect.value = key;
  if (settingsSelect && settingsSelect.value !== key) settingsSelect.value = key;
}

async function loadSelectedModel(presetKey, updateProgress) {
  const preset = MODEL_PRESETS[presetKey] || MODEL_PRESETS.pro;
  currentModelKey = presetKey;
  isModelReady = false;
  syncModelSelectors(presetKey);
  
  const vocabCacheKey = `vocab_${preset.key}_v1`;
  const modelCacheKey = `model_${preset.key}_v1`;
  
  if (updateProgress) updateProgress(5, `Loading ${preset.name} vocabulary...`);
  const vocab = await fetchAssetWithCache(vocabCacheKey, preset.vocabUrl, preset.remoteVocabUrl, null, true);
  
  if (updateProgress) updateProgress(15, `Loading ${preset.name} weights...`);
  const modelBlob = await fetchAssetWithCache(modelCacheKey, preset.modelUrl, preset.remoteModelUrl, (loaded, total) => {
    if (updateProgress) {
      const loadedMB = (loaded / (1024 * 1024)).toFixed(1);
      if (total > 0) {
        const pct = 15 + (loaded / total) * 65;
        const totalMB = (total / (1024 * 1024)).toFixed(1);
        updateProgress(pct, `Downloading ${preset.name} (${loadedMB}MB / ${totalMB}MB)...`);
      } else {
        const estPct = Math.min(78, 15 + (loaded / (54 * 1024 * 1024)) * 63);
        updateProgress(estPct, `Loading ${preset.name} (${loadedMB}MB loaded)...`);
      }
    }
  }, false);
  
  if (updateProgress) updateProgress(82, `Initializing ${preset.name} engine (CPU)...`);
  
  const modelBuffer = await modelBlob.arrayBuffer();
  await new Promise((resolve, reject) => {
    let timer = setTimeout(() => {
      reject(new Error("Neural network initialization timed out (25s)"));
    }, 25000);

    modelInitResolve = () => {
      clearTimeout(timer);
      resolve();
    };
    modelInitReject = (err) => {
      clearTimeout(timer);
      reject(err);
    };

    worker.postMessage({
      action: "init",
      activeModelName: preset.name,
      vocab: vocab,
      modelBuffer: modelBuffer,
      ep: preset.ep,
      modelConfig: preset.config
    }, [modelBuffer]);
  });
  
  activeModelName = preset.name;
  isModelReady = true;
  
  const modelDescEl = document.querySelector("#modelDesc");
  if (modelDescEl) {
    modelDescEl.textContent = preset.desc;
  }
}

// Initializer Flow
async function init() {
  const loadingOverlay = document.querySelector("#loadingOverlay");
  const loadProgress = document.querySelector("#loadProgress");
  const loadStatus = document.querySelector("#loadStatus");
  const loadPercent = document.querySelector("#loadPercent");
  
  const updateProgress = (pct, text) => {
    if (loadProgress) loadProgress.style.width = `${pct}%`;
    if (loadPercent) loadPercent.textContent = `${Math.round(pct)}%`;
    if (loadStatus) loadStatus.textContent = text;
  };
  
  updateProgress(2, "Initializing NanoMaestro engine...");
  
  try {
    updateReadouts();
    startBtn.disabled = true;
    
    const topSelect = document.querySelector("#topModelSelect");
    const settingsSelect = document.querySelector("#modelSelect");
    const selectedKey = (topSelect ? topSelect.value : null) || (settingsSelect ? settingsSelect.value : "pro");
    
    await loadSelectedModel(selectedKey, updateProgress);
    
    // Step 4. Loading Audio Samples
    updateProgress(90, `Loading ${currentInstrument} audio samples...`);
    await ensureInstrument();
    
    updateProgress(100, "Ready!");
    
    // Remove loading overlay and set status
    setTimeout(() => {
      if (loadingOverlay) {
        loadingOverlay.classList.add("fade-out");
      }
      setStatus(`Ready: ${activeModelName}`, "idle");
      startBtn.disabled = false;
      
      // Draw static visualizer background once
      resizeCanvas();
      drawVisualizer();
    }, 600);
    
  } catch (error) {
    console.error("Initialization failed:", error);
    if (loadStatus) {
      loadStatus.textContent = `Auto-load failed: ${error.message}`;
      loadStatus.style.color = "#ff8787";
    }
  }
}

// Bind Buttons and Input Controls
startBtn.disabled = true;
startBtn.addEventListener("click", () => void start());
stopBtn.addEventListener("click", stop);
if (newSongBtn) {
  newSongBtn.addEventListener("click", async () => {
    if (!running) return;
    newSongBtn.disabled = true;
    stop();
    setTimeout(async () => {
      try {
        await start();
      } catch (error) {
        console.error("New song restart failed:", error);
      }
    }, 150);
  });
}

if (recordBtn) {
  recordBtn.addEventListener("click", () => {
    if (isRecording) {
      stopRecording(true);
    } else {
      startRecording();
    }
  });
}

if (modelFileInput) {
  modelFileInput.addEventListener("change", async () => {
    try {
      await loadCustomUserFiles();
    } catch (error) {
      console.error(error);
      setStatus(`Load failed: ${error.message}`, "idle");
      startBtn.disabled = true;
    }
  });
}

if (vocabFileInput) {
  vocabFileInput.addEventListener("change", async () => {
    try {
      await loadCustomUserFiles();
    } catch (error) {
      console.error(error);
      setStatus(`Load failed: ${error.message}`, "idle");
      startBtn.disabled = true;
    }
  });
}

// Dynamic loading models logic removed since we load directly from directory.

tempInput.addEventListener("input", () => updateReadouts());
topKInput.addEventListener("input", () => updateReadouts());

volumeInput.addEventListener("input", (e) => {
  setVolume(Number(e.target.value));
});

reverbInput.addEventListener("input", (e) => {
  const wet = Number(e.target.value);
  if (reverb) reverb.wet.value = wet / 100;
  if (reverbValue) reverbValue.textContent = `${wet}%`;
});

delayInput.addEventListener("input", (e) => {
  const wet = Number(e.target.value);
  if (delay) delay.wet.value = wet / 100;
  if (delayValue) delayValue.textContent = `${wet}%`;
});

toggleButtons.forEach(btn => {
  btn.addEventListener("click", () => {
    if (btn.disabled) return;
    
    toggleButtons.forEach(b => b.classList.remove("active"));
    btn.classList.add("active");
    
    const value = btn.dataset.value;
    activePlayingStyle = value;
    if (value === "midi") {
      midiSeedControls.style.display = "block";
    } else {
      midiSeedControls.style.display = "none";
    }
  });
});

midiBpmModeButtons.forEach(btn => {
  btn.addEventListener("click", () => {
    if (btn.disabled) return;
    
    midiBpmModeButtons.forEach(b => b.classList.remove("active"));
    btn.classList.add("active");
    
    midiBpmMode = btn.dataset.value;
  });
});

midiFileInput.addEventListener("change", async (e) => {
  const file = e.target.files?.[0];
  if (!file) return;
  
  midiFileName = file.name;
  midiFileNameEl.textContent = `File: ${file.name}`;
  midiFileInfo.style.display = "block";
  midiFileStatsEl.textContent = "Parsing MIDI file...";
  
  try {
    const buffer = await file.arrayBuffer();
    const parser = new MidiParser(buffer);
    const parsed = parser.parse();
    
    // Tokenize
    const { tokens, maxBar } = tokenizeMidi(parsed, 64);
    midiTokens = tokens;
    midiBpm = parsed.bpm;
    midiTotalBars = Math.max(1, maxBar + 1);
    
    midiFileStatsEl.textContent = `BPM: ${parsed.bpm} | Total Bars: ${midiTotalBars} | Notes: ${parsed.notes.length}`;
    
    // Configure start bar slider
    const minBar = findMinStartBar(tokens);
    midiStartBarInput.min = String(minBar);
    midiStartBarInput.max = String(midiTotalBars - 1);
    midiStartBarInput.value = String(minBar);
    midiStartBar = minBar;
    midiStartBarValue.textContent = String(minBar);
    
    // Analyze and render timeline
    const barsData = analyzeMidiTokens(tokens);
    renderMidiTimeline(barsData);
    updateMidiStartBarText();
  } catch (error) {
    console.error(error);
    midiFileStatsEl.textContent = `Error: ${error.message}`;
  }
});

midiStartBarInput.addEventListener("input", (e) => {
  const val = Number(e.target.value);
  midiStartBar = val;
  midiStartBarValue.textContent = String(val);
  updateMidiTimelineHighlights();
  updateMidiStartBarText();
});

releaseInput.addEventListener("input", (e) => {
  const rel = Number(e.target.value);
  setRelease(rel);
});

const instrumentSelect = document.querySelector("#instrumentSelect");
if (instrumentSelect) {
  instrumentSelect.addEventListener("change", async (e) => {
    currentInstrument = e.target.value;
    if (synth) {
      const isAlreadyRunning = running;
      if (isAlreadyRunning) {
        // Temporarily show loading status during swap
        setStatus(`Loading ${currentInstrument} audio samples...`, "loading");
      }
      await createSampler(currentInstrument);
      if (isAlreadyRunning) {
        setStatus("Performing", "playing");
      }
    }
  });
}

const handleModelChange = async (e) => {
  const newKey = e.target.value;
  if (running) {
    stop();
  }
  startBtn.disabled = true;
  setStatus(`Switching to ${MODEL_PRESETS[newKey]?.name || newKey}...`, "loading");
  try {
    await loadSelectedModel(newKey, (pct, msg) => setStatus(msg, "loading"));
    setStatus(`Ready: ${activeModelName}`, "idle");
    startBtn.disabled = false;
  } catch (err) {
    console.error("Model switch failed:", err);
    setStatus(`Model switch failed: ${err.message}`, "idle");
  }
};

const topModelSelect = document.querySelector("#topModelSelect");
if (topModelSelect) {
  topModelSelect.addEventListener("change", handleModelChange);
}

const modelSelect = document.querySelector("#modelSelect");
if (modelSelect) {
  modelSelect.addEventListener("change", handleModelChange);
}

window.addEventListener("resize", resizeCanvas);

// Trigger setup on start
void init();