SubMiner/src/core/services/youtube/track-probe.ts

import { spawn } from 'node:child_process';
import type { YoutubeTrackOption } from '../../../types';
import { formatYoutubeTrackLabel, normalizeYoutubeLangCode, type YoutubeTrackKind } from './labels';

const YOUTUBE_TRACK_PROBE_TIMEOUT_MS = 15_000;

export type YoutubeTrackProbeResult = {
  videoId: string;
  title: string;
  tracks: YoutubeTrackOption[];
};

type YtDlpSubtitleEntry = Array<{ ext?: string; name?: string; url?: string }>;

type YtDlpInfo = {
  id?: string;
  title?: string;
  subtitles?: Record<string, YtDlpSubtitleEntry>;
  automatic_captions?: Record<string, YtDlpSubtitleEntry>;
};

function runCapture(
  command: string,
  args: string[],
  timeoutMs = YOUTUBE_TRACK_PROBE_TIMEOUT_MS,
): Promise<{ stdout: string; stderr: string }> {
  return new Promise((resolve, reject) => {
    const proc = spawn(command, args, { stdio: ['ignore', 'pipe', 'pipe'] });
    let stdout = '';
    let stderr = '';
    const timer = setTimeout(() => {
      proc.kill();
      reject(new Error(`yt-dlp timed out after ${timeoutMs}ms`));
    }, timeoutMs);
    proc.stdout.setEncoding('utf8');
    proc.stderr.setEncoding('utf8');
    proc.stdout.on('data', (chunk) => {
      stdout += String(chunk);
    });
    proc.stderr.on('data', (chunk) => {
      stderr += String(chunk);
    });
    proc.once('error', (error) => {
      clearTimeout(timer);
      reject(error);
    });
    proc.once('close', (code) => {
      clearTimeout(timer);
      if (code === 0) {
        resolve({ stdout, stderr });
        return;
      }
      reject(new Error(stderr.trim() || `yt-dlp exited with status ${code ?? 'unknown'}`));
    });
  });
}

function choosePreferredFormat(
  formats: YtDlpSubtitleEntry,
  kind: YoutubeTrackKind,
): { ext: string; url: string; title?: string } | null {
  const preferredOrder =
    kind === 'auto'
      ? ['srv3', 'srv2', 'srv1', 'vtt', 'srt', 'ttml', 'json3']
      : ['srt', 'vtt', 'srv3', 'srv2', 'srv1', 'ttml', 'json3'];
  for (const ext of preferredOrder) {
    const match = formats.find(
      (format) => typeof format.url === 'string' && format.url && format.ext === ext,
    );
    if (match?.url) {
      return { ext, url: match.url, title: match.name?.trim() || undefined };
    }
  }

  const fallback = formats.find((format) => typeof format.url === 'string' && format.url);
  if (!fallback?.url) {
    return null;
  }

  return {
    ext: fallback.ext?.trim() || 'vtt',
    url: fallback.url,
    title: fallback.name?.trim() || undefined,
  };
}

function toTracks(entries: Record<string, YtDlpSubtitleEntry> | undefined, kind: YoutubeTrackKind) {
  const tracks: YoutubeTrackOption[] = [];
  if (!entries) return tracks;
  for (const [language, formats] of Object.entries(entries)) {
    if (!Array.isArray(formats) || formats.length === 0) continue;
    const preferredFormat = choosePreferredFormat(formats, kind);
    if (!preferredFormat) continue;
    const sourceLanguage = language.trim() || language;
    const normalizedLanguage = normalizeYoutubeLangCode(sourceLanguage) || sourceLanguage;
    const title = preferredFormat.title;
    tracks.push({
      id: `${kind}:${sourceLanguage}`,
      language: normalizedLanguage,
      sourceLanguage,
      kind,
      title,
      label: formatYoutubeTrackLabel({ language: normalizedLanguage, kind, title }),
      downloadUrl: preferredFormat.url,
      fileExtension: preferredFormat.ext,
    });
  }
  return tracks;
}

export type { YoutubeTrackOption };

export async function probeYoutubeTracks(targetUrl: string): Promise<YoutubeTrackProbeResult> {
  const { stdout } = await runCapture('yt-dlp', ['--dump-single-json', '--no-warnings', targetUrl]);
  const trimmedStdout = stdout.trim();
  if (!trimmedStdout) {
    throw new Error('yt-dlp returned empty output while probing subtitle tracks');
  }
  let info: YtDlpInfo;
  try {
    info = JSON.parse(trimmedStdout) as YtDlpInfo;
  } catch (error) {
    const snippet = trimmedStdout.slice(0, 200);
    throw new Error(
      `Failed to parse yt-dlp output as JSON: ${
        error instanceof Error ? error.message : String(error)
      }${snippet ? `; stdout=${snippet}` : ''}`,
    );
  }
  const tracks = [...toTracks(info.subtitles, 'manual'), ...toTracks(info.automatic_captions, 'auto')];
  return {
    videoId: info.id || '',
    title: info.title || '',
    tracks,
  };
}