audio_codebuddy/server/src/modules/book-generator/book-generator.service.ts

/**
 * 书籍生成编排服务
 * 负责按顺序执行生成步骤并推送进度
 */

import { bookStore } from './book-generator.store';
import { pushBatchGenerationProgress } from '../../services/websocket.service.js';
import { createVideoProjectFromBook, generateVideoForProject } from '../video-generator/video-generator.service';
import { mergeChapterAudios } from '../player/player.service';
import { prisma } from '../../models';
import { advanceChapter, regenerateChapter } from './stage-manager';
import { estimateBookWords, estimateAudioMinutesFromWords, atomicReserveQuota, releaseQuota, getUserMonthlyCost, AUDIO_BILLING_CONFIG } from '../subscription/subscription.service';

// 步骤类型
export type GenerationStep = 'generate_content' | 'generate_audio' | 'merge_audio' | 'generate_video' | 'merge_video';

// 取消标志
const cancellationFlags = new Map<string, boolean>();

/**
 * 设置取消标志
 */
export function setCancellationFlag(taskId: string): void {
  cancellationFlags.set(taskId, true);
}

/**
 * 清除取消标志
 */
export function clearCancellationFlag(taskId: string): void {
  cancellationFlags.delete(taskId);
}

/**
 * 检查是否已取消
 */
export function isTaskCancelled(taskId: string): boolean {
  return cancellationFlags.get(taskId) === true;
}

/**
 * 批量生成编排器
 */
export class BatchGenerationOrchestrator {
  private taskId: string;
  private bookId: string;
  private steps: GenerationStep[];

  constructor(taskId: string, bookId: string, steps: GenerationStep[]) {
    this.taskId = taskId;
    this.bookId = bookId;
    this.steps = steps;
  }

  /**
   * 推送进度
   */
  private pushProgress(step: string, progress: number, message: string): void {
    pushBatchGenerationProgress(this.taskId, step, progress);
    console.log(`[BatchGen][${this.taskId}] ${step}: ${progress}% - ${message}`);
  }

  /**
   * 检查是否已取消
   */
  private checkCancellation(): void {
    if (isTaskCancelled(this.taskId)) {
      console.log(`[BatchGen][${this.taskId}] 任务已取消`);
      throw new Error('TASK_CANCELLED');
    }
  }

  /**
   * 执行所有步骤
   */
  async execute(): Promise<{ success: boolean; completedSteps: GenerationStep[]; failedStep?: string; error?: string }> {
    const completedSteps: GenerationStep[] = [];

    try {
      // 获取书籍信息
      const book = await bookStore.getById(this.bookId);
      if (!book) {
        return { success: false, completedSteps, failedStep: 'init', error: '书籍不存在' };
      }

      // 执行每个步骤
      for (let i = 0; i < this.steps.length; i++) {
        this.checkCancellation();

        const step = this.steps[i];
        const stepProgress = Math.round((i / this.steps.length) * 100);
        this.pushProgress(step, stepProgress, '开始执行');

        try {
          switch (step) {
            case 'generate_content':
              await this.executeGenerateContent();
              break;
            case 'generate_audio':
              await this.executeGenerateAudio();
              break;
            case 'merge_audio':
              await this.executeMergeAudio();
              break;
            case 'generate_video':
              await this.executeGenerateVideo();
              break;
            case 'merge_video':
              await this.executeMergeVideo();
              break;
            default:
              console.warn(`[BatchGen][${this.taskId}] 未知步骤: ${step}`);
          }

          completedSteps.push(step);
          this.pushProgress(step, 100, '执行完成');

          // 步骤间检查取消
          this.checkCancellation();

        } catch (error: any) {
          if (error.message === 'TASK_CANCELLED') {
            return { success: false, completedSteps, failedStep: step, error: '用户取消' };
          }
          console.error(`[BatchGen][${this.taskId}] 步骤 ${step} 执行失败:`, error);
          return { success: false, completedSteps, failedStep: step, error: error.message };
        }
      }

      // 更新书籍状态
      await bookStore.update(this.bookId, { progress: 100 });

      return { success: true, completedSteps };

    } catch (error: any) {
      console.error(`[BatchGen][${this.taskId}] 执行失败:`, error);
      return { success: false, completedSteps, error: error.message };
    } finally {
      // 清理取消标志
      clearCancellationFlag(this.taskId);
    }
  }

  /**
   * 执行内容生成步骤
   * 使用 LangGraph 生成书籍内容
   */
  private async executeGenerateContent(): Promise<void> {
    this.pushProgress('generate_content', 10, '检查书籍状态');

    const book = await bookStore.getById(this.bookId);
    if (!book) throw new Error('书籍不存在');

    // 如果书籍已有内容，则跳过
    const chapters = await bookStore.getChapterTree(this.bookId);
    const completedContent = chapters.filter((c: any) => c.genStage === 'content_completed');

    if (completedContent.length > 0) {
      this.pushProgress('generate_content', 50, `已有 ${completedContent.length} 个章节完成内容生成，跳过`);
      return;
    }

    // 预估费用并预留配额
    const userId = (book as any).userId;
    let reservedAmount = 0;
    if (userId) {
      try {
        const bookScale = book.bookScale || '1000';
        const wordEstimate = estimateBookWords(bookScale);
        const audioMinutes = estimateAudioMinutesFromWords(wordEstimate.avg);
        const estimatedCost = AUDIO_BILLING_CONFIG.pricing.monthly * audioMinutes + 0.1; // +LLM预估缓冲
        const quotaOk = await atomicReserveQuota(userId, estimatedCost);
        if (!quotaOk) {
          throw new Error('月度额度不足，请升级套餐');
        }
        reservedAmount = estimatedCost;
        console.log(`[BatchGen][${this.taskId}] 配额预留成功: ¥${estimatedCost.toFixed(4)}, userId=${userId}`);
      } catch (err: any) {
        console.error(`[BatchGen][${this.taskId}] 配额预留失败:`, err.message);
        throw err;
      }
    }

    this.pushProgress('generate_content', 20, '开始生成内容');

    // 调用 LangGraph 生成内容
    const { langGraphGenerator, resolveGenLevel } = await import('./index.js');

    // 异步执行生成，不阻塞
    langGraphGenerator.generate(this.bookId, book.description, book.bookScale || '1000', resolveGenLevel(book.bookScale || '1000'))
      .then(() => {
        console.log(`[BatchGen][${this.taskId}] 内容生成完成`);
      })
      .catch((err) => {
        console.error(`[BatchGen][${this.taskId}] 内容生成失败:`, err);
      });

    // 等待内容生成完成（轮询检查）
    let maxWaitTime = 3600 * 1000; // 最多等待60分钟
    let waited = 0;
    const checkInterval = 5000; // 每5秒检查一次
    let generationSuccess = false;

    try {
      while (waited < maxWaitTime) {
        this.checkCancellation();

        const currentChapters = await bookStore.getChapterTree(this.bookId);
        const leafNodes = currentChapters.filter((c: any) => c.level === Math.max(...currentChapters.map((ch: any) => ch.level || 0)));
        const completedCount = leafNodes.filter((c: any) => c.genStage === 'content_completed').length;
        const totalCount = leafNodes.length;

        if (totalCount > 0) {
          const progress = 20 + Math.round((completedCount / totalCount) * 60);
          this.pushProgress('generate_content', Math.min(progress, 90), `内容生成中: ${completedCount}/${totalCount}`);
        }

        // 检查是否全部完成
        if (totalCount > 0 && completedCount >= totalCount) {
          this.pushProgress('generate_content', 95, '内容生成完成');
          generationSuccess = true;
          return;
        }

        // 检查书籍状态
        const currentBook = await bookStore.getById(this.bookId);
        if (currentBook?.genStage === 'video_completed') {
          this.pushProgress('generate_content', 95, '内容生成完成');
          generationSuccess = true;
          return;
        }
        if (currentBook?.genStage === 'failed') {
          throw new Error('内容生成失败: ' + (currentBook.error || '未知错误'));
        }

        await this.sleep(checkInterval);
        waited += checkInterval;
      }
      throw new Error('内容生成超时');
    } finally {
      // 生成完成后，调整配额（多退少补）
      if (userId && reservedAmount > 0) {
        try {
          // 等待异步日志写入完成（AI调用日志是异步写入的）
          await this.sleep(2000);
          const actualCost = await getUserMonthlyCost(userId);
          const diff = actualCost - reservedAmount;
          if (Math.abs(diff) > 0.001) {
            if (diff > 0) {
              // 实际消耗超过预留，需要追加预留
              await atomicReserveQuota(userId, diff);
              console.log(`[BatchGen][${this.taskId}] 配额追加: ¥${diff.toFixed(4)}, userId=${userId}`);
            } else {
              // 实际消耗少于预留，退回多余部分
              await releaseQuota(userId, Math.abs(diff));
              console.log(`[BatchGen][${this.taskId}] 配额退回: ¥${Math.abs(diff).toFixed(4)}, userId=${userId}`);
            }
          }
        } catch (err) {
          console.error(`[BatchGen][${this.taskId}] 配额调整失败:`, err);
        }
      }
    }
  }

  /**
   * 执行音频生成步骤
   */
  private async executeGenerateAudio(): Promise<void> {
    this.pushProgress('generate_audio', 10, '开始生成音频');

    const chapters = await bookStore.getChapterTree(this.bookId);
    const maxLevel = chapters.length > 0
      ? Math.max(...chapters.map((c: any) => c.level || 0))
      : 0;

    // 获取叶节点
    const leafNodes = chapters.filter((c: any) => c.level === maxLevel);

    // 检查内容状态 - 叶节点必须有content且genStage为content_completed
    const leafNodesWithContent = leafNodes.filter((c: any) => c.content && c.genStage === 'content_completed');

    if (leafNodesWithContent.length === 0) {
      throw new Error('没有内容生成完成的章节，请先生成内容');
    }

    this.pushProgress('generate_audio', 20, `开始生成 ${leafNodesWithContent.length} 个章节音频`);

    // 异步生成所有叶节点音频
    const generationPromises: Promise<void>[] = [];
    for (const sub of leafNodesWithContent) {
      generationPromises.push(
        bookStore.generateChapterAudioById(sub.id, 1)
          .then(() => {
            console.log(`[BatchGen][${this.taskId}] 章节 ${sub.number} 音频生成完成`);
          })
          .catch((err) => {
            console.error(`[BatchGen][${this.taskId}] 章节 ${sub.number} 音频生成失败:`, err);
          }) as Promise<void>
      );
    }

    // 等待音频生成完成（轮询检查）
    let maxWaitTime = 3600 * 1000; // 最多等待60分钟
    let waited = 0;
    const checkInterval = 3000; // 每3秒检查一次

    while (waited < maxWaitTime) {
      this.checkCancellation();

      const currentChapters = await bookStore.getChapterTree(this.bookId);
      const currentLeafNodes = currentChapters.filter((c: any) => c.level === maxLevel);
      const completedCount = currentLeafNodes.filter((c: any) => c.audioUrl && c.audioUrl !== '').length;
      const totalCount = currentLeafNodes.length;

      if (totalCount > 0) {
        const progress = 20 + Math.round((completedCount / totalCount) * 70);
        this.pushProgress('generate_audio', Math.min(progress, 95), `音频生成中: ${completedCount}/${totalCount}`);
      }

      // 检查是否全部完成
      if (totalCount > 0 && completedCount >= totalCount) {
        this.pushProgress('generate_audio', 100, '音频生成完成');
        return;
      }

      await this.sleep(checkInterval);
      waited += checkInterval;
    }

    throw new Error('音频生成超时');
  }

  /**
   * 执行音频合并步骤
   */
  private async executeMergeAudio(): Promise<void> {
    this.pushProgress('merge_audio', 10, '开始合并音频');

    const chapters = await bookStore.getChapterTree(this.bookId);
    const maxLevel = chapters.length > 0
      ? Math.max(...chapters.map((c: any) => c.level || 0))
      : 0;

    if (maxLevel <= 1) {
      this.pushProgress('merge_audio', 100, '书籍层级不足，跳过音频合并');
      return;
    }

    // 获取叶节点
    const leafNodes = chapters.filter((c: any) => c.level === maxLevel);

    // 检查所有叶节点是否都有音频
    const leafNodesWithAudio = leafNodes.filter((c: any) => c.audioUrl && c.audioUrl !== '');

    if (leafNodesWithAudio.length !== leafNodes.length) {
      throw new Error(`并非所有章节都已完成音频生成，无法合并 (${leafNodesWithAudio.length}/${leafNodes.length})`);
    }

    this.pushProgress('merge_audio', 20, `开始合并 ${leafNodes.length} 个章节音频`);

    // 构建 sectionId → chapterId 映射（3层树需要）
    const sectionToChapter = new Map<number, number>();
    if (maxLevel >= 3) {
      const sections = chapters.filter((c: any) => c.level === 2);
      for (const sec of sections) {
        if (sec.parentId != null) {
          sectionToChapter.set(sec.id, sec.parentId);
        }
      }
    }

    // 按章（level=1）分组叶节点
    const groupedByChapter: { [key: number]: any[] } = {};
    leafNodesWithAudio.forEach(node => {
      let chapterId: number | null = null;
      if (maxLevel === 2) {
        chapterId = node.parentId;
      } else if (maxLevel === 3 && node.parentId != null) {
        chapterId = sectionToChapter.get(node.parentId) || null;
      }

      if (chapterId) {
        if (!groupedByChapter[chapterId]) {
          groupedByChapter[chapterId] = [];
        }
        groupedByChapter[chapterId].push(node);
      }
    });

    const chapterIds = Object.keys(groupedByChapter);
    const totalChapters = chapterIds.length;
    let processedChapters = 0;

    // 对每个章下的叶节点音频进行合并
    for (const chapterId of chapterIds) {
      this.checkCancellation();

      const chId = parseInt(chapterId);
      const childNodes = groupedByChapter[chId];

      if (childNodes.length > 0) {
        try {
          const mergedAudioUrl = await mergeChapterAudios(chId);

          if (mergedAudioUrl) {
            console.log(`[BatchGen][${this.taskId}] 章 ${chId} 音频合并完成`);
          }
        } catch (error) {
          console.error(`[BatchGen][${this.taskId}] 章 ${chId} 音频合并失败:`, error);
        }
      }

      processedChapters++;
      const progress = 20 + Math.round((processedChapters / totalChapters) * 70);
      this.pushProgress('merge_audio', Math.min(progress, 95), `音频合并中: ${processedChapters}/${totalChapters}`);
    }

    this.pushProgress('merge_audio', 100, '音频合并完成');
  }

  /**
   * 执行视频生成步骤
   */
  private async executeGenerateVideo(): Promise<void> {
    this.pushProgress('generate_video', 10, '开始生成视频');

    const chapters = await bookStore.getChapterTree(this.bookId);
    const maxLevel = chapters.length > 0
      ? Math.max(...chapters.map((c: any) => c.level || 0))
      : 0;

    // 获取叶节点
    const leafNodes = chapters.filter((c: any) => c.level === maxLevel);

    // 过滤出有音频的叶节点
    const leafNodesWithAudio = leafNodes.filter((c: any) => c.audioUrl && c.audioUrl !== '');

    if (leafNodesWithAudio.length === 0) {
      throw new Error('没有音频生成完成的章节，请先生成音频');
    }

    this.pushProgress('generate_video', 20, `开始生成 ${leafNodesWithAudio.length} 个章节视频`);

    // 设置所有有音频的叶节点视频状态为生成中
    for (const sub of leafNodesWithAudio) {
      await advanceChapter(sub.id, 'video_generating').catch(() => {});
    }

    // 异步生成所有叶节点视频
    for (const sub of leafNodesWithAudio) {
      this.checkCancellation();

      try {
        const project = await createVideoProjectFromBook(
          parseInt(this.bookId),
          sub.id,
          1
        );

        if (!project) {
          console.error(`[BatchGen][${this.taskId}] 章节 ${sub.number} 视频项目创建失败`);
          await regenerateChapter(sub.id, 'failed').catch(() => {});
          continue;
        }

        const result = await generateVideoForProject(project.id);

        if (result.success && result.outputUrl) {
          await bookStore.updateChapterById(sub.id, {
            videoUrl: result.outputUrl,
            videoDuration: result.duration,
          });
          await advanceChapter(sub.id, 'video_completed');
          console.log(`[BatchGen][${this.taskId}] 章节 ${sub.number} 视频生成成功`);
        } else {
          await regenerateChapter(sub.id, 'failed').catch(() => {});
        }
      } catch (error) {
        console.error(`[BatchGen][${this.taskId}] 章节 ${sub.number} 视频生成异常:`, error);
        await regenerateChapter(sub.id, 'failed').catch(() => {});
      }
    }

    // 等待视频生成完成（轮询检查）
    let maxWaitTime = 3600 * 1000; // 最多等待60分钟
    let waited = 0;
    const checkInterval = 5000; // 每5秒检查一次

    while (waited < maxWaitTime) {
      this.checkCancellation();

      const currentChapters = await bookStore.getChapterTree(this.bookId);
      const currentLeafNodes = currentChapters.filter((c: any) => c.level === maxLevel);
      const completedCount = currentLeafNodes.filter((c: any) => c.videoUrl && c.videoUrl !== '').length;
      const totalCount = currentLeafNodes.length;

      if (totalCount > 0) {
        const progress = 20 + Math.round((completedCount / totalCount) * 70);
        this.pushProgress('generate_video', Math.min(progress, 95), `视频生成中: ${completedCount}/${totalCount}`);
      }

      // 检查是否全部完成
      if (totalCount > 0 && completedCount >= totalCount) {
        this.pushProgress('generate_video', 100, '视频生成完成');
        return;
      }

      await this.sleep(checkInterval);
      waited += checkInterval;
    }

    throw new Error('视频生成超时');
  }

  /**
   * 执行视频合并步骤
   */
  private async executeMergeVideo(): Promise<void> {
    this.pushProgress('merge_video', 10, '开始合并视频');

    const chapters = await bookStore.getChapterTree(this.bookId);
    const maxLevel = chapters.length > 0
      ? Math.max(...chapters.map((c: any) => c.level || 0))
      : 0;

    if (maxLevel <= 1) {
      this.pushProgress('merge_video', 100, '书籍层级不足，跳过视频合并');
      return;
    }

    // 获取叶节点
    const leafNodes = chapters.filter((c: any) => c.level === maxLevel);

    // 检查所有叶节点是否都有视频
    const leafNodesWithVideo = leafNodes.filter((c: any) => c.videoUrl && c.videoUrl !== '');

    if (leafNodesWithVideo.length !== leafNodes.length) {
      throw new Error(`并非所有章节都已完成视频生成，无法合并 (${leafNodesWithVideo.length}/${leafNodes.length})`);
    }

    this.pushProgress('merge_video', 20, `开始合并 ${leafNodes.length} 个章节视频`);

    // 按父节点分组叶节点
    const groupedByParent: { [key: number]: any[] } = {};
    leafNodesWithVideo.forEach(node => {
      if (node.parentId != null) {
        if (!groupedByParent[node.parentId]) {
          groupedByParent[node.parentId] = [];
        }
        groupedByParent[node.parentId].push(node);
      }
    });

    const totalParents = Object.keys(groupedByParent).length;
    let processedParents = 0;

    // 对每个父节点下的叶节点视频进行合并
    for (const parentId in groupedByParent) {
      this.checkCancellation();

      const childNodes = groupedByParent[parentId];
      if (childNodes.length > 0) {
        const parentChapter = chapters.find((c: any) => c.id === parseInt(parentId));

        if (parentChapter) {
          try {
            // 简化处理：直接使用第一个视频或标记为已完成
            // 实际的视频合并需要FFmpeg处理
            const firstChildVideo = childNodes[0].videoUrl;
            if (firstChildVideo) {
              await bookStore.updateChapterById(parentChapter.id, {
                videoUrl: firstChildVideo,
              });
              console.log(`[BatchGen][${this.taskId}] 章节 ${parentChapter.number} 视频处理完成`);
            }
          } catch (error) {
            console.error(`[BatchGen][${this.taskId}] 章节 ${parentChapter.number} 视频处理失败:`, error);
          }
        }
      }

      processedParents++;
      const progress = 20 + Math.round((processedParents / totalParents) * 70);
      this.pushProgress('merge_video', Math.min(progress, 95), `视频处理中: ${processedParents}/${totalParents}`);
    }

    this.pushProgress('merge_video', 100, '视频合并完成');
  }

  /**
   * 休眠辅助函数
   */
  private sleep(ms: number): Promise<void> {
    return new Promise(resolve => setTimeout(resolve, ms));
  }
}

/**
 * 创建批量生成任务
 */
export async function createBatchGenerationTask(
  bookId: string,
  steps: GenerationStep[]
): Promise<{ taskId: string; orchestrator: BatchGenerationOrchestrator }> {
  const taskId = `batch_${bookId}_${Date.now()}`;
  const orchestrator = new BatchGenerationOrchestrator(taskId, bookId, steps);

  return { taskId, orchestrator };
}