// voicenote — fully client-side voice dictation.
// Mic capture (AudioWorklet, 16 kHz) → energy-based VAD segments utterances →
// a speech model in a Web Worker transcribes them (interim results while you
// speak, final text appended on silence). Document lives in localStorage.

import { MODELS, DEFAULT_MODEL, cachedModels, detectDevice } from './models.js';

const SR = 16000;
const DOC_KEY = 'voicenote-doc';
const MODEL_KEY = 'voicenote-model';

// VAD tuning (all sample counts at 16 kHz)
const PREROLL_SAMPLES = 0.32 * SR;      // audio kept before speech onset
const ENTER_FRAMES = 2;                 // consecutive loud frames to trigger speech
// Long enough that a dramatic mid-sentence pause does not split the sentence:
// short fragments transcribe badly, because the model loses the surrounding
// words it needs for context.
const SILENCE_FINAL = 0.9 * SR;         // trailing silence that ends an utterance
const SOFT_MAX = 20 * SR;               // after this, finalize at the next brief dip
const SOFT_DIP = 0.15 * SR;
const HARD_MAX = 28 * SR;               // hard cut (whisper's window is 30 s)
const MIN_SPEECH = 0.35 * SR;           // discard blips shorter than this
// Generous, because a trailing fricative ("...Americans") is quiet enough to
// read as silence, and trimming into it costs the end of the word.
const KEEP_TAIL = 0.4 * SR;             // trailing silence kept when trimming
const INTERIM_MIN = 0.8 * SR;           // utterance length before the first interim
const INTERIM_EVERY = 1.2 * SR;         // new audio between interim runs

// These models emit stock filler on near-silent audio; drop it when it is the
// entire result rather than letting it accumulate in the document.
const JUNK = new Set(['you', 'thank you.', 'thanks for watching!', 'thank you for watching!',
  '[blank_audio]', '(silence)', '.', 'bye.', 'the', 'so']);

const $ = (id) => document.getElementById(id);
const recordBtn = $('recordBtn');
const recordLabel = $('recordLabel');
const copyBtn = $('copyBtn');
const clearBtn = $('clearBtn');
const statusEl = $('status');
const modelSel = $('model');
const modelNote = $('modelNote');
const progressWrap = $('progressWrap');
const progressBar = $('progressBar');
const progressText = $('progressText');
const doc = $('doc');
const interimEl = $('interim');
const wordcountEl = $('wordcount');
const engineEl = $('engine');

// ---------- document ----------

doc.value = localStorage.getItem(DOC_KEY) || '';

let saveTimer = null;
function saveDoc() {
  clearTimeout(saveTimer);
  saveTimer = setTimeout(() => localStorage.setItem(DOC_KEY, doc.value), 250);
}

function updateWordCount() {
  const n = doc.value.trim().split(/\s+/).filter(Boolean).length;
  wordcountEl.textContent = `${n} word${n === 1 ? '' : 's'}`;
}

function autoresize() {
  doc.style.height = 'auto';
  doc.style.height = doc.scrollHeight + 'px';
}

doc.addEventListener('input', () => {
  saveDoc();
  updateWordCount();
  autoresize();
});

function appendFinal(text) {
  text = (text || '').trim();
  if (!text || JUNK.has(text.toLowerCase())) return;
  const sep = doc.value && !/\s$/.test(doc.value) ? ' ' : '';
  const focused = document.activeElement === doc;
  const selStart = doc.selectionStart;
  const selEnd = doc.selectionEnd;
  doc.value += sep + text;
  if (focused) doc.setSelectionRange(selStart, selEnd);
  saveDoc();
  updateWordCount();
  autoresize();
  if (!focused) window.scrollTo({ top: document.body.scrollHeight });
}

copyBtn.addEventListener('click', async () => {
  try {
    await navigator.clipboard.writeText(doc.value);
    copyBtn.textContent = 'Copied ✓';
    setTimeout(() => (copyBtn.textContent = 'Copy'), 1200);
  } catch {
    setStatus('clipboard access denied');
  }
});

clearBtn.addEventListener('click', () => {
  if (!doc.value.trim() || confirm('Clear the document? This cannot be undone.')) {
    doc.value = '';
    interimEl.textContent = '';
    localStorage.setItem(DOC_KEY, '');
    updateWordCount();
    autoresize();
    doc.focus();
  }
});

// ---------- model picker ----------

let modelId = localStorage.getItem(MODEL_KEY) || DEFAULT_MODEL;
if (!MODELS.some((m) => m.id === modelId)) modelId = DEFAULT_MODEL;

for (const m of MODELS) {
  const opt = document.createElement('option');
  opt.value = m.id;
  opt.textContent = `${m.label} — ${m.mb} MB`;
  modelSel.appendChild(opt);
}
modelSel.value = modelId;

function showModelNote() {
  modelNote.textContent = MODELS.find((m) => m.id === modelId)?.note ?? '';
}

// Label each option with the download it would actually cost on the backend
// we expect to use, and mark the ones already sitting in the browser cache.
async function refreshCacheLabels() {
  const cached = await cachedModels(device);
  for (const opt of modelSel.options) {
    const m = MODELS.find((x) => x.id === opt.value);
    opt.textContent = `${m.label} — ${cached.has(m.id) ? 'downloaded' : `${m.mb[device]} MB`}`;
  }
}

modelSel.addEventListener('change', () => {
  modelId = modelSel.value;
  localStorage.setItem(MODEL_KEY, modelId);
  showModelNote();

  // Drop the loaded model; anything queued for it is abandoned with the worker.
  if (worker) {
    worker.terminate();
    worker = null;
  }
  modelState = 'unloaded';
  outstanding = 0;
  interimEl.textContent = '';
  interimUtt = -1;
  engineEl.textContent = '';
  progressWrap.classList.add('hidden');

  if (recording) ensureModel(); // keep listening; load the new model right away
  updateStatus();
});

// ---------- transcription worker ----------

let worker = null;
let modelState = 'unloaded'; // unloaded | loading | ready | error
let outstanding = 0;         // transcription requests in flight
let nextReqId = 1;
let device = 'wasm';         // backend we expect to use; confirmed on 'ready'
let forceCpu = false;        // set once WebGPU has proven it cannot run here

// Finished utterances captured before a model was ready. They live here rather
// than in the worker so that restarting the worker — which is how the WebGPU
// fallback works — never costs the user a phrase they already spoke.
const queued = [];
let queuedSamples = 0;
const MAX_QUEUED = 120 * SR; // ~2 minutes of speech; drop the oldest beyond it

function fmtMB(b) { return (b / 1024 / 1024).toFixed(0); }

// Creates the worker and starts the model load. The worker is constructed
// synchronously and the load message queued first, so nothing races it.
function ensureModel() {
  if (modelState === 'ready' || modelState === 'loading') return;
  modelState = 'loading';
  worker = new Worker('worker.js', { type: 'module' });
  worker.onmessage = onWorkerMessage;
  worker.onerror = (e) => failLoad(e.message || 'worker failed to start');
  worker.postMessage({ type: 'load', model: modelId, device: forceCpu ? 'cpu' : 'auto' });
}

// A WebGPU failure needs a brand-new worker: ONNX Runtime will not hand back a
// working CPU backend in a process where WebGPU already failed to start.
function retryOnCpu(message) {
  console.warn('voicenote: WebGPU unavailable, restarting on CPU —', message);
  forceCpu = true;
  device = 'wasm';
  if (worker) { worker.terminate(); worker = null; }
  modelState = 'unloaded';
  outstanding = 0;
  progressWrap.classList.add('hidden');
  refreshCacheLabels();
  ensureModel();
  setStatus('WebGPU unavailable — loading on CPU instead…');
}

function failLoad(message) {
  modelState = 'error';
  progressWrap.classList.add('hidden');
  stopMic();
  syncButton();
  setStatus(`error: ${message}`);
}

function onWorkerMessage(e) {
  const msg = e.data;
  if (msg.type === 'progress') {
    progressWrap.classList.remove('hidden');
    progressBar.style.width = `${(100 * msg.loaded) / msg.total}%`;
    progressText.textContent = `${fmtMB(msg.loaded)} / ${fmtMB(msg.total)} MB`;
  } else if (msg.type === 'status') {
    setStatus(msg.message);
  } else if (msg.type === 'ready') {
    progressWrap.classList.add('hidden');
    modelState = 'ready';
    device = msg.device;
    engineEl.textContent =
      `${MODELS.find((m) => m.id === modelId).label} · ${msg.device === 'webgpu' ? 'WebGPU' : 'CPU (wasm)'}`;
    refreshCacheLabels();
    flushQueued();
    updateStatus();
  } else if (msg.type === 'result') {
    outstanding = Math.max(0, outstanding - 1);
    if (msg.final) {
      if (msg.utt === interimUtt) {
        interimEl.textContent = '';
        interimUtt = -1;
      }
      appendFinal(msg.text);
    } else if (msg.utt === currentUtt && inSpeech) {
      const t = (msg.text || '').trim();
      if (t && !JUNK.has(t.toLowerCase())) {
        interimEl.textContent = t + ' …';
        interimUtt = msg.utt;
      }
    }
    updateStatus();
  } else if (msg.type === 'error') {
    if (msg.retryOnCpu && !forceCpu) retryOnCpu(msg.message);
    else if (msg.fatal || modelState === 'loading') failLoad(msg.message);
    else setStatus(`transcription error: ${msg.message}`);
  }
}

function postAudio(samples, final, uttId) {
  outstanding++;
  const buf = samples.buffer;
  worker.postMessage({ type: 'transcribe', id: nextReqId++, utt: uttId, final, audio: buf }, [buf]);
}

function sendAudio(samples, final, uttId) {
  if (modelState === 'error') return;
  if (modelState === 'ready' && worker) {
    postAudio(samples, final, uttId);
    return;
  }
  if (!final) return; // interims are disposable; only hold on to finished phrases
  queued.push({ samples, uttId });
  queuedSamples += samples.length;
  while (queuedSamples > MAX_QUEUED && queued.length > 1) {
    queuedSamples -= queued.shift().samples.length;
  }
}

function flushQueued() {
  while (queued.length) {
    const q = queued.shift();
    queuedSamples -= q.samples.length;
    postAudio(q.samples, true, q.uttId);
  }
  queuedSamples = 0;
}

// ---------- mic capture + VAD ----------

let recording = false;
let stream = null;
let audioCtx = null;

let inSpeech = false;
let enterCount = 0;
let noiseFloor = 0.002;
let preroll = [];        // frames seen before speech onset
let prerollSamples = 0;
let utt = [];            // frames of the current utterance
let uttRms = [];         // rms per frame, for the trailing-silence trim
let uttSamples = 0;
let silenceRun = 0;      // samples of trailing sub-threshold audio
let sinceInterim = 0;
let currentUtt = 0;      // id of the open utterance
let interimUtt = -1;     // utterance whose interim text is on screen

function concatFrames(frames, totalLen) {
  const out = new Float32Array(totalLen);
  let o = 0;
  for (const f of frames) { out.set(f, o); o += f.length; }
  return out;
}

function resample(frame, fromRate) {
  if (fromRate === SR) return frame;
  const ratio = fromRate / SR;
  const n = Math.floor(frame.length / ratio);
  const out = new Float32Array(n);
  for (let i = 0; i < n; i++) {
    const pos = i * ratio;
    const i0 = Math.floor(pos);
    const frac = pos - i0;
    const i1 = Math.min(i0 + 1, frame.length - 1);
    out[i] = frame[i0] * (1 - frac) + frame[i1] * frac;
  }
  return out;
}

function onFrame(frame) {
  let sum = 0;
  for (let i = 0; i < frame.length; i++) sum += frame[i] * frame[i];
  const rms = Math.sqrt(sum / frame.length);

  const enterTh = Math.max(noiseFloor * 3, 0.006);
  const exitTh = Math.max(noiseFloor * 2, 0.004);

  if (!inSpeech) {
    // adapt the noise floor from quiet frames only
    if (rms < enterTh) noiseFloor = Math.max(0.0001, 0.95 * noiseFloor + 0.05 * rms);
    preroll.push(frame);
    prerollSamples += frame.length;
    while (prerollSamples > PREROLL_SAMPLES && preroll.length > 1) {
      prerollSamples -= preroll[0].length;
      preroll.shift();
    }
    enterCount = rms > enterTh ? enterCount + 1 : 0;
    if (enterCount >= ENTER_FRAMES) {
      inSpeech = true;
      enterCount = 0;
      currentUtt++;
      utt = preroll.slice();
      uttRms = preroll.map(() => rms); // approximate; only the tail is used
      uttSamples = prerollSamples;
      preroll = [];
      prerollSamples = 0;
      silenceRun = 0;
      sinceInterim = 0;
      updateStatus();
    }
    return;
  }

  utt.push(frame);
  uttRms.push(rms);
  uttSamples += frame.length;
  sinceInterim += frame.length;
  silenceRun = rms < exitTh ? silenceRun + frame.length : 0;

  const done =
    silenceRun >= SILENCE_FINAL ||
    (uttSamples >= SOFT_MAX && silenceRun >= SOFT_DIP) ||
    uttSamples >= HARD_MAX;

  if (done) {
    finalizeUtterance();
  } else if (
    modelState === 'ready' &&
    outstanding === 0 &&
    uttSamples >= INTERIM_MIN &&
    sinceInterim >= INTERIM_EVERY
  ) {
    sinceInterim = 0;
    sendAudio(concatFrames(utt, uttSamples), false, currentUtt);
  }
}

function finalizeUtterance() {
  if (!inSpeech) return;
  inSpeech = false;

  // trim trailing silence down to KEEP_TAIL
  let trailing = 0;
  const exitTh = Math.max(noiseFloor * 2, 0.004);
  for (let i = uttRms.length - 1; i >= 0 && uttRms[i] < exitTh; i--) {
    trailing += utt[i].length;
  }
  while (utt.length > 1 && trailing - utt[utt.length - 1].length >= KEEP_TAIL) {
    const f = utt.pop();
    uttRms.pop();
    uttSamples -= f.length;
    trailing -= f.length;
  }

  const voiced = uttSamples - Math.min(trailing, uttSamples);
  if (voiced >= MIN_SPEECH) {
    sendAudio(concatFrames(utt, uttSamples), true, currentUtt);
  } else if (interimUtt === currentUtt) {
    interimEl.textContent = '';
    interimUtt = -1;
  }
  utt = [];
  uttRms = [];
  uttSamples = 0;
  silenceRun = 0;
  updateStatus();
}

async function startMic() {
  stream = await navigator.mediaDevices.getUserMedia({
    audio: { channelCount: 1, echoCancellation: true, noiseSuppression: true, autoGainControl: true },
  });
  audioCtx = new AudioContext({ sampleRate: SR });
  if (audioCtx.state === 'suspended') await audioCtx.resume();
  await audioCtx.audioWorklet.addModule('worklet.js');
  const src = audioCtx.createMediaStreamSource(stream);
  const node = new AudioWorkletNode(audioCtx, 'capture');
  const rate = audioCtx.sampleRate;
  node.port.onmessage = (e) => {
    if (recording) onFrame(resample(e.data, rate));
  };
  src.connect(node);
  node.connect(audioCtx.destination); // worklet outputs silence; keeps the graph pulling
  recording = true;
}

function stopMic() {
  recording = false;
  if (inSpeech) finalizeUtterance();
  if (stream) { stream.getTracks().forEach((t) => t.stop()); stream = null; }
  if (audioCtx) { audioCtx.close(); audioCtx = null; }
  inSpeech = false;
  enterCount = 0;
  preroll = [];
  prerollSamples = 0;
}

// ---------- UI state ----------

function setStatus(text) { statusEl.textContent = text; }

function syncButton() {
  recordBtn.classList.toggle('live', recording);
  recordBtn.classList.toggle('speech', recording && inSpeech);
  recordLabel.textContent = recording ? 'Pause' : 'Record';
}

function updateStatus() {
  syncButton();
  if (modelState === 'error') return; // keep the error message on screen
  if (modelState === 'loading') {
    const n = queued.length;
    setStatus(
      'loading speech model (one-time download, then cached)…' +
        (n ? ` — ${n} phrase${n === 1 ? '' : 's'} waiting` : ''),
    );
  } else if (recording) {
    setStatus(inSpeech ? 'hearing speech…' : outstanding > 0 ? 'transcribing…' : 'listening…');
  } else if (outstanding > 0) {
    setStatus('transcribing…');
  } else if (modelState === 'ready') {
    setStatus('paused — press Record to continue');
  } else {
    setStatus('');
  }
}

recordBtn.addEventListener('click', async () => {
  if (recording) {
    stopMic();
    updateStatus();
    return;
  }
  recordBtn.disabled = true;
  try {
    // Open the mic first, while still close to the user gesture, then kick off
    // the model load in parallel; audio captured meanwhile is queued.
    await startMic();
    ensureModel();
    updateStatus();
  } catch (err) {
    console.error(err);
    stopMic();
    syncButton();
    setStatus(
      err.name === 'NotAllowedError'
        ? 'microphone access denied — allow it and try again'
        : `error: ${err.message}`,
    );
  } finally {
    recordBtn.disabled = false;
  }
});

window.addEventListener('beforeunload', () => {
  clearTimeout(saveTimer);
  localStorage.setItem(DOC_KEY, doc.value);
});

updateWordCount();
autoresize();
showModelNote();
setStatus('');

// Sizes depend on the backend, so settle that before labelling the picker.
detectDevice().then((d) => {
  device = d;
  refreshCacheLabels();
});