app.js

/**
 * app.js
 * ------
 * Browser-side voice agent client.
 *
 * Flow
 * ----
 *  getUserMedia  ──►  AudioWorklet capture (16 kHz int16)
 *                          │
 *                    WebSocket.send(binary)
 *                          │
 *                    ◄──── Server ─────►
 *                          │
 *              WebSocket.onmessage(binary)  ← PCM int16 TTS audio
 *                          │
 *                    playback queue  ──►  AudioContext.playBuffer()
 *
 * States:  idle → listening → processing → speaking → idle
 */

'use strict';

// ── Config ──────────────────────────────────────────────────────────────── //
const SERVER_SR   = 16_000;   // Server sends 16 kHz PCM
const WS_URL      = (() => {
  const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
  return `${proto}//${location.host}/ws/audio`;
})();

// ── State ────────────────────────────────────────────────────────────────── //
let audioCtx        = null;
let captureNode     = null;
let micStream       = null;
let ws              = null;
let appState        = 'idle';   // idle | listening | processing | speaking

let audioQueue      = [];       // Array<AudioBuffer> pending playback
let isPlaying       = false;    // Is the drain loop running?
let conversation    = [];       // [{role, text}] for UI transcript

// ── DOM refs ─────────────────────────────────────────────────────────────── //
const btnToggle     = document.getElementById('btn-toggle');
const statusDot     = document.getElementById('status-dot');
const statusText    = document.getElementById('status-text');
const canvasEl      = document.getElementById('visualiser');
const ctx2d         = canvasEl.getContext('2d');
const transcriptEl  = document.getElementById('transcript');
const latencyEl     = document.getElementById('latency');

// ── Entry point ───────────────────────────────────────────────────────────── //
btnToggle.addEventListener('click', () => {
  if (appState === 'idle') {
    startSession();
  } else {
    stopSession();
  }
});

// ── Session start/stop ───────────────────────────────────────────────────── //
async function startSession () {
  try {
    await initAudio();
    await initWebSocket();
    setState('listening');
    btnToggle.textContent = '⏹ Stop';
    ws.send(JSON.stringify({ type: 'start' }));
    startVisualiser();
  } catch (err) {
    console.error('startSession:', err);
    setStatus('error', `Mic / WebSocket error: ${err.message}`);
  }
}

function stopSession () {
  if (ws) {
    ws.send(JSON.stringify({ type: 'stop' }));
    ws.close();
  }
  teardownAudio();
  setState('idle');
  btnToggle.textContent = '🎤 Start';
  stopVisualiser();
}

// ── Audio setup ───────────────────────────────────────────────────────────── //
async function initAudio () {
  // Request mic with constraints that help speech quality
  micStream = await navigator.mediaDevices.getUserMedia({
    audio: {
      channelCount:    1,
      echoCancellation: true,
      noiseSuppression: true,
      autoGainControl:  true,
    },
  });

  audioCtx = new AudioContext();

  // Load the capture worklet
  await audioCtx.audioWorklet.addModule('/audio-capture-worklet.js');

  const source  = audioCtx.createMediaStreamSource(micStream);
  captureNode   = new AudioWorkletNode(audioCtx, 'audio-capture-processor');

  captureNode.port.onmessage = ({ data }) => {
    if (data.type === 'audio' && ws && ws.readyState === WebSocket.OPEN) {
      ws.send(data.pcm);           // ArrayBuffer of int16 PCM
    }
  };

  source.connect(captureNode);
  // captureNode intentionally NOT connected to destination (no feedback)
  captureNode.port.postMessage({ type: 'start' });
}

function teardownAudio () {
  if (captureNode) {
    captureNode.port.postMessage({ type: 'stop' });
    captureNode.disconnect();
    captureNode = null;
  }
  if (micStream) {
    micStream.getTracks().forEach(t => t.stop());
    micStream = null;
  }
  if (audioCtx) {
    audioCtx.close();
    audioCtx = null;
  }
  audioQueue  = [];
  isPlaying   = false;
}

// ── WebSocket ─────────────────────────────────────────────────────────────── //
function initWebSocket () {
  return new Promise((resolve, reject) => {
    ws = new WebSocket(WS_URL);
    ws.binaryType = 'arraybuffer';

    ws.onopen    = () => { console.log('WS open'); resolve(); };
    ws.onerror   = (e) => reject(new Error('WebSocket error'));
    ws.onclose   = () => {
      if (appState !== 'idle') stopSession();
    };

    ws.onmessage = (evt) => {
      if (evt.data instanceof ArrayBuffer) {
        onAudioChunk(evt.data);       // Binary: PCM audio from TTS
      } else {
        onControlMessage(JSON.parse(evt.data));   // Text: JSON status
      }
    };
  });
}

// ── Inbound control messages ─────────────────────────────────────────────── //
function onControlMessage (msg) {
  switch (msg.type) {

    case 'status':
      setStatus(appState, msg.message);
      break;

    case 'transcript':
      appendTranscript('user', msg.text);
      setState('processing');
      break;

    case 'agent_start':
      setState('speaking');
      break;

    case 'agent_done':
      appendTranscript('agent', msg.text);
      if (msg.latency_ms) {
        latencyEl.textContent = `Latency: ${msg.latency_ms} ms`;
      }
      // State returns to 'listening' when the playback queue drains
      break;

    case 'interrupted':
      audioQueue  = [];
      isPlaying   = false;
      setState('listening');
      break;

    case 'error':
      console.error('Server error:', msg.message);
      setStatus('error', msg.message);
      setState('listening');
      break;
  }
}

// ── Audio playback queue ─────────────────────────────────────────────────── //
function onAudioChunk (arrayBuffer) {
  if (!audioCtx) return;

  // PCM int16, mono, 16 kHz  →  float32 AudioBuffer at AudioContext.sampleRate
  const int16        = new Int16Array(arrayBuffer);
  const float32      = new Float32Array(int16.length);
  for (let i = 0; i < int16.length; i++) float32[i] = int16[i] / 32_768;

  const resampled    = resample(float32, SERVER_SR, audioCtx.sampleRate);
  const audioBuf     = audioCtx.createBuffer(1, resampled.length, audioCtx.sampleRate);
  audioBuf.copyToChannel(resampled, 0);

  audioQueue.push(audioBuf);
  if (!isPlaying) drainQueue();
}

function drainQueue () {
  if (audioQueue.length === 0) {
    isPlaying = false;
    if (appState === 'speaking') setState('listening');
    return;
  }

  isPlaying       = true;
  const buf       = audioQueue.shift();
  const src       = audioCtx.createBufferSource();
  src.buffer      = buf;
  src.connect(audioCtx.destination);
  src.onended     = drainQueue;   // chain next buffer immediately
  src.start();
}

// ── Resampler (linear interpolation) ─────────────────────────────────────── //
function resample (input, inSR, outSR) {
  if (inSR === outSR) return input;
  const ratio     = inSR / outSR;
  const outLen    = Math.round(input.length / ratio);
  const output    = new Float32Array(outLen);
  for (let i = 0; i < outLen; i++) {
    const src     = i * ratio;
    const lo      = Math.floor(src);
    const hi      = Math.min(lo + 1, input.length - 1);
    const frac    = src - lo;
    output[i]     = input[lo] + frac * (input[hi] - input[lo]);
  }
  return output;
}

// ── State management ──────────────────────────────────────────────────────── //
const STATE_LABELS = {
  idle:       'Ready',
  listening:  'Listening…',
  processing: 'Thinking…',
  speaking:   'Speaking…',
};

const STATE_COLORS = {
  idle:       '#6b7280',
  listening:  '#22c55e',
  processing: '#f59e0b',
  speaking:   '#3b82f6',
};

function setState (newState) {
  appState = newState;
  statusDot.style.background = STATE_COLORS[newState] ?? '#6b7280';
  statusText.textContent     = STATE_LABELS[newState] ?? newState;
}

function setStatus (_state, message) {
  statusText.textContent = message;
}

// ── Transcript ────────────────────────────────────────────────────────────── //
function appendTranscript (role, text) {
  conversation.push({ role, text });
  const div   = document.createElement('div');
  div.className = `msg msg-${role}`;
  div.innerHTML = `<span class="msg-role">${role === 'user' ? 'You' : 'Agent'}</span><span class="msg-text">${escHtml(text)}</span>`;
  transcriptEl.appendChild(div);
  transcriptEl.scrollTop = transcriptEl.scrollHeight;
}

function escHtml (s) {
  return s.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;');
}

// ── Canvas visualiser ─────────────────────────────────────────────────────── //
let analyser    = null;
let visRAF      = null;

function startVisualiser () {
  if (!audioCtx || analyser) return;

  analyser             = audioCtx.createAnalyser();
  analyser.fftSize     = 256;
  const source         = audioCtx.createMediaStreamSource(micStream);
  source.connect(analyser);

  const dataArr        = new Uint8Array(analyser.frequencyBinCount);
  const W              = canvasEl.width;
  const H              = canvasEl.height;
  const BAR_COUNT      = 48;
  const BAR_W          = W / BAR_COUNT - 1;

  function draw () {
    visRAF = requestAnimationFrame(draw);
    analyser.getByteFrequencyData(dataArr);

    ctx2d.clearRect(0, 0, W, H);

    const color = STATE_COLORS[appState] ?? '#6b7280';
    ctx2d.fillStyle = color + '99';   // 60% opacity

    for (let i = 0; i < BAR_COUNT; i++) {
      const binIdx = Math.floor(i * dataArr.length / BAR_COUNT);
      const v      = dataArr[binIdx] / 255;
      const barH   = Math.max(2, v * H * 0.9);
      const x      = i * (BAR_W + 1);
      const y      = (H - barH) / 2;
      ctx2d.beginPath();
      ctx2d.roundRect(x, y, BAR_W, barH, 2);
      ctx2d.fill();
    }
  }
  draw();
}

function stopVisualiser () {
  if (visRAF) { cancelAnimationFrame(visRAF); visRAF = null; }
  if (analyser) { analyser.disconnect(); analyser = null; }
  ctx2d.clearRect(0, 0, canvasEl.width, canvasEl.height);
}

// ── Resize canvas ─────────────────────────────────────────────────────────── //
function resizeCanvas () {
  canvasEl.width  = canvasEl.offsetWidth;
  canvasEl.height = canvasEl.offsetHeight;
}
window.addEventListener('resize', resizeCanvas);
resizeCanvas();