FabledScribe/frontend/src/composables/useVad.ts

import { ref, readonly } from 'vue'
import type { MicVAD } from '@ricky0123/vad-web'

export interface UseVadOptions {
  onSpeechEnd: () => void
  onNoSpeech: () => void
  onVadError?: (message: string) => void
  graceMs?: number
  minRecordingMs?: number
}

/**
 * Silero VAD-based speech detector.
 *
 * Replaces amplitude-based silence detection. Uses the Silero VAD v5 ONNX
 * model via `@ricky0123/vad-web`. The package ships the model + audio
 * worklet, and depends on `onnxruntime-web`'s WASM. Vite is configured
 * (see `vite.config.ts`) to serve these assets from the site root, which
 * is why we pass `baseAssetPath: "/"` and `onnxWASMBasePath: "/"`.
 *
 * `amplitude` is computed separately from a Web Audio API analyser node
 * on the same stream — it drives the mic pulse animation and is not part
 * of the stop decision. VAD's `onSpeechEnd` is the stop trigger.
 *
 * If VAD never detects speech during a session, calling `stopAndCheck()`
 * fires `onNoSpeech` instead of treating the recording as transcribable.
 */
export function useVad(options: UseVadOptions) {
  const {
    onSpeechEnd,
    onNoSpeech,
    graceMs = 1500,
    minRecordingMs = 500,
  } = options

  const speaking = ref(false)
  const amplitude = ref(0)
  const loaded = ref(false)

  let micVad: MicVAD | null = null
  let audioCtx: AudioContext | null = null
  let analyserInterval: ReturnType<typeof setInterval> | null = null
  let speechDetected = false
  let speechStartedAt = 0
  let recordingStartedAt = 0
  let stopped = false

  async function start(stream: MediaStream): Promise<void> {
    await stop()
    stopped = false
    speechDetected = false
    speechStartedAt = 0
    recordingStartedAt = Date.now()

    // Visual-only amplitude signal. This does NOT drive the stop decision.
    audioCtx = new AudioContext()
    await audioCtx.resume().catch(() => {})
    const source = audioCtx.createMediaStreamSource(stream)
    const analyser = audioCtx.createAnalyser()
    analyser.fftSize = 1024
    source.connect(analyser)
    const samples = new Float32Array(analyser.fftSize)

    analyserInterval = setInterval(() => {
      analyser.getFloatTimeDomainData(samples)
      let sumSq = 0
      for (let i = 0; i < samples.length; i++) sumSq += samples[i] * samples[i]
      const rms = Math.sqrt(sumSq / samples.length)
      amplitude.value = Math.min(rms * 5, 1)
    }, 100)

    try {
      const { MicVAD } = await import('@ricky0123/vad-web')
      micVad = await MicVAD.new({
        model: 'v5',
        baseAssetPath: '/',
        onnxWASMBasePath: '/',
        // MicVAD manages its own audio graph, so we hand it the same stream
        // the MediaRecorder is using. It won't consume or alter the stream.
        getStream: async () => stream,
        onSpeechStart: () => {
          speaking.value = true
          if (!speechDetected) {
            speechDetected = true
            speechStartedAt = Date.now()
          }
        },
        onSpeechEnd: () => {
          speaking.value = false
          if (stopped) return
          const now = Date.now()
          const totalElapsed = now - recordingStartedAt
          const sinceSpeechStart = speechStartedAt > 0 ? now - speechStartedAt : 0
          if (totalElapsed >= minRecordingMs && sinceSpeechStart >= graceMs) {
            stopped = true
            onSpeechEnd()
          }
        },
        onVADMisfire: () => {
          // Speech-like blip that was too short to count. Reset the local flag
          // so a true "no speech" session still hits onNoSpeech.
          speaking.value = false
        },
      })
      await micVad.start()
      loaded.value = true
    } catch (e) {
      console.error('VAD init failed:', e)
      options.onVadError?.(e instanceof Error ? e.message : String(e))
    }
  }

  async function stop(): Promise<{ hadSpeech: boolean }> {
    const had = speechDetected
    stopped = true
    speaking.value = false
    amplitude.value = 0

    if (analyserInterval !== null) {
      clearInterval(analyserInterval)
      analyserInterval = null
    }
    if (audioCtx) {
      await audioCtx.close().catch(() => {})
      audioCtx = null
    }
    if (micVad) {
      try {
        await micVad.pause()
        await micVad.destroy()
      } catch {
        // Already destroyed or failed — nothing to clean up further.
      }
      micVad = null
    }
    return { hadSpeech: had }
  }

  async function stopAndCheck(): Promise<void> {
    const { hadSpeech } = await stop()
    if (!hadSpeech) {
      onNoSpeech()
    }
  }

  return {
    speaking: readonly(speaking),
    amplitude: readonly(amplitude),
    loaded: readonly(loaded),
    start,
    stop,
    stopAndCheck,
  }
}