import { comfyWsUrl } from '../comfy'
import { saveRecord } from './store'

/** Listen only to this client's prompt; history polling remains the completion authority. */
export function watchProgress(record: any, job: any) {
  let ws: WebSocket | undefined, timer: ReturnType<typeof setTimeout>, stopped = false
  const connect = () => {
    if (stopped) return
    try {
      ws = new WebSocket(comfyWsUrl(job.clientId))
      ws.addEventListener('message', event => {
        if (typeof event.data !== 'string') return
        try {
          const {type,data} = JSON.parse(event.data)
          if (!record.promptId || (data?.prompt_id && data.prompt_id !== record.promptId)) return
          if (type === 'progress' && Number(data.max) > 0) {
            record.progress = {step:Number(data.value),maxStep:Number(data.max),percent:Math.round(Number(data.value)/Number(data.max)*100),updatedAt:Date.now()}
            saveRecord(record)
          }
        } catch { /* Binary previews and malformed events cannot interrupt a job. */ }
      })
      ws.addEventListener('error', () => { ws?.close() })
      ws.addEventListener('close', () => { if (!stopped) timer=setTimeout(connect,3000) })
    } catch { if (!stopped) timer=setTimeout(connect,3000) }
  }
  connect()
  return () => { stopped=true;clearTimeout(timer);ws?.close() }
}