// Browser-side ASR client. Connects to /asr (proxied to mochi-asr on i7:9000),
// captures the mic, downsamples in an AudioWorklet, emits transcript events.

import { asrWsUrl } from './api'

export type AsrEvent =
  | { kind: 'partial'; text: string }
  | { kind: 'final'; text: string }
  | { kind: 'error'; message: string }

export type AsrEventHandler = (e: AsrEvent) => void

export class AsrSession {
  private ws: WebSocket | null = null
  private audioCtx: AudioContext | null = null
  private node: AudioWorkletNode | null = null
  private mediaStream: MediaStream | null = null
  private source: MediaStreamAudioSourceNode | null = null
  private handler: AsrEventHandler
  private active = false

  constructor(handler: AsrEventHandler) {
    this.handler = handler
  }

  async start(): Promise<void> {
    if (this.active) return
    this.active = true

    this.mediaStream = await navigator.mediaDevices.getUserMedia({
      audio: { channelCount: 1, echoCancellation: true, noiseSuppression: true },
    })

    this.audioCtx = new AudioContext()
    await this.audioCtx.audioWorklet.addModule('/asr-worklet.js')
    this.node = new AudioWorkletNode(this.audioCtx, 'asr-worklet', {
      numberOfInputs: 1,
      numberOfOutputs: 0,
      channelCount: 1,
    })
    this.source = this.audioCtx.createMediaStreamSource(this.mediaStream)
    this.source.connect(this.node)

    this.ws = new WebSocket(asrWsUrl())
    this.ws.binaryType = 'arraybuffer'

    const flushQueue: ArrayBuffer[] = []
    let wsReady = false

    this.node.port.onmessage = (ev: MessageEvent<ArrayBuffer>) => {
      if (!this.ws) return
      if (wsReady && this.ws.readyState === WebSocket.OPEN) {
        this.ws.send(ev.data)
      } else {
        flushQueue.push(ev.data)
      }
    }

    this.ws.addEventListener('open', () => {
      wsReady = true
      while (flushQueue.length) {
        const buf = flushQueue.shift()!
        this.ws?.send(buf)
      }
    })

    this.ws.addEventListener('message', (ev) => {
      if (typeof ev.data !== 'string') return
      let data: { text?: string; final?: boolean; error?: string }
      try {
        data = JSON.parse(ev.data)
      } catch {
        return
      }
      if (data.error) {
        this.handler({ kind: 'error', message: data.error })
        return
      }
      if (typeof data.text === 'string') {
        this.handler(
          data.final
            ? { kind: 'final', text: data.text }
            : { kind: 'partial', text: data.text },
        )
      }
    })

    this.ws.addEventListener('error', () => {
      this.handler({ kind: 'error', message: 'asr ws error' })
    })
  }

  async stop(): Promise<void> {
    if (!this.active) return
    this.active = false

    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify({ type: 'flush' }))
    }

    if (this.source) this.source.disconnect()
    if (this.node) this.node.disconnect()
    if (this.mediaStream) {
      for (const track of this.mediaStream.getTracks()) track.stop()
    }
    if (this.audioCtx) await this.audioCtx.close()

    this.source = null
    this.node = null
    this.mediaStream = null
    this.audioCtx = null

    // Give upstream a moment to send the {"final": true} response.
    const ws = this.ws
    this.ws = null
    setTimeout(() => ws?.close(), 1500)
  }
}
