agent-ui/frontend/src/services/whisperSocket.ts

/**
 * Singleton Whisper WebSocket Service
 * One shared connection used by all voice components (FloatingVoice, useVoiceInput, etc.)
 */

import { ref } from 'vue'
import { resolveEndpoints } from '../config/endpoints'
import { apiFetch } from '@/lib/tauri'

export type WhisperStatus = 'offline' | 'loading' | 'ready'

type TranscriptionCallback = (msg: {
  success?: boolean
  text?: string
  error?: string
  partial?: boolean
  model?: string
  device?: string
}) => void

// ====== Singleton state ======
const status = ref<WhisperStatus>('loading')
let socket: WebSocket | null = null
let reconnectTimer: number | null = null
const listeners = new Set<TranscriptionCallback>()

// ====== Connection management ======

function connect() {
  if (socket?.readyState === WebSocket.OPEN || socket?.readyState === WebSocket.CONNECTING) return

  console.log('[WhisperSocket] Connecting to', resolveEndpoints().whisper)
  socket = new WebSocket(resolveEndpoints().whisper)

  const timeout = setTimeout(() => {
    if (socket && socket.readyState !== WebSocket.OPEN) {
      console.error('[WhisperSocket] Connection timeout (10s)')
      socket.close()
      status.value = 'loading'
    }
  }, 10000)

  socket.onopen = () => {
    clearTimeout(timeout)
    console.log('[WhisperSocket] Connected')
    status.value = 'ready'
  }

  socket.onmessage = (event) => {
    try {
      const msg = JSON.parse(event.data)

      if (msg.type === 'ready') {
        console.log('[WhisperSocket] Server ready:', msg.model, msg.device)
        status.value = 'ready'
      } else if (msg.type === 'transcription') {
        // Broadcast to all listeners
        for (const cb of listeners) {
          cb(msg)
        }
      }
    } catch (e) {
      console.error('[WhisperSocket] Message parse error:', e)
    }
  }

  socket.onclose = () => {
    console.log('[WhisperSocket] Closed, will reconnect...')
    socket = null
    status.value = 'loading'
    scheduleReconnect()
  }

  socket.onerror = (e) => {
    console.error('[WhisperSocket] Error:', e)
    status.value = 'loading'
  }
}

function scheduleReconnect() {
  if (reconnectTimer) return
  reconnectTimer = window.setTimeout(() => {
    reconnectTimer = null
    checkStatusAndConnect()
  }, 2000)
}

async function checkStatusAndConnect() {
  try {
    const res = await apiFetch('/api/whisper/status')
    const data = await res.json()
    if (data.running) {
      connect()
    } else {
      status.value = 'loading'
      scheduleReconnect()
    }
  } catch {
    status.value = 'loading'
    scheduleReconnect()
  }
}

// ====== Public API ======

/** Initialize the singleton connection (call once at app startup) */
export function initWhisperSocket() {
  checkStatusAndConnect()
}

/** Send audio for transcription */
export function sendAudio(base64: string, language: string, partial: boolean) {
  if (socket?.readyState === WebSocket.OPEN) {
    socket.send(JSON.stringify({
      type: 'transcribe',
      audio: base64,
      language,
      partial
    }))
  } else {
    console.warn('[WhisperSocket] Not connected, dropping audio')
  }
}

/** Subscribe to transcription results. Returns unsubscribe function. */
export function onTranscription(callback: TranscriptionCallback): () => void {
  listeners.add(callback)
  return () => listeners.delete(callback)
}

/** Get reactive status */
export function getWhisperStatus() {
  return status
}

/** Check if socket is connected */
export function isConnected(): boolean {
  return socket?.readyState === WebSocket.OPEN
}

/** Force reconnect (e.g. when user toggles Whisper) */
export async function reconnect() {
  if (status.value === 'loading' && socket?.readyState === WebSocket.CONNECTING) return

  status.value = 'loading'
  if (socket) {
    socket.close()
    socket = null
  }

  try {
    const res = await apiFetch('/api/whisper/toggle', { method: 'POST' })
    const data = await res.json()
    if (data.running) {
      connect()
    } else {
      // Poll until ready
      const poll = async () => {
        for (let i = 0; i < 60; i++) {
          await new Promise(r => setTimeout(r, 2000))
          try {
            const s = await apiFetch('/api/whisper/status')
            const d = await s.json()
            if (d.running) {
              connect()
              return
            }
          } catch { /* retry */ }
        }
        status.value = 'offline'
      }
      poll()
    }
  } catch {
    status.value = 'loading'
    scheduleReconnect()
  }
}