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;
  peakThreshold: number;
}

interface RawHeatmapSegment {
  start_seconds?: number;
  start_time?: number;
  end_seconds?: number;
  end_time?: number;
  intensity?: number;
  heat?: number;
  value?: number;
}

interface ProcessedSegment {
  start: number;
  end: number;
  intensity: number;
  peakScore: number;
}

interface VideoInfo {
  title: string;
  duration: number;
  heatmap?: RawHeatmapSegment[];
}

async function getVideoInfo(url: string): Promise<VideoInfo> {
  return new Promise((resolve, reject) => {
    const ytDlp = spawn("yt-dlp", [
      "--dump-json",
      "--no-download",
      "--compat-option",
      "no-youtube-channel-redirect",
      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",
          duration: info.duration || 0,
          heatmap: info.heatmap,
        });
      } 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 downloadSegment(
  url: string,
  outputPath: string,
  startTime: number,
  endTime: number,
  format: string
): Promise<void> {
  const section = `*${startTime.toFixed(3)}-${endTime.toFixed(3)}`;

  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 {
  if (!Number.isFinite(seconds) || seconds < 0) {
    return "0:00";
  }
  const mins = Math.floor(seconds / 60);
  const secs = Math.floor(seconds % 60);
  return `${mins}:${secs.toString().padStart(2, "0")}`;
}

function getStartTime(segment: RawHeatmapSegment): number {
  return segment.start_seconds ?? segment.start_time ?? 0;
}

function getEndTime(segment: RawHeatmapSegment): number {
  return segment.end_seconds ?? segment.end_time ?? 0;
}

function getIntensity(segment: RawHeatmapSegment): number {
  return segment.intensity ?? segment.heat ?? segment.value ?? 0;
}

function findPeakSegments(
  segments: RawHeatmapSegment[],
  threshold: number = 0.3
): ProcessedSegment[] {
  if (segments.length < 3) {
    return [];
  }

  // Convert to processed format
  const processed = segments
    .map(seg => ({
      start: getStartTime(seg),
      end: getEndTime(seg),
      intensity: getIntensity(seg),
      peakScore: 0,
    }))
    .filter(seg => 
      Number.isFinite(seg.start) && 
      Number.isFinite(seg.end) && 
      Number.isFinite(seg.intensity)
    );

  if (processed.length < 3) {
    return [];
  }

  // Calculate peak score for each segment
  // A peak is where intensity is significantly higher than neighbors
  const scored = processed.map((seg, i) => {
    const prevIntensity = i > 0 ? processed[i - 1].intensity : seg.intensity;
    const nextIntensity = i < processed.length - 1 ? processed[i + 1].intensity : seg.intensity;
    
    // Peak score = how much higher this segment is compared to average of neighbors
    const avgNeighborIntensity = (prevIntensity + nextIntensity) / 2;
    const peakScore = avgNeighborIntensity > 0 
      ? (seg.intensity - avgNeighborIntensity) / avgNeighborIntensity 
      : 0;
    
    return {
      ...seg,
      peakScore,
    };
  });

  // Filter segments that are true peaks (higher than neighbors)
  const peaks = scored.filter(seg => seg.peakScore > threshold);

  // Sort by peak score (highest peaks first)
  peaks.sort((a, b) => b.peakScore - a.peakScore);

  return peaks;
}

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

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

  // Get video info with heatmap data from YouTube
  console.log("Fetching video information from YouTube...");
  const info = await getVideoInfo(url);
  const safeTitle = sanitizeFilename(info.title);

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

  // Check for heatmap data
  if (!info.heatmap || info.heatmap.length === 0) {
    console.log("\nNo heatmap data available for this video.");
    console.log("Downloading full video instead...");
    
    const outputPath = join(outputDir, `${safeTitle}.%(ext)s`);
    await downloadSegment(url, outputPath, 0, info.duration, format);
    return;
  }

  console.log(`\nHeatmap data found: ${info.heatmap.length} segments`);

  // Find peak segments (segments that stand out from their neighbors)
  const peakSegments = findPeakSegments(info.heatmap, peakThreshold);

  if (peakSegments.length === 0) {
    console.log("No significant peak segments found.");
    console.log("Downloading full video...");
    const outputPath = join(outputDir, `${safeTitle}.%(ext)s`);
    await downloadSegment(url, outputPath, 0, info.duration, format);
    return;
  }

  // Get the top peak segment
  const topPeak = peakSegments[0];

  console.log(`\nTop peak segment:`);
  console.log(`  Time: ${formatTime(topPeak.start)} - ${formatTime(topPeak.end)}`);
  console.log(`  Duration: ${formatTime(topPeak.end - topPeak.start)}`);
  console.log(`  Peak Score: ${(topPeak.peakScore * 100).toFixed(1)}%`);
  console.log(`  Base Intensity: ${(topPeak.intensity * 100).toFixed(1)}%`);

  // Download the peak segment
  const outputPath = join(outputDir, `${safeTitle}_peak.%(ext)s`);

  console.log(`\nDownloading peak segment...`);
  await downloadSegment(url, outputPath, topPeak.start, topPeak.end, format);

  // Save segment info
  const segmentInfoPath = join(outputDir, `${safeTitle}_peak_info.txt`);
  let segmentInfo = `# ${info.title}\n\n`;
  segmentInfo += `Peak segment (stands out from surrounding content):\n`;
  segmentInfo += `  Start: ${formatTime(topPeak.start)} (${topPeak.start.toFixed(1)}s)\n`;
  segmentInfo += `  End: ${formatTime(topPeak.end)} (${topPeak.end.toFixed(1)}s)\n`;
  segmentInfo += `  Duration: ${formatTime(topPeak.end - topPeak.start)}\n`;
  segmentInfo += `  Peak Score: ${(topPeak.peakScore * 100).toFixed(1)}%\n`;
  segmentInfo += `  Intensity: ${(topPeak.intensity * 100).toFixed(1)}%\n\n`;

  if (peakSegments.length > 1) {
    segmentInfo += `Other peaks:\n`;
    for (let i = 1; i < Math.min(peakSegments.length, 5); i++) {
      const seg = peakSegments[i];
      segmentInfo += `  ${formatTime(seg.start)} - ${formatTime(seg.end)} (score: ${(seg.peakScore * 100).toFixed(1)}%)\n`;
    }
  }

  writeFileSync(segmentInfoPath, segmentInfo);
  console.log(`\nSegment info saved to: ${segmentInfoPath}`);
  console.log("Download complete!");
}