SubMiner/launcher/youtube.ts

import fs from "node:fs";
import path from "node:path";
import os from "node:os";
import type { Args, SubtitleCandidate, YoutubeSubgenOutputs } from "./types.js";
import { YOUTUBE_SUB_EXTENSIONS, YOUTUBE_AUDIO_EXTENSIONS } from "./types.js";
import { log } from "./log.js";
import {
  resolvePathMaybe, uniqueNormalizedLangCodes,
  escapeRegExp, normalizeBasename, runExternalCommand, commandExists,
} from "./util.js";
import { state } from "./mpv.js";

function toYtdlpLangPattern(langCodes: string[]): string {
  return langCodes.map((lang) => `${lang}.*`).join(",");
}

function filenameHasLanguageTag(filenameLower: string, langCode: string): boolean {
  const escaped = escapeRegExp(langCode);
  const pattern = new RegExp(`(^|[._-])${escaped}([._-]|$)`);
  return pattern.test(filenameLower);
}

function classifyLanguage(
  filename: string,
  primaryLangCodes: string[],
  secondaryLangCodes: string[],
): "primary" | "secondary" | null {
  const lower = filename.toLowerCase();
  const primary = primaryLangCodes.some((code) =>
    filenameHasLanguageTag(lower, code),
  );
  const secondary = secondaryLangCodes.some((code) =>
    filenameHasLanguageTag(lower, code),
  );
  if (primary && !secondary) return "primary";
  if (secondary && !primary) return "secondary";
  return null;
}

function preferredLangLabel(langCodes: string[], fallback: string): string {
  return uniqueNormalizedLangCodes(langCodes)[0] || fallback;
}

function sourceTag(source: SubtitleCandidate["source"]): string {
  if (source === "manual" || source === "auto") return `ytdlp-${source}`;
  if (source === "whisper-translate") return "whisper-translate";
  return "whisper";
}

function pickBestCandidate(candidates: SubtitleCandidate[]): SubtitleCandidate | null {
  if (candidates.length === 0) return null;
  const scored = [...candidates].sort((a, b) => {
    const sourceA = a.source === "manual" ? 1 : 0;
    const sourceB = b.source === "manual" ? 1 : 0;
    if (sourceA !== sourceB) return sourceB - sourceA;
    const srtA = a.ext === ".srt" ? 1 : 0;
    const srtB = b.ext === ".srt" ? 1 : 0;
    if (srtA !== srtB) return srtB - srtA;
    return b.size - a.size;
  });
  return scored[0];
}

function scanSubtitleCandidates(
  tempDir: string,
  knownSet: Set<string>,
  source: "manual" | "auto",
  primaryLangCodes: string[],
  secondaryLangCodes: string[],
): SubtitleCandidate[] {
  const entries = fs.readdirSync(tempDir);
  const out: SubtitleCandidate[] = [];
  for (const name of entries) {
    const fullPath = path.join(tempDir, name);
    if (knownSet.has(fullPath)) continue;
    let stat: fs.Stats;
    try {
      stat = fs.statSync(fullPath);
    } catch {
      continue;
    }
    if (!stat.isFile()) continue;
    const ext = path.extname(fullPath).toLowerCase();
    if (!YOUTUBE_SUB_EXTENSIONS.has(ext)) continue;
    const lang = classifyLanguage(name, primaryLangCodes, secondaryLangCodes);
    if (!lang) continue;
    out.push({ path: fullPath, lang, ext, size: stat.size, source });
  }
  return out;
}

async function convertToSrt(
  inputPath: string,
  tempDir: string,
  langLabel: string,
): Promise<string> {
  if (path.extname(inputPath).toLowerCase() === ".srt") return inputPath;
  const outputPath = path.join(tempDir, `converted.${langLabel}.srt`);
  await runExternalCommand("ffmpeg", ["-y", "-loglevel", "error", "-i", inputPath, outputPath]);
  return outputPath;
}

function findAudioFile(tempDir: string, preferredExt: string): string | null {
  const entries = fs.readdirSync(tempDir);
  const audioFiles: Array<{ path: string; ext: string; mtimeMs: number }> = [];
  for (const name of entries) {
    const fullPath = path.join(tempDir, name);
    let stat: fs.Stats;
    try {
      stat = fs.statSync(fullPath);
    } catch {
      continue;
    }
    if (!stat.isFile()) continue;
    const ext = path.extname(name).toLowerCase();
    if (!YOUTUBE_AUDIO_EXTENSIONS.has(ext)) continue;
    audioFiles.push({ path: fullPath, ext, mtimeMs: stat.mtimeMs });
  }
  if (audioFiles.length === 0) return null;
  const preferred = audioFiles.find((entry) => entry.ext === `.${preferredExt.toLowerCase()}`);
  if (preferred) return preferred.path;
  audioFiles.sort((a, b) => b.mtimeMs - a.mtimeMs);
  return audioFiles[0].path;
}

async function runWhisper(
  whisperBin: string,
  modelPath: string,
  audioPath: string,
  language: string,
  translate: boolean,
  outputPrefix: string,
): Promise<string> {
  const args = [
    "-m",
    modelPath,
    "-f",
    audioPath,
    "--output-srt",
    "--output-file",
    outputPrefix,
    "--language",
    language,
  ];
  if (translate) args.push("--translate");
  await runExternalCommand(whisperBin, args, {
    commandLabel: "whisper",
    streamOutput: true,
  });
  const outputPath = `${outputPrefix}.srt`;
  if (!fs.existsSync(outputPath)) {
    throw new Error(`whisper output not found: ${outputPath}`);
  }
  return outputPath;
}

async function convertAudioForWhisper(inputPath: string, tempDir: string): Promise<string> {
  const wavPath = path.join(tempDir, "whisper-input.wav");
  await runExternalCommand("ffmpeg", [
    "-y",
    "-loglevel",
    "error",
    "-i",
    inputPath,
    "-ar",
    "16000",
    "-ac",
    "1",
    "-c:a",
    "pcm_s16le",
    wavPath,
  ]);
  if (!fs.existsSync(wavPath)) {
    throw new Error(`Failed to prepare whisper audio input: ${wavPath}`);
  }
  return wavPath;
}

export function resolveWhisperBinary(args: Args): string | null {
  const explicit = args.whisperBin.trim();
  if (explicit) return resolvePathMaybe(explicit);
  if (commandExists("whisper-cli")) return "whisper-cli";
  return null;
}

export async function generateYoutubeSubtitles(
  target: string,
  args: Args,
  onReady?: (lang: "primary" | "secondary", pathToLoad: string) => Promise<void>,
): Promise<YoutubeSubgenOutputs> {
  const outDir = path.resolve(resolvePathMaybe(args.youtubeSubgenOutDir));
  fs.mkdirSync(outDir, { recursive: true });

  const primaryLangCodes = uniqueNormalizedLangCodes(args.youtubePrimarySubLangs);
  const secondaryLangCodes = uniqueNormalizedLangCodes(args.youtubeSecondarySubLangs);
  const primaryLabel = preferredLangLabel(primaryLangCodes, "primary");
  const secondaryLabel = preferredLangLabel(secondaryLangCodes, "secondary");
  const secondaryCanUseWhisperTranslate =
    secondaryLangCodes.includes("en") || secondaryLangCodes.includes("eng");
  const ytdlpManualLangs = toYtdlpLangPattern([
    ...primaryLangCodes,
    ...secondaryLangCodes,
  ]);

  const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "subminer-yt-subgen-"));
  const knownFiles = new Set<string>();
  let keepTemp = args.youtubeSubgenKeepTemp;

  const publishTrack = async (
    lang: "primary" | "secondary",
    source: SubtitleCandidate["source"],
    selectedPath: string,
    basename: string,
  ): Promise<string> => {
    const langLabel = lang === "primary" ? primaryLabel : secondaryLabel;
    const taggedPath = path.join(
      outDir,
      `${basename}.${langLabel}.${sourceTag(source)}.srt`,
    );
    const aliasPath = path.join(outDir, `${basename}.${langLabel}.srt`);
    fs.copyFileSync(selectedPath, taggedPath);
    fs.copyFileSync(taggedPath, aliasPath);
    log(
      "info",
      args.logLevel,
      `Generated subtitle (${langLabel}, ${source}) -> ${aliasPath}`,
    );
    if (onReady) await onReady(lang, aliasPath);
    return aliasPath;
  };

  try {
    log("debug", args.logLevel, `YouTube subtitle temp dir: ${tempDir}`);
    const meta = await runExternalCommand(
      "yt-dlp",
      ["--dump-single-json", "--no-warnings", target],
      {
        captureStdout: true,
        logLevel: args.logLevel,
        commandLabel: "yt-dlp:meta",
      },
      state.youtubeSubgenChildren,
    );
    const metadata = JSON.parse(meta.stdout) as { id?: string };
    const videoId = metadata.id || `${Date.now()}`;
    const basename = normalizeBasename(videoId, videoId);

    await runExternalCommand(
      "yt-dlp",
      [
        "--skip-download",
        "--no-warnings",
        "--write-subs",
        "--sub-format",
        "srt/vtt/best",
        "--sub-langs",
        ytdlpManualLangs,
        "-o",
        path.join(tempDir, "%(id)s.%(ext)s"),
        target,
      ],
      {
        allowFailure: true,
        logLevel: args.logLevel,
        commandLabel: "yt-dlp:manual-subs",
        streamOutput: true,
      },
      state.youtubeSubgenChildren,
    );

    const manualSubs = scanSubtitleCandidates(
      tempDir,
      knownFiles,
      "manual",
      primaryLangCodes,
      secondaryLangCodes,
    );
    for (const sub of manualSubs) knownFiles.add(sub.path);
    let primaryCandidates = manualSubs.filter((entry) => entry.lang === "primary");
    let secondaryCandidates = manualSubs.filter(
      (entry) => entry.lang === "secondary",
    );

    const missingAuto: string[] = [];
    if (primaryCandidates.length === 0)
      missingAuto.push(toYtdlpLangPattern(primaryLangCodes));
    if (secondaryCandidates.length === 0)
      missingAuto.push(toYtdlpLangPattern(secondaryLangCodes));

    if (missingAuto.length > 0) {
      await runExternalCommand(
        "yt-dlp",
        [
          "--skip-download",
          "--no-warnings",
          "--write-auto-subs",
          "--sub-format",
          "srt/vtt/best",
          "--sub-langs",
          missingAuto.join(","),
          "-o",
          path.join(tempDir, "%(id)s.%(ext)s"),
          target,
        ],
        {
          allowFailure: true,
          logLevel: args.logLevel,
          commandLabel: "yt-dlp:auto-subs",
          streamOutput: true,
        },
        state.youtubeSubgenChildren,
      );

      const autoSubs = scanSubtitleCandidates(
        tempDir,
        knownFiles,
        "auto",
        primaryLangCodes,
        secondaryLangCodes,
      );
      for (const sub of autoSubs) knownFiles.add(sub.path);
      primaryCandidates = primaryCandidates.concat(
        autoSubs.filter((entry) => entry.lang === "primary"),
      );
      secondaryCandidates = secondaryCandidates.concat(
        autoSubs.filter((entry) => entry.lang === "secondary"),
      );
    }

    let primaryAlias = "";
    let secondaryAlias = "";
    const selectedPrimary = pickBestCandidate(primaryCandidates);
    const selectedSecondary = pickBestCandidate(secondaryCandidates);

    if (selectedPrimary) {
      const srt = await convertToSrt(selectedPrimary.path, tempDir, primaryLabel);
      primaryAlias = await publishTrack(
        "primary",
        selectedPrimary.source,
        srt,
        basename,
      );
    }
    if (selectedSecondary) {
      const srt = await convertToSrt(
        selectedSecondary.path,
        tempDir,
        secondaryLabel,
      );
      secondaryAlias = await publishTrack(
        "secondary",
        selectedSecondary.source,
        srt,
        basename,
      );
    }

    const needsPrimaryWhisper = !selectedPrimary;
    const needsSecondaryWhisper = !selectedSecondary && secondaryCanUseWhisperTranslate;
    if (needsPrimaryWhisper || needsSecondaryWhisper) {
      const whisperBin = resolveWhisperBinary(args);
      const modelPath = args.whisperModel.trim()
        ? path.resolve(resolvePathMaybe(args.whisperModel.trim()))
        : "";
      const hasWhisperFallback = !!whisperBin && !!modelPath && fs.existsSync(modelPath);

      if (!hasWhisperFallback) {
        log(
          "warn",
          args.logLevel,
          "Whisper fallback is not configured; continuing with available subtitle tracks.",
        );
      } else {
        try {
          await runExternalCommand(
            "yt-dlp",
            [
              "-f",
              "bestaudio/best",
              "--extract-audio",
              "--audio-format",
              args.youtubeSubgenAudioFormat,
              "--no-warnings",
              "-o",
              path.join(tempDir, "%(id)s.%(ext)s"),
              target,
            ],
            {
              logLevel: args.logLevel,
              commandLabel: "yt-dlp:audio",
              streamOutput: true,
            },
            state.youtubeSubgenChildren,
          );
          const audioPath = findAudioFile(tempDir, args.youtubeSubgenAudioFormat);
          if (!audioPath) {
            throw new Error("Audio extraction succeeded, but no audio file was found.");
          }
          const whisperAudioPath = await convertAudioForWhisper(audioPath, tempDir);

          if (needsPrimaryWhisper) {
            try {
              const primaryPrefix = path.join(tempDir, `${basename}.${primaryLabel}`);
              const primarySrt = await runWhisper(
                whisperBin!,
                modelPath,
                whisperAudioPath,
                args.youtubeWhisperSourceLanguage,
                false,
                primaryPrefix,
              );
              primaryAlias = await publishTrack(
                "primary",
                "whisper",
                primarySrt,
                basename,
              );
            } catch (error) {
              log(
                "warn",
                args.logLevel,
                `Failed to generate primary subtitle via whisper fallback: ${(error as Error).message}`,
              );
            }
          }

          if (needsSecondaryWhisper) {
            try {
              const secondaryPrefix = path.join(
                tempDir,
                `${basename}.${secondaryLabel}`,
              );
              const secondarySrt = await runWhisper(
                whisperBin!,
                modelPath,
                whisperAudioPath,
                args.youtubeWhisperSourceLanguage,
                true,
                secondaryPrefix,
              );
              secondaryAlias = await publishTrack(
                "secondary",
                "whisper-translate",
                secondarySrt,
                basename,
              );
            } catch (error) {
              log(
                "warn",
                args.logLevel,
                `Failed to generate secondary subtitle via whisper fallback: ${(error as Error).message}`,
              );
            }
          }
        } catch (error) {
          log(
            "warn",
            args.logLevel,
            `Whisper fallback pipeline failed: ${(error as Error).message}`,
          );
        }
      }
    }

    if (!secondaryCanUseWhisperTranslate && !selectedSecondary) {
      log(
        "warn",
        args.logLevel,
        `Secondary subtitle language (${secondaryLabel}) has no whisper translate fallback; relying on yt-dlp subtitles only.`,
      );
    }

    if (!primaryAlias && !secondaryAlias) {
      throw new Error("Failed to generate any subtitle tracks.");
    }
    if (!primaryAlias || !secondaryAlias) {
      log(
        "warn",
        args.logLevel,
        `Generated partial subtitle result: primary=${primaryAlias ? "ok" : "missing"}, secondary=${secondaryAlias ? "ok" : "missing"}`,
      );
    }

    return {
      basename,
      primaryPath: primaryAlias || undefined,
      secondaryPath: secondaryAlias || undefined,
    };
  } catch (error) {
    keepTemp = true;
    throw error;
  } finally {
    if (keepTemp) {
      log("warn", args.logLevel, `Keeping subtitle temp dir: ${tempDir}`);
    } else {
      try {
        fs.rmSync(tempDir, { recursive: true, force: true });
      } catch {
        // ignore cleanup failures
      }
    }
  }
}