douyin-archive/lib/asr-batch-jobs.ts

import { randomUUID } from "node:crypto";
import { transcriptAweme } from "@/app/api/asr";
import { resolveEffectiveContentTags } from "@/lib/asr-tags";
import { prisma } from "@/lib/prisma";

export const ASR_BATCH_SCOPES = [
  "missingTags",
  "pending",
  "transcribed",
  "tagged",
  "all",
  "selected",
] as const;

export type AsrBatchScope = (typeof ASR_BATCH_SCOPES)[number];
export type AsrBatchJobStatus =
  | "queued"
  | "running"
  | "paused"
  | "completed"
  | "canceled"
  | "failed";
export type AsrBatchItemStatus =
  | "pending"
  | "running"
  | "success"
  | "failed"
  | "canceled";
export type AsrBatchControlAction = "pause" | "resume" | "cancel";

export type CreateAsrBatchJobInput = {
  scope?: unknown;
  rangeStart?: unknown;
  rangeEnd?: unknown;
  limit?: unknown;
  awemeIds?: unknown;
};

export type AsrBatchJobItemSnapshot = {
  index: number;
  awemeId: string;
  desc: string;
  status: AsrBatchItemStatus;
  tags: string[];
  error: string | null;
  startedAt: string | null;
  finishedAt: string | null;
  durationMs: number | null;
};

export type AsrBatchJobSnapshot = {
  id: string;
  status: AsrBatchJobStatus;
  scope: AsrBatchScope;
  selection: {
    sourceCount: number;
    matchedCount: number;
    selectedCount: number;
    rangeStart: number;
    rangeEnd: number;
    requestedIdsCount: number;
  };
  progress: {
    total: number;
    completed: number;
    failed: number;
    canceled: number;
    pending: number;
    running: number;
    processed: number;
    percent: number;
  };
  currentItem: AsrBatchJobItemSnapshot | null;
  pauseRequested: boolean;
  cancelRequested: boolean;
  error: string | null;
  logs: { at: string; message: string; awemeId: string | null }[];
  createdAt: string;
  startedAt: string | null;
  updatedAt: string;
  finishedAt: string | null;
  items?: AsrBatchJobItemSnapshot[];
};

type CandidateVideo = {
  awemeId: string;
  desc: string;
  createdAt: string;
  hasTranscript: boolean;
  contentTags: string[];
};

type InternalAsrBatchJob = Omit<AsrBatchJobSnapshot, "progress" | "items"> & {
  items: AsrBatchJobItemSnapshot[];
  runner?: Promise<void>;
};

type AsrBatchStore = {
  jobs: Map<string, InternalAsrBatchJob>;
};

const MAX_SOURCE_VIDEOS = 5000;
const MAX_JOB_ITEMS = 1000;
const MAX_JOB_HISTORY = 20;
const MAX_LOGS = 160;

const globalForAsrBatchJobs = globalThis as typeof globalThis & {
  __asrBatchJobs?: AsrBatchStore;
};

if (!globalForAsrBatchJobs.__asrBatchJobs) {
  globalForAsrBatchJobs.__asrBatchJobs = { jobs: new Map() };
}

const store = globalForAsrBatchJobs.__asrBatchJobs;

export class AsrBatchJobError extends Error {
  status: number;

  constructor(message: string, status = 400) {
    super(message);
    this.name = "AsrBatchJobError";
    this.status = status;
  }
}

export function isTerminalAsrBatchStatus(status: AsrBatchJobStatus) {
  return status === "completed" || status === "canceled" || status === "failed";
}

function nowIso() {
  return new Date().toISOString();
}

function normalizeScope(scope: unknown): AsrBatchScope {
  return typeof scope === "string" &&
    ASR_BATCH_SCOPES.includes(scope as AsrBatchScope)
    ? (scope as AsrBatchScope)
    : "missingTags";
}

function normalizePositiveInt(value: unknown) {
  const numberValue =
    typeof value === "number"
      ? value
      : typeof value === "string" && value.trim().length > 0
        ? Number.parseInt(value, 10)
        : Number.NaN;

  return Number.isFinite(numberValue) && numberValue > 0
    ? Math.floor(numberValue)
    : null;
}

function normalizeAwemeIds(value: unknown) {
  if (!Array.isArray(value)) return [];

  return [...new Set(value)]
    .filter(
      (id): id is string => typeof id === "string" && id.trim().length > 0,
    )
    .map((id) => id.trim())
    .slice(0, MAX_JOB_ITEMS);
}

function appendJobLog(
  job: InternalAsrBatchJob,
  message: string,
  awemeId: string | null = null,
) {
  job.logs.push({ at: nowIso(), message, awemeId });
  if (job.logs.length > MAX_LOGS) {
    job.logs.splice(0, job.logs.length - MAX_LOGS);
  }
  job.updatedAt = nowIso();
}

function calculateProgress(job: InternalAsrBatchJob) {
  const total = job.items.length;
  const completed = job.items.filter(
    (item) => item.status === "success",
  ).length;
  const failed = job.items.filter((item) => item.status === "failed").length;
  const canceled = job.items.filter(
    (item) => item.status === "canceled",
  ).length;
  const running = job.items.filter((item) => item.status === "running").length;
  const pending = job.items.filter((item) => item.status === "pending").length;
  const processed = completed + failed + canceled;

  return {
    total,
    completed,
    failed,
    canceled,
    pending,
    running,
    processed,
    percent: total > 0 ? Math.round((processed / total) * 100) : 0,
  };
}

function serializeJob(
  job: InternalAsrBatchJob,
  includeItems = false,
): AsrBatchJobSnapshot {
  const currentItem =
    job.items.find((item) => item.status === "running") ?? null;
  const snapshot: AsrBatchJobSnapshot = {
    id: job.id,
    status: job.status,
    scope: job.scope,
    selection: job.selection,
    progress: calculateProgress(job),
    currentItem,
    pauseRequested: job.pauseRequested,
    cancelRequested: job.cancelRequested,
    error: job.error,
    logs: job.logs,
    createdAt: job.createdAt,
    startedAt: job.startedAt,
    updatedAt: job.updatedAt,
    finishedAt: job.finishedAt,
  };

  if (includeItems) snapshot.items = job.items;
  return snapshot;
}

function trimJobHistory() {
  const jobs = [...store.jobs.values()].sort(
    (left, right) =>
      new Date(right.createdAt).getTime() - new Date(left.createdAt).getTime(),
  );

  for (const job of jobs.slice(MAX_JOB_HISTORY)) {
    if (isTerminalAsrBatchStatus(job.status)) store.jobs.delete(job.id);
  }
}

function getActiveJob() {
  return [...store.jobs.values()].find(
    (job) => !isTerminalAsrBatchStatus(job.status),
  );
}

async function readCandidateVideos(awemeIds: string[]) {
  const videos = await prisma.video.findMany({
    ...(awemeIds.length > 0 ? { where: { aweme_id: { in: awemeIds } } } : {}),
    orderBy: { created_at: "desc" },
    take: awemeIds.length > 0 ? awemeIds.length : MAX_SOURCE_VIDEOS,
    select: {
      aweme_id: true,
      desc: true,
      created_at: true,
      manual_content_tags: true,
      transcript: {
        select: {
          id: true,
          content_tags: true,
        },
      },
    },
  });

  const orderMap = new Map(awemeIds.map((awemeId, index) => [awemeId, index]));
  const sortedVideos =
    awemeIds.length > 0
      ? [...videos].sort(
          (left, right) =>
            (orderMap.get(left.aweme_id) ?? 0) -
            (orderMap.get(right.aweme_id) ?? 0),
        )
      : videos;

  return sortedVideos.map<CandidateVideo>((video) => ({
    awemeId: video.aweme_id,
    desc: video.desc,
    createdAt: video.created_at.toISOString(),
    hasTranscript: Boolean(video.transcript?.id),
    contentTags: resolveEffectiveContentTags(
      video.manual_content_tags,
      video.transcript?.content_tags,
    ),
  }));
}

function filterCandidatesByScope(
  candidates: CandidateVideo[],
  scope: AsrBatchScope,
) {
  if (scope === "selected" || scope === "all") return candidates;

  return candidates.filter((candidate) => {
    const hasTags = candidate.contentTags.length > 0;
    if (scope === "missingTags") return !hasTags;
    if (scope === "pending") return !candidate.hasTranscript;
    if (scope === "transcribed") return candidate.hasTranscript;
    if (scope === "tagged") return hasTags;
    return true;
  });
}

function resolveRange(matchedCount: number, input: CreateAsrBatchJobInput) {
  const rangeStart = normalizePositiveInt(input.rangeStart) ?? 1;
  const explicitRangeEnd = normalizePositiveInt(input.rangeEnd);
  const limit = normalizePositiveInt(input.limit);
  const requestedRangeEnd =
    explicitRangeEnd ?? (limit ? rangeStart + limit - 1 : matchedCount);
  const rangeEnd = Math.min(
    Math.max(requestedRangeEnd, rangeStart),
    matchedCount,
  );

  return { rangeStart, rangeEnd };
}

async function selectJobItems(input: CreateAsrBatchJobInput) {
  const awemeIds = normalizeAwemeIds(input.awemeIds);
  const scope = awemeIds.length > 0 ? "selected" : normalizeScope(input.scope);
  const sourceCandidates = await readCandidateVideos(awemeIds);
  const matchedCandidates = filterCandidatesByScope(sourceCandidates, scope);
  const { rangeStart, rangeEnd } = resolveRange(
    matchedCandidates.length,
    input,
  );
  const selectedCandidates = matchedCandidates
    .slice(rangeStart - 1, rangeEnd)
    .slice(0, MAX_JOB_ITEMS);

  return {
    scope,
    rangeStart,
    rangeEnd,
    sourceCount: sourceCandidates.length,
    matchedCount: matchedCandidates.length,
    requestedIdsCount: awemeIds.length,
    selectedCandidates,
  };
}

function launchRunner(job: InternalAsrBatchJob) {
  if (job.runner) return;

  const runner = runAsrBatchJob(job.id).finally(() => {
    const latestJob = store.jobs.get(job.id);
    if (latestJob?.runner === runner) latestJob.runner = undefined;
  });
  job.runner = runner;
}

async function runAsrBatchJob(jobId: string) {
  const job = store.jobs.get(jobId);
  if (!job) return;

  job.status = "running";
  job.startedAt ??= nowIso();
  job.updatedAt = nowIso();
  appendJobLog(job, "任务开始执行");

  try {
    for (;;) {
      if (job.cancelRequested) {
        finalizeCanceledJob(job, "任务已取消");
        return;
      }

      if (job.pauseRequested) {
        job.status = "paused";
        job.updatedAt = nowIso();
        appendJobLog(job, "任务已暂停");
        return;
      }

      const nextItem = job.items.find((item) => item.status === "pending");
      if (!nextItem) {
        job.status = "completed";
        job.finishedAt = nowIso();
        job.updatedAt = nowIso();
        appendJobLog(job, "任务已完成");
        return;
      }

      await runJobItem(job, nextItem);
    }
  } catch (error) {
    job.status = "failed";
    job.error = error instanceof Error ? error.message : "任务执行失败";
    job.finishedAt = nowIso();
    job.updatedAt = nowIso();
    appendJobLog(job, job.error);
  }
}

async function runJobItem(
  job: InternalAsrBatchJob,
  item: AsrBatchJobItemSnapshot,
) {
  const startedAt = Date.now();
  item.status = "running";
  item.startedAt = nowIso();
  item.error = null;
  job.updatedAt = nowIso();
  appendJobLog(
    job,
    `开始处理第 ${item.index}/${job.items.length} 个视频`,
    item.awemeId,
  );

  try {
    const result = await transcriptAweme(item.awemeId, {
      log: (message) => appendJobLog(job, message, item.awemeId),
    });
    item.status = "success";
    item.tags = result.content_tags;
    appendJobLog(
      job,
      `处理完成: ${result.content_tags.join("、")}`,
      item.awemeId,
    );
  } catch (error) {
    item.status = "failed";
    item.error = error instanceof Error ? error.message : "处理失败";
    appendJobLog(job, `处理失败: ${item.error}`, item.awemeId);
  } finally {
    item.finishedAt = nowIso();
    item.durationMs = Date.now() - startedAt;
    job.updatedAt = nowIso();
  }
}

function finalizeCanceledJob(job: InternalAsrBatchJob, message: string) {
  for (const item of job.items) {
    if (item.status === "pending") {
      item.status = "canceled";
      item.finishedAt = nowIso();
    }
  }
  job.status = "canceled";
  job.finishedAt = nowIso();
  job.updatedAt = nowIso();
  appendJobLog(job, message);
}

export async function createAsrBatchJob(input: CreateAsrBatchJobInput) {
  trimJobHistory();

  const activeJob = getActiveJob();
  if (activeJob) {
    throw new AsrBatchJobError("已有批量任务正在执行或暂停中", 409);
  }

  const selection = await selectJobItems(input);
  if (selection.selectedCandidates.length === 0) {
    throw new AsrBatchJobError("没有匹配的视频可处理", 400);
  }

  const createdAt = nowIso();
  const job: InternalAsrBatchJob = {
    id: randomUUID(),
    status: "queued",
    scope: selection.scope,
    selection: {
      sourceCount: selection.sourceCount,
      matchedCount: selection.matchedCount,
      selectedCount: selection.selectedCandidates.length,
      rangeStart: selection.rangeStart,
      rangeEnd: selection.rangeEnd,
      requestedIdsCount: selection.requestedIdsCount,
    },
    currentItem: null,
    pauseRequested: false,
    cancelRequested: false,
    error: null,
    logs: [],
    createdAt,
    startedAt: null,
    updatedAt: createdAt,
    finishedAt: null,
    items: selection.selectedCandidates.map((candidate, index) => ({
      index: index + 1,
      awemeId: candidate.awemeId,
      desc: candidate.desc,
      status: "pending",
      tags: [],
      error: null,
      startedAt: null,
      finishedAt: null,
      durationMs: null,
    })),
  };

  store.jobs.set(job.id, job);
  appendJobLog(
    job,
    `任务已创建，匹配 ${selection.matchedCount} 个，选中 ${selection.selectedCandidates.length} 个`,
  );
  launchRunner(job);

  return serializeJob(job, true);
}

export function listAsrBatchJobs() {
  return [...store.jobs.values()]
    .sort(
      (left, right) =>
        new Date(right.createdAt).getTime() -
        new Date(left.createdAt).getTime(),
    )
    .map((job) => serializeJob(job));
}

export function getActiveAsrBatchJob() {
  const activeJob = getActiveJob();
  return activeJob ? serializeJob(activeJob) : null;
}

export function getAsrBatchJob(jobId: string, includeItems = true) {
  const job = store.jobs.get(jobId);
  return job ? serializeJob(job, includeItems) : null;
}

export function controlAsrBatchJob(jobId: string, action: unknown) {
  const job = store.jobs.get(jobId);
  if (!job) throw new AsrBatchJobError("任务不存在", 404);

  if (action === "pause") {
    if (isTerminalAsrBatchStatus(job.status)) return serializeJob(job, true);
    job.pauseRequested = true;
    job.cancelRequested = false;
    if (!job.runner && job.status !== "paused") job.status = "paused";
    appendJobLog(job, "已请求暂停，当前视频处理完后暂停");
    return serializeJob(job, true);
  }

  if (action === "resume" || action === "continue") {
    if (isTerminalAsrBatchStatus(job.status)) {
      throw new AsrBatchJobError("已结束的任务不能继续", 409);
    }
    if (job.cancelRequested) {
      throw new AsrBatchJobError("已请求取消的任务不能继续", 409);
    }
    job.pauseRequested = false;
    job.status = "running";
    appendJobLog(job, "任务继续执行");
    launchRunner(job);
    return serializeJob(job, true);
  }

  if (action === "cancel") {
    if (isTerminalAsrBatchStatus(job.status)) return serializeJob(job, true);
    job.cancelRequested = true;
    job.pauseRequested = false;
    if (!job.runner || job.status === "paused" || job.status === "queued") {
      finalizeCanceledJob(job, "任务已取消");
    } else {
      appendJobLog(job, "已请求取消，当前视频处理完后停止");
    }
    return serializeJob(job, true);
  }

  throw new AsrBatchJobError("不支持的任务操作", 400);
}