srt-streamer/electron/ffmpeg.js

'use strict'
const { EventEmitter } = require('events')
const { spawn } = require('child_process')
const path = require('path')
const fs = require('fs')

class FFmpegManager extends EventEmitter {
  constructor() {
    super()
    this.streams = new Map()        // id -> { process, config }
    this._manualStops = new Set()   // IDs stopped by user (no reconnect)
    this._reconnectTimers = new Map() // id -> setTimeout handle
    this.tolbekProcess = null
    this.ffmpegPath = this._resolveFfmpegPath()
  }

  _resolveFfmpegPath() {
    const resourcesPath = process.resourcesPath || path.join(__dirname, '..')
    const candidates = [
      path.join(resourcesPath, 'ffmpeg-bin', process.platform === 'win32' ? 'ffmpeg.exe' : 'ffmpeg'),
      path.join(__dirname, '..', 'ffmpeg-bin', process.platform === 'win32' ? 'ffmpeg.exe' : 'ffmpeg'),
    ]
    for (const p of candidates) {
      if (fs.existsSync(p)) { console.log('[FFmpeg] Using bundled:', p); return p }
    }
    console.log('[FFmpeg] Using system ffmpeg')
    return process.platform === 'win32' ? 'ffmpeg.exe' : 'ffmpeg'
  }

  // ─── Build FFmpeg args ────────────────────────────────────────────────────────
  _buildStreamArgs(config) {
    const args = []
    const platform = process.platform
    const { videoSource, audioSource } = config

    // ── VIDEO INPUT ──────────────────────────────────────────────────────────────
    if (videoSource.type === 'desktop') {
      if (platform === 'win32') {
        args.push('-f', 'gdigrab')
        args.push('-draw_mouse', videoSource.hideCursor ? '0' : '1')
        args.push('-framerate', String(config.framerate || 30))

        // Per-monitor capture: use stored bounds if available
        if (videoSource.monitorBounds) {
          const { x, y, width, height } = videoSource.monitorBounds
          args.push('-offset_x', String(x))
          args.push('-offset_y', String(y))
          args.push('-video_size', `${width}x${height}`)
        }
        args.push('-i', 'desktop')
      } else if (platform === 'darwin') {
        args.push('-f', 'avfoundation')
        args.push('-framerate', String(config.framerate || 30))
        args.push('-capture_cursor', videoSource.hideCursor ? '0' : '1')
        args.push('-i', `${videoSource.screenIndex ?? '1'}:none`)
      } else {
        args.push('-f', 'x11grab')
        const display = process.env.DISPLAY || ':0'
        args.push('-framerate', String(config.framerate || 30))
        if (videoSource.hideCursor) args.push('-draw_mouse', '0')
        // Per-monitor: use bounds if available
        if (videoSource.monitorBounds) {
          const { x, y, width, height } = videoSource.monitorBounds
          args.push('-s', `${width}x${height}`)
          args.push('-i', `${display}+${x},${y}`)
        } else {
          args.push('-i', display)
        }
      }
    } else if (videoSource.type === 'window') {
      if (platform === 'win32') {
        args.push('-f', 'gdigrab')
        args.push('-draw_mouse', videoSource.hideCursor ? '0' : '1')
        args.push('-framerate', String(config.framerate || 30))
        args.push('-i', `title=${videoSource.windowTitle || videoSource.deviceName}`)
      } else if (platform === 'darwin') {
        args.push('-f', 'avfoundation')
        args.push('-framerate', String(config.framerate || 30))
        args.push('-i', `${videoSource.deviceIndex || '1'}:none`)
      } else {
        args.push('-f', 'x11grab')
        args.push('-framerate', String(config.framerate || 30))
        args.push('-i', process.env.DISPLAY || ':0')
      }
    } else if (videoSource.type === 'device') {
      if (platform === 'win32') {
        args.push('-f', 'dshow')
        args.push('-framerate', String(config.framerate || 30))
        if (config.resolution) args.push('-video_size', config.resolution)
        args.push('-i', `video=${videoSource.deviceName}`)
      } else if (platform === 'darwin') {
        args.push('-f', 'avfoundation')
        args.push('-framerate', String(config.framerate || 30))
        if (config.resolution) args.push('-video_size', config.resolution)
        args.push('-i', `${videoSource.deviceIndex}:none`)
      } else {
        args.push('-f', 'v4l2')
        args.push('-framerate', String(config.framerate || 30))
        if (config.resolution) args.push('-video_size', config.resolution)
        args.push('-i', videoSource.devicePath || '/dev/video0')
      }
    } else {
      // No video — black frame
      args.push('-f', 'lavfi', '-i', 'color=black:640x480:rate=25')
    }

    // ── AUDIO INPUT ──────────────────────────────────────────────────────────────
    if (audioSource && audioSource.type !== 'none') {
      if (platform === 'win32') {
        if (audioSource.deviceName === '__wasapi_loopback__') {
          // Native WASAPI loopback — captures all system audio, no extra software needed
          args.push('-f', 'wasapi', '-loopback', '1', '-i', '')
        } else {
          args.push('-f', 'dshow')
          args.push('-i', `audio=${audioSource.deviceName}`)
        }
      } else if (platform === 'darwin') {
        args.push('-f', 'avfoundation')
        args.push('-i', `:${audioSource.deviceIndex || '0'}`)
      } else {
        if (audioSource.type === 'system') {
          args.push('-f', 'pulse')
          args.push('-i', audioSource.deviceName || 'default.monitor')
        } else {
          args.push('-f', 'alsa')
          args.push('-i', audioSource.deviceName || 'default')
        }
      }
    }

    // ── VIDEO ENCODING ────────────────────────────────────────────────────────────
    const videoCodec = this._selectVideoCodec(config.hwAccel)
    args.push('-c:v', videoCodec)

    if (videoCodec === 'h264_nvenc') {
      args.push('-preset', 'p4', '-tune', 'ull', '-rc', 'cbr')
    } else if (videoCodec === 'h264_qsv') {
      args.push('-preset', 'faster')
    } else if (videoCodec === 'h264_amf') {
      args.push('-quality', 'speed', '-rc', 'cbr')
    } else if (videoCodec === 'h264_videotoolbox') {
      args.push('-realtime', '1')
    } else {
      args.push('-preset', 'ultrafast', '-tune', 'zerolatency')
    }

    const videoBitrate = config.videoBitrate || '2000k'
    args.push('-b:v', videoBitrate, '-maxrate', videoBitrate, '-bufsize', videoBitrate)

    if (config.resolution) {
      args.push('-vf', `scale=${config.resolution.replace('x', ':')}`)
    }
    args.push('-g', String((config.framerate || 30) * 2))

    // ── AUDIO ENCODING ────────────────────────────────────────────────────────────
    if (audioSource && audioSource.type !== 'none') {
      args.push('-c:a', 'aac', '-b:a', config.audioBitrate || '128k', '-ar', '44100')
    } else {
      args.push('-an')
    }

    // ── OUTPUT (SRT / tee with recording) ─────────────────────────────────────────
    const { serverAddress, port, srtMode, latency } = config
    const address = serverAddress.replace(/^(srt:\/\/|rtmp:\/\/|http:\/\/|https:\/\/)/, '')
    const lat  = latency || 200
    const mode = srtMode || 'caller'
    const srtUrl = `srt://${address}:${port}?mode=${mode}&latency=${lat * 1000}&pkt_size=1316`

    if (config.record && config.recordFolder) {
      // Tee muxer: stream to SRT AND write segmented file simultaneously
      const segSecs = Math.max(1, config.recordSegmentMinutes || 60) * 60
      const safeName = (config.name || config.id).replace(/[^a-zA-Z0-9_-]/g, '_')

      // Build file path — convert Windows backslashes, escape drive colon for tee muxer
      let folder = (config.recordFolder + '').replace(/\\/g, '/')
      if (process.platform === 'win32') {
        folder = folder.replace(/^([A-Za-z]):/, '$1\\:')
      }
      const filePath = `${folder}/${safeName}_%Y%m%d_%H%M%S.ts`
      const fileOpts = `f=segment:segment_time=${segSecs}:strftime=1:reset_timestamps=1`
      const teeOut = `[f=mpegts]${srtUrl}|[${fileOpts}]${filePath}`

      args.push('-f', 'tee', teeOut)
    } else {
      args.push('-f', 'mpegts', srtUrl)
    }

    return args
  }

  _selectVideoCodec(hwAccel) {
    if (!hwAccel || hwAccel === 'none') return 'libx264'
    if (hwAccel === 'auto') {
      if (process.platform === 'darwin') return 'h264_videotoolbox'
      return 'h264_nvenc'  // will fallback to software if unavailable
    }
    return {
      nvenc:        'h264_nvenc',
      qsv:          'h264_qsv',
      amf:          'h264_amf',
      videotoolbox: 'h264_videotoolbox',
      vaapi:        'h264_vaapi',
      software:     'libx264',
    }[hwAccel] || 'libx264'
  }

  // ─── Public: start stream ────────────────────────────────────────────────────
  async startStream(config) {
    const { id } = config

    // Cancel any pending reconnect timer
    if (this._reconnectTimers.has(id)) {
      clearTimeout(this._reconnectTimers.get(id))
      this._reconnectTimers.delete(id)
    }

    // Kill existing process for this id if any
    if (this.streams.has(id)) {
      const existing = this.streams.get(id)
      try { existing.process.kill() } catch {}
      this.streams.delete(id)
    }

    // Mark as NOT a manual stop (so reconnect will fire on unexpected exits)
    this._manualStops.delete(id)

    return this._spawn(config, 0, true)
  }

  // ─── Internal: spawn ffmpeg process ─────────────────────────────────────────
  _spawn(config, attempt, initial) {
    const { id } = config

    return new Promise((resolve, reject) => {
      const args = this._buildStreamArgs(config)
      console.log(`[FFmpeg] Starting stream ${id} (attempt ${attempt}):`, this.ffmpegPath, args.join(' '))

      const proc = spawn(this.ffmpegPath, args, { stdio: ['pipe', 'pipe', 'pipe'] })

      let started = false
      let errorBuffer = ''

      proc.stderr.on('data', (data) => {
        const text = data.toString()
        errorBuffer += text
        this.emit('log', { id, text })

        if (!started && (text.includes('Output #0') || text.includes('frame='))) {
          started = true
          if (initial) resolve({ pid: proc.pid })
        }
      })

      proc.on('error', (err) => {
        console.error(`[FFmpeg] Spawn error ${id}:`, err)
        if (initial && !started) reject(err)
        this.emit('error', { id, error: err.message })
        this.streams.delete(id)
      })

      proc.on('exit', (code, signal) => {
        console.log(`[FFmpeg] Stream ${id} exited: code=${code} signal=${signal} attempt=${attempt}`)
        const wasManual = this._manualStops.has(id)
        this.streams.delete(id)
        this.emit('ended', { id, code, signal })

        // ── Initial start failed — try hwAccel fallback first ─────────────────
        if (initial && !started && code !== 0 && config.hwAccel !== 'software') {
          console.log(`[FFmpeg] HW encode failed for ${id}, retrying with software…`)
          this._spawn({ ...config, hwAccel: 'software' }, 0, true)
            .then(resolve).catch(reject)
          return
        }
        if (initial && !started) {
          reject(new Error(`FFmpeg exited with code ${code}\n${errorBuffer}`))
          return
        }

        // ── Auto-reconnect on unexpected exit ─────────────────────────────────
        if (!wasManual) {
          const delayMs = Math.min(2000 * Math.pow(1.5, attempt), 30000)
          const nextAttempt = attempt + 1
          console.log(`[FFmpeg] Scheduling reconnect for ${id} in ${delayMs}ms (attempt ${nextAttempt})`)
          this.emit('reconnecting', { id, attempt: nextAttempt, delayMs })

          const timer = setTimeout(() => {
            this._reconnectTimers.delete(id)
            if (!this._manualStops.has(id)) {
              this._spawn(config, nextAttempt, false)
            }
          }, delayMs)
          this._reconnectTimers.set(id, timer)
        }
      })

      // Startup timeout — resolve even if FFmpeg hasn't printed "Output #0" yet
      setTimeout(() => {
        if (!started) {
          started = true
          if (initial) resolve({ pid: proc.pid })
        }
      }, 5000)

      this.streams.set(id, { process: proc, config })
    })
  }

  // ─── Public: stop stream ─────────────────────────────────────────────────────
  stopStream(id) {
    // Cancel any pending reconnect
    if (this._reconnectTimers.has(id)) {
      clearTimeout(this._reconnectTimers.get(id))
      this._reconnectTimers.delete(id)
    }

    // Mark as manual stop so reconnect won't fire
    this._manualStops.add(id)

    const stream = this.streams.get(id)
    if (stream) {
      console.log(`[FFmpeg] Stopping stream ${id}`)
      stream.process.kill(process.platform === 'win32' ? 'SIGTERM' : 'SIGINT')
      this.streams.delete(id)
    }
  }

  stopAllStreams() {
    for (const [id] of this.streams) this.stopStream(id)
    this.stopTolbek()
  }

  getActiveCount() {
    return this.streams.size + (this.tolbekProcess ? 1 : 0)
  }

  getActiveStreams() {
    const result = []
    for (const [id, { config }] of this.streams) result.push({ id, config })
    return result
  }

  // ─── Tolbek SRT Receiver ─────────────────────────────────────────────────────
  async startTolbek(config) {
    if (this.tolbekProcess) this.stopTolbek()

    const { port, latency, mode, serverAddress } = config
    const lat = (latency || 200) * 1000
    const srtMode = mode || 'listener'

    let srtUrl
    if (srtMode === 'listener') {
      srtUrl = `srt://0.0.0.0:${port}?mode=listener&latency=${lat}`
    } else if (srtMode === 'caller') {
      const addr = (serverAddress || '').replace(/^srt:\/\//, '')
      srtUrl = `srt://${addr}:${port}?mode=caller&latency=${lat}`
    } else {
      const addr = (serverAddress || '0.0.0.0').replace(/^srt:\/\//, '')
      srtUrl = `srt://${addr}:${port}?mode=rendezvous&latency=${lat}`
    }

    const args = [
      '-fflags', 'nobuffer', '-flags', 'low_delay',
      '-f', 'mpegts', '-i', srtUrl,
      '-f', 'null', '-'
    ]

    console.log('[FFmpeg] Starting Tolbek receiver:', args.join(' '))

    return new Promise((resolve, reject) => {
      const proc = spawn(this.ffmpegPath, args, { stdio: ['pipe', 'pipe', 'pipe'] })
      let started = false

      proc.stderr.on('data', (data) => {
        const text = data.toString()
        this.emit('log', { id: 'tolbek', text })
        if (!started) { started = true; resolve({ pid: proc.pid }) }
      })
      proc.on('error', (err) => {
        if (!started) reject(err)
        this.emit('error', { id: 'tolbek', error: err.message })
        this.tolbekProcess = null
      })
      proc.on('exit', (code) => {
        this.tolbekProcess = null
        this.emit('ended', { id: 'tolbek', code })
      })

      setTimeout(() => { if (!started) { started = true; resolve({ pid: proc.pid }) } }, 3000)

      this.tolbekProcess = proc
    })
  }

  stopTolbek() {
    if (this.tolbekProcess) {
      console.log('[FFmpeg] Stopping Tolbek receiver')
      this.tolbekProcess.kill(process.platform === 'win32' ? 'SIGTERM' : 'SIGINT')
      this.tolbekProcess = null
    }
  }
}

module.exports = new FFmpegManager()