teidesu-scripts/utils/mp4-streaming.ts

import { assert, ConditionVariable } from '@fuman/utils'
import { load } from 'cheerio'
import { ffetch } from './fetch.ts'
import { writeWebStreamToFile } from './fs.ts'

interface SimpleMpd {
  codecs: string
  segments: string[]
}

export function parseSimpleMpd(xml: string): SimpleMpd {
  const $ = load(xml, { xml: true })

  const period = $('Period')
  assert(period.length === 1, 'expected exactly one period')

  const adaptations = period.find('AdaptationSet')
  assert(adaptations.length === 1, 'expected exactly one adaptation set')

  const representation = adaptations.find('Representation')
  assert(representation.length === 1, 'expected exactly one representation')

  const segmentTemplate = representation.find('SegmentTemplate')
  assert(segmentTemplate.length === 1, 'expected exactly one segment template')

  const initUrl = segmentTemplate.attr('initialization')
  const templateUrl = segmentTemplate.attr('media')
  const startNum = segmentTemplate.attr('startNumber')

  assert(initUrl !== undefined, 'expected initialization url')
  assert(templateUrl !== undefined, 'expected template url')
  assert(!templateUrl.match(/\$(RepresentationID|Bandwidth|Time)\$/), 'unsupported template url')
  assert(startNum !== undefined, 'expected start number')

  const timeline = segmentTemplate.find('SegmentTimeline')
  assert(timeline.length === 1, 'expected exactly one segment timeline')

  const segments = timeline.find('S')
  assert(segments.length > 0, 'expected at least one segment')

  const segmentUrls: string[] = [initUrl]

  let segmentNum = Number(startNum)
  for (const segment of segments) {
    const duration = $(segment).attr('d')
    assert(duration !== undefined, 'expected duration')
    const r = $(segment).attr('r')
    const repeats = r ? Number.parseInt(r) + 1 : 1

    for (let i = 0; i < repeats; i++) {
      segmentUrls.push(templateUrl.replace('$Number$', String(segmentNum)))
      segmentNum++
    }
  }

  return {
    codecs: representation.attr('codecs')!,
    segments: segmentUrls,
  }
}

export function parseSimpleHls(m3u8: string): string[] {
  let initUrl: string | undefined
  const segments: string[] = []

  const lines = m3u8.split('\n')

  for (let i = 0; i < lines.length; i++) {
    const line = lines[i]
    if (line.startsWith('#EXT-X-MAP:URI=')) {
      initUrl = JSON.parse(line.slice('#EXT-X-MAP:URI='.length))
    } else if (line.startsWith('#EXTINF:')) {
      const segmentUrl = lines[i + 1]
      segments.push(segmentUrl)
      i++
    } else if (line.startsWith('#EXT-X-ENDLIST')) {
      break
    }
  }

  if (initUrl) {
    segments.unshift(initUrl)
  }

  return segments
}

export function concatSegments(options: {
  segments: string[]
  fetch: (url: string) => Promise<Uint8Array>
  poolSize?: number
}): ReadableStream {
  const { segments, fetch, poolSize = 8 } = options

  let nextSegmentIdx = 0
  let nextWorkerSegmentIdx = 0
  const nextSegmentCv = new ConditionVariable()
  const buffer: Record<number, Uint8Array> = {}

  const downloadSegment = async (idx = nextWorkerSegmentIdx++) => {
    // console.log('downloading segment %s', idx)
    const url = segments[idx]
    const chunk = await fetch(url)
    buffer[idx] = chunk

    if (idx === nextSegmentIdx) {
      nextSegmentCv.notify()
    }

    if (nextWorkerSegmentIdx < segments.length) {
      return downloadSegment()
    }
  }

  let error: unknown
  void Promise.all(Array.from({
    length: Math.min(poolSize, segments.length),
  }, downloadSegment))
    .catch((e) => {
      error = e
      nextSegmentCv.notify()
    })

  return new ReadableStream({
    async start(controller) {
      while (true) {
        await nextSegmentCv.wait()
        if (error) {
          controller.error(error)
          return
        }

        while (nextSegmentIdx in buffer) {
          const buf = buffer[nextSegmentIdx]
          delete buffer[nextSegmentIdx]
          nextSegmentIdx++
          controller.enqueue(buf)
        }

        if (nextSegmentIdx >= segments.length) {
          controller.close()
          return
        }
      }
    },
  })
}