AeroFetch/src/main/probe.ts

import { existsSync } from 'fs'
import { getYtdlpPath, YTDLP_MISSING_MSG } from './binaries'
import { execFileAsync } from './lib/exec'
import { PROBE_MAX_BUFFER, PROBE_TIMEOUT_MS } from './constants'
import { fmtBytes } from '@shared/format'
import { cleanError } from './log'
import { assertHttpUrl } from './url'
import { entryUrl, fmtDuration, type RawEntry } from './core/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)
  // No `ext`: the final container for the auto-best pick depends on the user's
  // videoContainer setting (mp4/mkv/webm), unknown here -- claiming 'mp4' would
  // mislabel mkv/webm outputs (L62). hasAudio is always true (best merges audio).
  options.unshift({ id: BEST_FORMAT_ID, label: 'Best available', 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 async function probeMedia(url: string): Promise<ProbeResult> {
  const ytdlp = getYtdlpPath()
  if (!existsSync(ytdlp)) {
    return { ok: false, error: YTDLP_MISSING_MSG }
  }
  let target: string
  try {
    target = assertHttpUrl(url) // normalised form (audit F5)
  } catch (e) {
    return { ok: false, error: (e as Error).message }
  }

  // `--` terminates option parsing so the URL can never be read as a flag.
  const r = await execFileAsync(ytdlp, ['-J', '--flat-playlist', '--no-warnings', '--', target], {
    maxBuffer: PROBE_MAX_BUFFER,
    timeout: PROBE_TIMEOUT_MS
  })
  if (!r.ok) {
    const msg = r.timedOut
      ? 'Timed out fetching video info. Check the link or your connection.'
      : cleanError(r.stderr) || r.error?.message || ''
    return { ok: false, error: msg }
  }
  try {
    const data = JSON.parse(r.stdout) as RawInfo
    if (data._type === 'playlist' || Array.isArray(data.entries)) {
      const playlist = buildPlaylist(data)
      if (playlist.count === 0) {
        return { ok: false, error: 'This playlist has no downloadable entries.' }
      }
      return { ok: true, kind: 'playlist', playlist }
    }
    return { ok: true, kind: 'video', info: buildInfo(data) }
  } catch {
    return { ok: false, error: 'Could not parse video info from yt-dlp.' }
  }
}