AeroFetch/src/main/download.ts

import { spawn, execFile, type ChildProcess } from 'child_process'
import { existsSync } from 'fs'
import { join } from 'path'
import { app, BrowserWindow, Notification, type WebContents } from 'electron'
import {
  getYtdlpPath,
  getBinDir,
  getAria2cPath,
  getFfmpegPath,
  getFfprobePath,
  getSystem32Path
} from './binaries'
import { getSettings, getDownloadArchivePath } from './settings'
import { getCookiesFilePath } from './cookies'
import { listTemplates } from './templates'
import { assertHttpUrl } from './url'
import { isSafeOutputDir } from './validation'
import {
  buildArgs,
  selectExtraArgs,
  formatCommandLine,
  collectionOutputTemplate
} from './buildArgs'
import { cleanError } from './log'
import { addErrorLog } from './errorlog'
import {
  IpcChannels,
  type StartDownloadOptions,
  type StartDownloadResult,
  type CommandPreviewResult,
  type DownloadEvent,
  type DownloadMeta,
  type DownloadProgress,
  type Settings
} from '@shared/ipc'

interface ActiveDownload {
  child: ChildProcess
  canceled: boolean
}

const active = new Map<string, ActiveDownload>()

// --- Formatting helpers (raw yt-dlp numbers → human strings) ----------------

function num(s?: string): number | undefined {
  if (!s || s === 'NA') return undefined
  const n = Number(s)
  return Number.isFinite(n) ? n : undefined
}

export function fmtBytes(bytes: number): string {
  if (bytes < 1024) return `${bytes} B`
  const units = ['KB', 'MB', 'GB', 'TB']
  let v = bytes / 1024
  let i = 0
  while (v >= 1024 && i < units.length - 1) {
    v /= 1024
    i++
  }
  return `${v.toFixed(v >= 100 ? 0 : 1)} ${units[i]}`
}

function fmtSpeed(bytesPerSec?: number): string | undefined {
  if (bytesPerSec == null) return undefined
  return `${fmtBytes(bytesPerSec)}/s`
}

function fmtEta(seconds?: number): string | undefined {
  if (seconds == null) return undefined
  const s = Math.max(0, Math.round(seconds))
  const m = Math.floor(s / 60)
  const r = s % 60
  return `${m}:${String(r).padStart(2, '0')}`
}

function parseProgress(rest: string): DownloadProgress | null {
  const parts = rest.split('|')
  if (parts.length < 6) return null
  const [status, dl, total, totalEst, speed, eta] = parts
  const downloaded = num(dl)
  const totalBytes = num(total) ?? num(totalEst)
  let progress = 0
  if (totalBytes && downloaded != null) progress = Math.min(1, downloaded / totalBytes)
  return {
    status: status || 'downloading',
    progress,
    speed: fmtSpeed(num(speed)),
    eta: fmtEta(num(eta)),
    sizeLabel: totalBytes ? fmtBytes(totalBytes) : undefined
  }
}

function send(wc: WebContents, ev: DownloadEvent): void {
  if (!wc.isDestroyed()) wc.send(IpcChannels.downloadEvent, ev)
}

/** Native OS notification on completion/failure, gated by Settings.notifyOnComplete. */
function notify(wc: WebContents, title: string, body: string): void {
  if (!getSettings().notifyOnComplete || !Notification.isSupported()) return
  const n = new Notification({ title, body })
  n.on('click', () => {
    if (wc.isDestroyed()) return
    const win = BrowserWindow.fromWebContents(wc)
    if (win) {
      if (win.isMinimized()) win.restore()
      win.show()
      win.focus()
    }
  })
  n.show()
}

function logFailure(opts: StartDownloadOptions, title: string | undefined, error: string): void {
  addErrorLog({
    id: opts.id,
    title,
    url: opts.url,
    kind: opts.kind,
    error,
    occurredAt: Date.now()
  })
}

// --- Best-effort metadata probe (runs alongside the download) ---------------

function probeMeta(ytdlp: string, url: string): Promise<DownloadMeta | null> {
  return new Promise((resolve) => {
    execFile(
      ytdlp,
      [
        '--no-playlist',
        '--no-warnings',
        '--skip-download',
        '--print',
        'title',
        '--print',
        'uploader',
        '--print',
        'duration_string',
        '--',
        url
      ],
      { windowsHide: true, maxBuffer: 4 * 1024 * 1024, timeout: 30_000 },
      (err, stdout) => {
        if (err) return resolve(null)
        const [title, uploader, duration] = stdout.split('\n').map((l) => l.trim())
        const clean = (v?: string): string | undefined =>
          v && v !== 'NA' ? v : undefined
        resolve({
          title: clean(title),
          channel: clean(uploader),
          durationLabel: clean(duration)
        })
      }
    )
  })
}

// --- Argv construction (shared by startDownload and the command preview) ---

// A per-download override (opts.extraArgs, even '') wins over the persisted
// default template — but BOTH are gated on the customCommandEnabled consent flag
// (see selectExtraArgs / audit F2). The gate is enforced here in main, not just
// in the renderer UI, so the renderer can't be trusted to apply it.
function resolveExtraArgs(opts: StartDownloadOptions, settings: Settings): string[] {
  return selectExtraArgs({
    customCommandEnabled: settings.customCommandEnabled,
    perDownloadExtraArgs: opts.extraArgs,
    defaultTemplateId: settings.defaultTemplateId,
    templates: listTemplates()
  })
}

/** Resolve settings + per-download overrides into the full yt-dlp argv. */
export function buildCommand(opts: StartDownloadOptions): string[] {
  const settings = getSettings()
  // A per-download outputDir override must clear the same safety check the
  // persisted setting does (absolute path only); a renderer-supplied override is
  // otherwise untrusted — it dictates where downloaded files get written. An
  // unsafe override is ignored in favour of the validated setting, then the OS
  // Downloads folder. (audit F4)
  const override = opts.outputDir?.trim()
  const outDir =
    (override && isSafeOutputDir(override) ? override : '') ||
    settings.outputDir ||
    app.getPath('downloads')
  // A collection (media-manager) download is filed into <channel>/<playlist>/
  // <NNN> - <title> folders; an ordinary download uses the flat filenameTemplate.
  const filenameTemplate = settings.filenameTemplate?.trim() || '%(title)s.%(ext)s'
  const outputTemplate = opts.collection
    ? collectionOutputTemplate(outDir, opts.collection, '%(title)s.%(ext)s')
    : join(outDir, filenameTemplate)
  // Per-download override wins; otherwise use the persisted defaults.
  const options = opts.options ?? settings.downloadOptions
  // Silently fall back to yt-dlp's own downloader if aria2c.exe wasn't dropped
  // into resources/bin — the toggle shouldn't turn into a hard error.
  const aria2cPath = settings.useAria2c && existsSync(getAria2cPath()) ? getAria2cPath() : undefined
  // Same idea: 'login' cookies only apply once the sign-in window has actually
  // exported a file; otherwise the download proceeds cookie-less rather than failing.
  const cookiesFile =
    settings.cookieSource === 'login' && existsSync(getCookiesFilePath())
      ? getCookiesFilePath()
      : undefined
  const access = {
    proxy: settings.proxy,
    rateLimit: settings.rateLimit,
    aria2cPath,
    cookiesFromBrowser: settings.cookieSource === 'browser' ? settings.cookiesBrowser : undefined,
    cookiesFile,
    restrictFilenames: settings.restrictFilenames,
    downloadArchivePath: settings.downloadArchive ? getDownloadArchivePath() : undefined
  }
  const extraArgs = resolveExtraArgs(opts, settings)
  return buildArgs(opts, outputTemplate, options, getBinDir(), access, extraArgs)
}

/** Build the exact command line for the current form state, without running it. */
export function previewCommand(opts: StartDownloadOptions): CommandPreviewResult {
  let normalized: StartDownloadOptions
  try {
    // Use the parser-normalised URL (audit F5), so the previewed command matches
    // exactly what startDownload would spawn.
    normalized = { ...opts, url: assertHttpUrl(opts.url) }
  } catch (e) {
    return { ok: false, error: (e as Error).message }
  }
  try {
    return { ok: true, command: formatCommandLine(getYtdlpPath(), buildCommand(normalized)) }
  } catch (e) {
    return { ok: false, error: (e as Error).message }
  }
}

// --- Public API -------------------------------------------------------------

export function startDownload(
  wc: WebContents,
  opts: StartDownloadOptions
): StartDownloadResult {
  const ytdlp = getYtdlpPath()
  if (!existsSync(ytdlp)) {
    return {
      ok: false,
      error: `yt-dlp.exe not found at ${ytdlp}\nDrop it into resources/bin/ (see the README there).`
    }
  }
  // ffmpeg is used by nearly every download (merge, audio extract, thumbnail/
  // metadata embed) and ffprobe by duration-aware post-processing (SponsorBlock-
  // remove, --force-keyframes-at-cuts, --split-chapters). Assert both up front so a
  // missing binary is a clear AeroFetch message, not a cryptic mid-download yt-dlp
  // postprocessing error (e.g. "Unable to determine video duration: ffprobe not found").
  const missingBins: string[] = []
  if (!existsSync(getFfmpegPath())) missingBins.push('ffmpeg.exe')
  if (!existsSync(getFfprobePath())) missingBins.push('ffprobe.exe')
  if (missingBins.length > 0) {
    return {
      ok: false,
      error:
        `${missingBins.join(' and ')} not found in ${getBinDir()}\n` +
        `Add the ffmpeg build's binaries to resources/bin/ (see the README there).`
    }
  }
  // Reject anything that isn't an http(s) URL before it reaches yt-dlp's argv,
  // and replace opts.url with the parser-normalised form so the exact string we
  // validated is the one that gets spawned/probed — not a raw variant carrying
  // interior tabs/newlines or leading control chars that URL parsing silently
  // tolerates. (audit F5)
  try {
    opts = { ...opts, url: assertHttpUrl(opts.url) }
  } catch (e) {
    return { ok: false, error: (e as Error).message }
  }
  if (active.has(opts.id)) {
    return { ok: false, error: 'A download with this id is already running.' }
  }
  // Defence-in-depth: the renderer's pump() already caps concurrency, but the main
  // process shouldn't trust it — a buggy or compromised renderer could otherwise
  // spawn unbounded yt-dlp processes. Enforce the same cap here on active spawns.
  const maxConcurrent = getSettings().maxConcurrent
  if (active.size >= maxConcurrent) {
    return { ok: false, error: 'Max concurrent downloads reached. Wait for a slot to free up.' }
  }

  let child: ChildProcess
  try {
    child = spawn(ytdlp, buildCommand(opts), { windowsHide: true })
  } catch (e) {
    return { ok: false, error: (e as Error).message }
  }

  const rec: ActiveDownload = { child, canceled: false }
  active.set(opts.id, rec)

  // Title/channel/duration are kept locally so the completion notification and
  // error log can show a real title instead of just the raw URL.
  let resolvedTitle: string | undefined
  if (opts.meta && (opts.meta.title || opts.meta.channel || opts.meta.durationLabel)) {
    // The renderer already probed this URL and passed the metadata along — reuse
    // it instead of spawning a redundant second yt-dlp probe (audit P2).
    resolvedTitle = opts.meta.title
    send(wc, { type: 'meta', id: opts.id, meta: opts.meta })
  } else {
    // No pre-probed metadata (e.g. a direct paste that skipped the probe) — fetch
    // it in parallel so the card fills in quickly.
    probeMeta(ytdlp, opts.url).then((meta) => {
      if (meta?.title) resolvedTitle = meta.title
      if (meta && active.has(opts.id)) send(wc, { type: 'meta', id: opts.id, meta })
    })
  }

  let stdoutBuf = ''
  let stderrTail = ''
  let filePath: string | undefined
  // 'error' and 'close' can both fire for one process; only act on the first.
  let settled = false

  child.stdout?.on('data', (chunk: Buffer) => {
    stdoutBuf += chunk.toString()
    let nl: number
    while ((nl = stdoutBuf.indexOf('\n')) >= 0) {
      const line = stdoutBuf.slice(0, nl).replace(/\r$/, '')
      stdoutBuf = stdoutBuf.slice(nl + 1)
      if (line.startsWith('prog|')) {
        const p = parseProgress(line.slice('prog|'.length))
        if (p) send(wc, { type: 'progress', id: opts.id, progress: p })
      } else if (line.startsWith('path|')) {
        filePath = line.slice('path|'.length).trim()
      }
    }
  })

  child.stderr?.on('data', (chunk: Buffer) => {
    stderrTail = (stderrTail + chunk.toString()).slice(-4000)
  })

  child.on('error', (err) => {
    if (settled) return
    settled = true
    active.delete(opts.id)
    if (!rec.canceled) {
      send(wc, { type: 'error', id: opts.id, error: err.message })
      logFailure(opts, resolvedTitle, err.message)
      notify(wc, resolvedTitle ?? 'Download failed', err.message)
    }
  })

  child.on('close', (code) => {
    if (settled) return
    settled = true
    active.delete(opts.id)
    if (rec.canceled) return // renderer already showed 'canceled' optimistically
    if (code === 0) {
      send(wc, { type: 'done', id: opts.id, filePath })
      notify(wc, resolvedTitle ?? 'Download complete', 'Finished downloading.')
    } else {
      const msg = cleanError(stderrTail) || `yt-dlp exited with code ${code}`
      send(wc, { type: 'error', id: opts.id, error: msg })
      logFailure(opts, resolvedTitle, msg)
      notify(wc, resolvedTitle ?? 'Download failed', msg)
    }
  })

  return { ok: true }
}

export function cancelDownload(id: string): void {
  const rec = active.get(id)
  if (!rec) return
  rec.canceled = true
  const pid = rec.child.pid
  if (pid != null) {
    // Kill the whole tree (/T) so the spawned ffmpeg child dies too. Resolve
    // taskkill from System32 by absolute path, never the bare name, so a planted
    // taskkill.exe on PATH / in the CWD can't run in its place. (audit F3)
    execFile(
      getSystem32Path('taskkill.exe'),
      ['/pid', String(pid), '/T', '/F'],
      { windowsHide: true },
      () => {}
    )
  } else {
    rec.child.kill()
  }
}