Tdarr-plugin/LocalFlowPlugins/subtitles/extractSubtitlesAndOcr/1.0.0/index.js

'use strict';

const path = require('path');
const fs   = require('fs');

/* eslint-disable no-param-reassign */

/**
 * extractSubtitlesAndOcr
 * Extracts subtitle streams for ITA and ENG languages:
 *   - Text-based (subrip, srt, ass, ssa, webvtt): direct FFmpeg extraction
 *   - Bitmap-based (hdmv_pgs_subtitle, dvd_subtitle): FFmpeg .sup extract
 *     then SubtitleEdit CLI OCR -> .srt (with graceful fallback on failure)
 *
 * Output files go next to the source file.
 * Naming: <baseName>.<lang>[.forced][.default][.sdh].srt
 */

const TEXT_CODECS    = new Set(['subrip', 'srt', 'ass', 'ssa', 'webvtt', 'mov_text']);
const BITMAP_CODECS  = new Set(['hdmv_pgs_subtitle', 'dvd_subtitle', 'pgssub', 'dvbsub']);

const details = () => ({
  name: 'Extract Subtitles & OCR Fallback',
  description:
    'Extracts ITA/ENG subtitle streams to SRT files next to the source. '
    + 'Text subtitles are extracted directly via FFmpeg. '
    + 'Bitmap subtitles (PGS/VOBSUB) are OCR-d via SubtitleEdit CLI. '
    + 'OCR failures are logged as warnings without stopping the pipeline.',
  style: { borderColor: '#00BCD4' },
  tags: 'subtitles',
  isStartPlugin: false,
  pType: '',
  requiresVersion: '2.11.01',
  sidebarPosition: 3,
  icon: 'faClosedCaptioning',
  inputs: [
    {
      label: 'SubtitleEdit Executable Path',
      name: 'subtitleEditPath',
      type: 'string',
      defaultValue: 'SubtitleEdit',
      inputUI: { type: 'text' },
      tooltip:
        'Name or full path of the SubtitleEdit CLI executable. '
        + 'Must be in system PATH or provide absolute path. '
        + 'Used for OCR of PGS/VOBSUB streams on Windows workers.',
    },
    {
      label: 'Languages to Extract (CSV)',
      name: 'languages',
      type: 'string',
      defaultValue: 'ita,eng',
      inputUI: { type: 'text' },
      tooltip: 'Comma-separated list of ISO 639-2 language codes. Example: ita,eng',
    },
  ],
  outputs: [
    { number: 1, tooltip: 'Continue to next plugin (warnings in log if OCR failed)' },
  ],
});

// ─── Helpers ──────────────────────────────────────────────────────────────────

const normLang = (tag) => {
  const l = (tag || '').toLowerCase().trim();
  if (l === 'it') return 'ita';
  if (l === 'en') return 'eng';
  return l;
};

const buildSuffix = (lang, stream) => {
  const parts = [lang];
  if (stream.disposition && stream.disposition.forced)           parts.push('forced');
  if (stream.disposition && stream.disposition.default)          parts.push('default');
  if (stream.disposition && stream.disposition.hearing_impaired) parts.push('sdh');
  // Also check title tag for SDH
  const title = ((stream.tags && stream.tags.title) || '').toLowerCase();
  if (!parts.includes('sdh') && title.includes('sdh'))           parts.push('sdh');
  return parts.join('.');
};

const runFFmpeg = async (spawnArgs, workDir, ffmpegPath, jobLog, inputFileObj) => {
  const { CLI } = require('../../../../FlowHelpers/1.0.0/cliUtils');
  const cli = new CLI({
    cli: ffmpegPath, spawnArgs, WorkDir: workDir, jobLog, inputFileObj,
  });
  return cli.runCli();
};

// ─── Plugin ───────────────────────────────────────────────────────────────────

const plugin = async (args) => {
  const { jobLog, inputFileObj, variables, workDir } = args;

  const inputFile = inputFileObj._id;
  const sourceDir = path.dirname(inputFile);
  const baseName  = path.basename(inputFile, path.extname(inputFile));

  jobLog(`[extractSubtitlesAndOcr] Source: ${inputFile}\n`);

  const subtitleEditPath = args.inputs.subtitleEditPath || 'SubtitleEdit';
  const targetLangs = (args.inputs.languages || 'ita,eng')
    .split(',').map((l) => l.trim().toLowerCase()).filter(Boolean);

  const streams = (inputFileObj.ffProbeData && inputFileObj.ffProbeData.streams) || [];
  const subtitleStreams = streams.filter((s) => s.codec_type === 'subtitle');

  if (subtitleStreams.length === 0) {
    jobLog('[extractSubtitlesAndOcr] No subtitle streams found.\n');
    return { outputNumber: 1, outputFileObj: inputFileObj, variables };
  }

  jobLog(`[extractSubtitlesAndOcr] Found ${subtitleStreams.length} subtitle stream(s).\n`);

  for (const stream of subtitleStreams) {
    const lang = normLang((stream.tags && stream.tags.language) || '');
    if (!targetLangs.includes(lang)) {
      jobLog(`[extractSubtitlesAndOcr] Stream #${stream.index} lang="${lang}" — not in target list, skipping.\n`);
      continue;
    }

    const codec  = (stream.codec_name || '').toLowerCase();
    const suffix = buildSuffix(lang, stream);
    const outSrt = path.join(sourceDir, `${baseName}.${suffix}.srt`);

    jobLog(`[extractSubtitlesAndOcr] Stream #${stream.index} lang=${lang} codec=${codec} -> ${path.basename(outSrt)}\n`);

    // ── Text-based subtitles ──────────────────────────────────────────────
    if (TEXT_CODECS.has(codec)) {
      const { cliExitCode } = await runFFmpeg(
        ['-y', '-i', inputFile, '-map', `0:${stream.index}`, '-c:s', 'srt', outSrt],
        workDir, args.ffmpegPath, jobLog, inputFileObj,
      );
      if (cliExitCode !== 0) {
        jobLog(`[extractSubtitlesAndOcr] WARNING: FFmpeg failed for stream #${stream.index} (exit ${cliExitCode}).\n`);
      } else {
        jobLog(`[extractSubtitlesAndOcr] Extracted: ${path.basename(outSrt)}\n`);
      }
      continue;
    }

    // ── Bitmap-based subtitles (PGS / VOBSUB) ────────────────────────────
    if (BITMAP_CODECS.has(codec)) {
      const tmpSup = path.join(workDir, `${baseName}_s${stream.index}.sup`);

      // Step 1: Extract to .sup
      const { cliExitCode: extractExit } = await runFFmpeg(
        ['-y', '-i', inputFile, '-map', `0:${stream.index}`, '-c:s', 'copy', tmpSup],
        workDir, args.ffmpegPath, jobLog, inputFileObj,
      );

      if (extractExit !== 0) {
        jobLog(`[extractSubtitlesAndOcr] WARNING: .sup extraction failed for stream #${stream.index}.\n`);
        continue;
      }

      // Step 2: OCR with SubtitleEdit
      try {
        const { CLI } = require('../../../../FlowHelpers/1.0.0/cliUtils');
        const seCli = new CLI({
          cli: subtitleEditPath,
          spawnArgs: [
            '/convert', tmpSup, 'subrip',
            '/track-number:0',
            `/outputfolder:${sourceDir}`,
          ],
          WorkDir: workDir,
          jobLog,
          inputFileObj,
        });
        const { cliExitCode: seExit } = await seCli.runCli();

        if (seExit !== 0) {
          jobLog(`[extractSubtitlesAndOcr] WARNING: SubtitleEdit OCR failed (exit ${seExit}) for stream #${stream.index}.\n`);
        } else {
          // SubtitleEdit names the output after the .sup filename
          const seOut = path.join(sourceDir, path.basename(tmpSup, '.sup') + '.srt');
          if (fs.existsSync(seOut)) {
            fs.renameSync(seOut, outSrt);
            jobLog(`[extractSubtitlesAndOcr] OCR complete: ${path.basename(outSrt)}\n`);
          }
        }
      } catch (err) {
        jobLog(`[extractSubtitlesAndOcr] WARNING: SubtitleEdit not available or crashed: ${err.message}\n`);
      } finally {
        // Clean up temp .sup
        try { fs.unlinkSync(tmpSup); } catch (_) { /* ignore */ }
      }
      continue;
    }

    jobLog(`[extractSubtitlesAndOcr] Stream #${stream.index} codec "${codec}" not supported — skipping.\n`);
  }

  jobLog('[extractSubtitlesAndOcr] Done.\n');
  return { outputNumber: 1, outputFileObj: inputFileObj, variables };
};

module.exports.details = details;
module.exports.plugin = plugin;