ytdlp-segment-downloader/src/cli/downloader.ts

import { spawn } from "child_process";
import { writeFileSync, mkdirSync, existsSync } from "fs";
import { join } from "path";

export interface DownloadOptions {
  url: string;
  outputDir: string;
  format: string;
  extractChapters: boolean;
}

interface MostWatchedSegment {
  start: number;
  end: number;
  intensity?: number;
}

interface VideoInfo {
  title: string;
  chapters: Array<{ title: string; start_time: number; end_time: number }>;
  heatmap?: Array<{ start_seconds: number; end_seconds: number; intensity: number }>;
  duration: number;
}

async function getVideoInfo(url: string): Promise<VideoInfo> {
  return new Promise((resolve, reject) => {
    const ytDlp = spawn("yt-dlp", [
      "--dump-json",
      "--no-download",
      url,
    ]);

    let stdout = "";
    let stderr = "";

    ytDlp.stdout.on("data", (data) => {
      stdout += data.toString();
    });

    ytDlp.stderr.on("data", (data) => {
      stderr += data.toString();
    });

    ytDlp.on("close", (code) => {
      if (code !== 0) {
        reject(new Error(`yt-dlp failed: ${stderr}`));
        return;
      }

      try {
        const info = JSON.parse(stdout);

        resolve({
          title: info.title || "video",
          chapters: info.chapters || [],
          heatmap: info.heatmap || [],
          duration: info.duration || 0,
        });
      } catch (parseError) {
        reject(new Error(`Failed to parse video info: ${parseError}`));
      }
    });

    ytDlp.on("error", (err) => {
      reject(new Error(`Failed to run yt-dlp: ${err.message}`));
    });
  });
}

async function downloadSection(
  url: string,
  outputPath: string,
  section: string,
  format: string
): Promise<void> {
  return new Promise((resolve, reject) => {
    const ytDlp = spawn("yt-dlp", [
      "-f", format,
      "--download-sections", section,
      "-o", outputPath,
      url,
    ]);

    let stderr = "";

    ytDlp.stderr.on("data", (data) => {
      stderr += data.toString();
    });

    ytDlp.on("close", (code) => {
      if (code !== 0) {
        reject(new Error(`yt-dlp failed: ${stderr}`));
        return;
      }
      resolve();
    });

    ytDlp.on("error", (err) => {
      reject(new Error(`Failed to run yt-dlp: ${err.message}`));
    });
  });
}

function sanitizeFilename(filename: string): string {
  return filename
    .replace(/[^a-zA-Z0-9\s\-_]/g, "")
    .replace(/\s+/g, "_")
    .substring(0, 100);
}

function formatTime(seconds: number): string {
  const mins = Math.floor(seconds / 60);
  const secs = Math.floor(seconds % 60);
  return `${mins}:${secs.toString().padStart(2, "0")}`;
}

function getMostWatchedSegments(
  heatmap: Array<{ start_seconds: number; end_seconds: number; intensity: number }>,
  duration: number,
  topN: number = 1
): MostWatchedSegment[] {
  if (!heatmap || heatmap.length === 0) {
    return [];
  }

  // Sort by intensity (most watched first)
  const sorted = [...heatmap].sort((a, b) => b.intensity - a.intensity);

  // Get top N segments
  const topSegments = sorted.slice(0, topN);

  return topSegments.map((segment) => ({
    start: segment.start_seconds,
    end: segment.end_seconds,
    intensity: segment.intensity,
  }));
}

export async function downloadVideoSegments(options: DownloadOptions): Promise<void> {
  const { url, outputDir, format } = options;

  // Create output directory if it doesn't exist
  if (!existsSync(outputDir)) {
    mkdirSync(outputDir, { recursive: true });
  }

  // Get video info including most watched segments from YouTube API
  console.log("Fetching video information from YouTube API...");
  const info = await getVideoInfo(url);
  const safeTitle = sanitizeFilename(info.title);

  console.log(`Video: ${info.title}`);
  console.log(`Duration: ${formatTime(info.duration)}`);

  // Try to get most watched segments from heatmap data
  const mostWatchedSegments = getMostWatchedSegments(info.heatmap || [], info.duration);

  // If no heatmap data, fall back to chapters
  if (mostWatchedSegments.length === 0 && info.chapters.length > 0) {
    console.log("\nNo most watched segments found. Falling back to chapters...");
    const chapter = info.chapters[0]; // Download first chapter as most relevant
    const outputPath = join(outputDir, `${safeTitle}_most_watched.%(ext)s`);
    const section = `*${formatTime(chapter.start_time || 0)}-${formatTime(chapter.end_time || 60)}`;

    console.log(`Downloading chapter: ${chapter.title || "First Chapter"}`);
    await downloadSection(url, outputPath, section, format);
    return;
  }

  if (mostWatchedSegments.length === 0) {
    console.log("No segments found. Downloading full video...");
    const outputPath = join(outputDir, `${safeTitle}.%(ext)s`);
    await downloadSection(url, outputPath, "*", format);
    return;
  }

  // Download the most watched segment
  const topSegment = mostWatchedSegments[0];
  const outputPath = join(outputDir, `${safeTitle}_most_watched.%(ext)s`);
  const section = `*${formatTime(topSegment.start)}-${formatTime(topSegment.end)}`;

  console.log(`\nMost watched segment: ${formatTime(topSegment.start)} - ${formatTime(topSegment.end)}`);
  console.log(`Duration: ${formatTime(topSegment.end - topSegment.start)}`);
  console.log(`Intensity: ${((topSegment.intensity || 0) * 100).toFixed(1)}%`);
  console.log(`\nDownloading most watched segment...`);

  await downloadSection(url, outputPath, section, format);

  // Save segment info
  const segmentInfoPath = join(outputDir, `${safeTitle}_most_watched.txt`);
  let segmentInfo = `# ${info.title}\n\n`;
  segmentInfo += `Most watched segment:\n`;
  segmentInfo += `  Start: ${formatTime(topSegment.start)}\n`;
  segmentInfo += `  End: ${formatTime(topSegment.end)}\n`;
  segmentInfo += `  Duration: ${formatTime(topSegment.end - topSegment.start)}\n`;
  segmentInfo += `  Intensity: ${((topSegment.intensity || 0) * 100).toFixed(1)}%\n`;

  if (mostWatchedSegments.length > 1) {
    segmentInfo += `\nOther top segments:\n`;
    for (let i = 1; i < mostWatchedSegments.length; i++) {
      const seg = mostWatchedSegments[i];
      segmentInfo += `  ${formatTime(seg.start)} - ${formatTime(seg.end)} (${((seg.intensity || 0) * 100).toFixed(1)}%)\n`;
    }
  }

  writeFileSync(segmentInfoPath, segmentInfo);
  console.log(`Segment info saved to: ${segmentInfoPath}`);
}