AeroFetch/src/main/probe.ts

import { execFile } from 'child_process'
import { existsSync } from 'fs'
import { getYtdlpPath } from './binaries'
import { fmtBytes } from './download'
import { cleanError } from './log'
import { assertHttpUrl } from './url'
import { entryUrl, fmtDuration, type RawEntry } from './indexerCore'
import {
  BEST_FORMAT_ID,
  type ProbeResult,
  type MediaInfo,
  type FormatOption,
  type PlaylistInfo,
  type PlaylistEntry
} from '@shared/ipc'

// The slice of `yt-dlp -J` output we actually read.
interface RawFormat {
  format_id: string
  ext?: string
  vcodec?: string
  acodec?: string
  height?: number
  fps?: number
  tbr?: number
  filesize?: number
  filesize_approx?: number
}

interface RawInfo {
  _type?: string
  title?: string
  uploader?: string
  channel?: string
  duration_string?: string
  thumbnail?: string
  formats?: RawFormat[]
  entries?: RawEntry[]
}

/** Quality proxy for picking the best format at a given height. */
function score(f: RawFormat): number {
  return (f.tbr ?? 0) + (f.ext === 'mp4' ? 0.5 : 0)
}

function toOption(f: RawFormat): FormatOption {
  const size = f.filesize ?? f.filesize_approx
  const fpsTag = f.fps && f.fps > 30 ? String(Math.round(f.fps)) : ''
  const sizeLabel = size ? fmtBytes(size) : undefined
  const parts = [`${f.height}p${fpsTag}`, f.ext, sizeLabel].filter(Boolean) as string[]
  return {
    id: f.format_id,
    label: parts.join(' · '),
    height: f.height,
    ext: f.ext,
    filesizeLabel: sizeLabel,
    hasAudio: !!f.acodec && f.acodec !== 'none'
  }
}

/** One option per distinct height (best at that height), best-first, plus auto. */
function buildVideoFormats(raw: RawFormat[]): FormatOption[] {
  const bestByHeight = new Map<number, RawFormat>()
  for (const f of raw) {
    if (!f.vcodec || f.vcodec === 'none' || !f.height) continue
    const cur = bestByHeight.get(f.height)
    if (!cur || score(f) > score(cur)) bestByHeight.set(f.height, f)
  }
  const options = [...bestByHeight.values()]
    .sort((a, b) => (b.height ?? 0) - (a.height ?? 0))
    .map(toOption)
  options.unshift({ id: BEST_FORMAT_ID, label: 'Best available', ext: 'mp4', hasAudio: true })
  return options
}

function buildInfo(data: RawInfo): MediaInfo {
  return {
    title: data.title || 'Untitled',
    channel: data.channel || data.uploader || undefined,
    durationLabel: data.duration_string || undefined,
    thumbnail: data.thumbnail || undefined,
    formats: buildVideoFormats(data.formats ?? [])
  }
}

function buildPlaylist(data: RawInfo): PlaylistInfo {
  const entries: PlaylistEntry[] = []
  ;(data.entries ?? []).forEach((e, i) => {
    const url = entryUrl(e)
    if (!url) return // skip entries we can't turn into a downloadable URL
    entries.push({
      index: i + 1,
      id: e.id || String(i + 1),
      title: e.title || `Item ${i + 1}`,
      url,
      durationLabel: fmtDuration(e.duration),
      uploader: e.uploader || e.channel || undefined
    })
  })
  return {
    title: data.title || 'Playlist',
    uploader: data.uploader || data.channel || undefined,
    count: entries.length,
    entries
  }
}

/**
 * Probe a URL with `yt-dlp -J --flat-playlist`. Resolves to a single video (with
 * its format list) or, when the URL is a playlist, the flat list of its entries.
 * `--flat-playlist` is a no-op for a lone video, so one call covers both cases.
 */
export function probeMedia(url: string): Promise<ProbeResult> {
  const ytdlp = getYtdlpPath()
  if (!existsSync(ytdlp)) {
    return Promise.resolve({
      ok: false,
      error: `yt-dlp.exe not found at ${ytdlp}\nDrop it into resources/bin/ (see the README there).`
    })
  }
  let target: string
  try {
    target = assertHttpUrl(url) // normalised form (audit F5)
  } catch (e) {
    return Promise.resolve({ ok: false, error: (e as Error).message })
  }

  return new Promise((resolve) => {
    execFile(
      ytdlp,
      // `--` terminates option parsing so the URL can never be read as a flag.
      ['-J', '--flat-playlist', '--no-warnings', '--', target],
      { windowsHide: true, maxBuffer: 64 * 1024 * 1024, timeout: 60_000 },
      (err, stdout, stderr) => {
        if (err) {
          // execFile sets `killed` when it terminated the process on timeout.
          const msg = (err as { killed?: boolean }).killed
            ? 'Timed out fetching video info. Check the link or your connection.'
            : cleanError(stderr) || err.message
          resolve({ ok: false, error: msg })
          return
        }
        try {
          const data = JSON.parse(stdout) as RawInfo
          if (data._type === 'playlist' || Array.isArray(data.entries)) {
            const playlist = buildPlaylist(data)
            if (playlist.count === 0) {
              resolve({ ok: false, error: 'This playlist has no downloadable entries.' })
              return
            }
            resolve({ ok: true, kind: 'playlist', playlist })
          } else {
            resolve({ ok: true, kind: 'video', info: buildInfo(data) })
          }
        } catch {
          resolve({ ok: false, error: 'Could not parse video info from yt-dlp.' })
        }
      }
    )
  })
}