douyin-archive/app/api/media.ts

import { execFile } from "child_process";
import { promises as fs } from "fs";
import os from "os";
import path from "path";
import { promisify } from "util";

const execFileAsync = promisify(execFile);

/**
 * 使用 ffmpeg 从视频二进制中提取第一帧，返回 JPEG buffer
 */
export async function extractFirstFrame(
  videoBuffer: Buffer,
): Promise<{ buffer: Buffer; contentType: string; ext: string } | null> {
  const ffmpegCmd = process.env.FFMPEG_PATH || "ffmpeg";
  const tmpDir = os.tmpdir();
  const base = `dy_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;
  const inPath = path.join(tmpDir, `${base}.mp4`);
  const outPath = path.join(tmpDir, `${base}.jpg`);

  try {
    await fs.writeFile(inPath, videoBuffer);
    const args = [
      "-hide_banner",
      "-loglevel",
      "error",
      "-ss",
      "0",
      "-i",
      inPath,
      "-frames:v",
      "1",
      "-q:v",
      "2",
      "-f",
      "image2",
      "-y",
      outPath,
    ];
    await execFileAsync(ffmpegCmd, args, { windowsHide: true });
    const img = await fs.readFile(outPath);
    return { buffer: img, contentType: "image/jpeg", ext: "jpg" };
  } catch (e: any) {
    if (
      e &&
      (e.code === "ENOENT" ||
        /not found|is not recognized/i.test(String(e.message)))
    ) {
      console.warn(
        "系统未检测到 ffmpeg，可安装并配置 PATH 或设置 FFMPEG_PATH 后启用封面提取。",
      );
      return null;
    }
    throw e;
  } finally {
    try {
      await fs.unlink(inPath);
    } catch {}
    try {
      await fs.unlink(outPath);
    } catch {}
  }
}

/**
 * 使用 ffprobe 获取视频时长（毫秒）
 */
export async function getVideoDuration(
  videoBuffer: Buffer,
): Promise<number | null> {
  const ffprobeCmd = process.env.FFPROBE_PATH || "ffprobe";
  const tmpDir = os.tmpdir();
  const base = `dy_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;
  const inPath = path.join(tmpDir, `${base}.mp4`);

  try {
    await fs.writeFile(inPath, videoBuffer);
    const args = [
      "-v",
      "error",
      "-show_entries",
      "format=duration",
      "-of",
      "default=noprint_wrappers=1:nokey=1",
      inPath,
    ];
    const { stdout } = await execFileAsync(ffprobeCmd, args, {
      windowsHide: true,
    });
    const durationSeconds = parseFloat(stdout.trim());
    if (isNaN(durationSeconds)) return null;
    return Math.round(durationSeconds * 1000); // 转换为毫秒
  } catch (e: any) {
    if (
      e &&
      (e.code === "ENOENT" ||
        /not found|is not recognized/i.test(String(e.message)))
    ) {
      console.warn(
        "系统未检测到 ffprobe，可安装并配置 PATH 或设置 FFPROBE_PATH 后启用时长提取。",
      );
      return null;
    }
    console.warn(`获取视频时长失败: ${e?.message || e}`);
    return null;
  } finally {
    try {
      await fs.unlink(inPath);
    } catch {}
  }
}

/**
 * 使用 ffmpeg 将视频 Buffer 或本地文件转换为音频，返回音频 Buffer
 *
 * 默认输出 mp3，可通过 opts.format 指定 'mp3' | 'aac' | 'wav'
 * - mp3: 使用 libmp3lame，16 kHz 单声道，contentType 为 'audio/mpeg'
 * - aac: 输出 m4a 容器（AAC 编码），contentType 为 'audio/mp4'
 * - wav: PCM 16-bit，44100Hz，双声道，contentType 为 'audio/wav'
 *
 * 若系统未安装 ffmpeg 或未在 PATH 中，返回 null 并给出提示。
 */
export async function extractAudio(
  videoBuffer: Buffer | string,
  opts?: { format?: "mp3" | "aac" | "wav"; bitrateKbps?: number },
): Promise<{ buffer: Buffer; contentType: string; ext: string } | null> {
  const ffmpegCmd = process.env.FFMPEG_PATH || "ffmpeg";
  const tmpDir = os.tmpdir();
  const base = `dy_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;
  const inPath = path.join(tmpDir, `${base}.mp4`);

  const inputPath = typeof videoBuffer === "string" ? videoBuffer : inPath;
  const format = opts?.format ?? "mp3";
  const bitrate = Math.max(32, Math.min(512, opts?.bitrateKbps ?? 192)); // 安全范围 32~512 kbps

  // 根据目标格式设置输出路径、MIME 与编码参数
  let outPath = "";
  let contentType = "";
  let ext = "";
  let codecArgs: string[] = [];

  if (format === "mp3") {
    ext = "mp3";
    contentType = "audio/mpeg";
    outPath = path.join(tmpDir, `${base}.${ext}`);
    codecArgs = ["-c:a", "libmp3lame", "-b:a", `${bitrate}k`];
  } else if (format === "aac") {
    // 使用 m4a 容器更通用
    ext = "m4a";
    contentType = "audio/mp4";
    outPath = path.join(tmpDir, `${base}.${ext}`);
    codecArgs = [
      "-c:a",
      "aac",
      "-b:a",
      `${bitrate}k`,
      "-movflags",
      "+faststart",
    ];
  } else {
    // wav
    ext = "wav";
    contentType = "audio/wav";
    outPath = path.join(tmpDir, `${base}.${ext}`);
    codecArgs = [
      "-f",
      "wav",
      "-acodec",
      "pcm_s16le",
      "-ar",
      "44100",
      "-ac",
      "2",
    ];
  }

  try {
    if (typeof videoBuffer !== "string") await fs.writeFile(inPath, videoBuffer);
    const args = [
      "-hide_banner",
      "-loglevel",
      "error",
      "-i",
      inputPath,
      "-map",
      "0:a:0",
      "-ac",
      "1",
      "-ar",
      "16000",
      "-vn", // 丢弃视频流
      ...codecArgs,
      "-y",
      outPath,
    ];
    await execFileAsync(ffmpegCmd, args, {
      windowsHide: true,
      timeout: 60 * 60 * 1000,
    });
    const audio = await fs.readFile(outPath);
    return { buffer: audio, contentType, ext };
  } catch (e: any) {
    if (
      e &&
      (e.code === "ENOENT" ||
        /not found|is not recognized/i.test(String(e.message)))
    ) {
      console.warn(
        "系统未检测到 ffmpeg，可安装并配置 PATH 或设置 FFMPEG_PATH 后启用音频提取。",
      );
      return null;
    }
    // 一些环境可能缺少特定编码器（如 libmp3lame），提示并抛出原始错误
    console.warn(`提取音频失败: ${e?.message || e}`);
    throw e;
  } finally {
    try {
      await fs.unlink(inPath);
    } catch {}
    try {
      if (outPath) await fs.unlink(outPath);
    } catch {}
  }
}