LaiTool_PRO/src/main/service/book/subBookHandle/bookPromptHandle.ts

import { OperateBookType, PromptMergeType } from '@/define/enum/bookEnum'
import { Book } from '@/define/model/book/book'
import { errorMessage, SendReturnMessage, successMessage } from '@/public/generalTools'

import { cloneDeep, isEmpty } from 'lodash'
import { AiReasonCommon } from '../../aiReason/aiReasonCommon'
import { DEFINE_STRING } from '@/define/ipcDefineString'
import { GeneralResponse } from '@/define/model/generalResponse'
import { ExecuteConcurrently } from '@/define/Tools/common'
import { OptionKeyName } from '@/define/enum/option'
import { optionSerialization } from '../../option/optionSerialization'
import { SettingModal } from '@/define/model/setting'
import { MJServiceHandle } from '@/main/service/mj/mjServiceHandle'
import { BookBasicHandle } from './bookBasicHandle'
import { PresetCategory } from '@/define/data/presetData'
import { ValidateJsonAndParse } from '@/define/Tools/validate'
import { BookTask } from '@/define/model/book/bookTask'
import { SDServiceHandle } from '../../sd/sdServiceHandle'
import { aiHandle } from '../../ai'
import { AICharacterAnalyseRequestData } from '@/define/data/aiData/aiPrompt/CharacterAndScene/aiCharacterAnalyseRequestData'
import { AISceneAnalyseRequestData } from '@/define/data/aiData/aiPrompt/CharacterAndScene/aiSceneAnalyseRequestData'
import { t } from '@/i18n'
import { OptionRealmService } from '@/define/db/service/optionService'
import { PresetModel } from '@/define/model/preset'
import { AiInferenceModelModel } from '@/define/data/aiData/aiData'

export class BookPromptHandle extends BookBasicHandle {
  aiReasonCommon: AiReasonCommon
  mjServiceHandle: MJServiceHandle
  sdServiceHandle: SDServiceHandle

  constructor() {
    super()
    this.aiReasonCommon = new AiReasonCommon()
    this.mjServiceHandle = new MJServiceHandle()
    this.sdServiceHandle = new SDServiceHandle()
  }

  /**
   * 为小说分镜生成AI提示词
   *
   * 该方法根据操作类型获取相应的分镜数据，然后使用AI推理为每个分镜生成提示词。
   * 支持三种操作模式：
   * 1. 对整个小说任务的所有分镜进行处理（BOOKTASK）
   * 2. 对单个指定分镜进行处理（BOOKTASKDETAIL）
   * 3. 对指定分镜及其后续所有分镜进行处理（UNDERBOOKTASK）
   *
   * 生成过程中会实时向前端发送进度通知，并根据系统设置控制并发处理数量。
   *
   * @param {string} id - 根据operateBookType不同，可能是小说任务ID或分镜ID
   * @param {OperateBookType} operateBookType - 操作类型，决定处理范围
   * @param {boolean} coverData - 是否覆盖已有提示词数据：
   *                             true-处理所有分镜, false-只处理空白提示词的分镜
   * @returns {Promise<SuccessItem | ErrorItem>} 操作结果，成功或失败的标准化响应
   *
   * @throws 如果找不到指定分镜数据或操作类型未知，将抛出异常
   *
   * @example
   * // 为整个小说任务生成提示词，不覆盖已有数据
   * const result = await bookPromptHandle.OriginalGetAiPrompt(
   *   "task-123",
   *   OperateBookType.BOOKTASK,
   *   false
   * );
   */
  OriginalGetAiPrompt = async (
    id: string,
    operateBookType: OperateBookType,
    coverData: boolean
  ): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> => {
    try {
      let bookTask: Book.SelectBookTask = {} as Book.SelectBookTask
      let bookTaskDetails: Book.SelectBookTaskDetail[] = []
      let allBookTaskDetails: Book.SelectBookTaskDetail[] = []

      await this.InitBookBasicHandle()
      if (operateBookType == OperateBookType.BOOKTASK) {
        bookTask = await this.bookTaskService.GetBookTaskDataById(id, true)
        allBookTaskDetails = await this.bookTaskDetailService.GetBookTaskDetailDataByCondition({
          bookTaskId: id
        })
        if (!coverData) {
          // 不覆盖数据，只推理空白提示词
          bookTaskDetails = allBookTaskDetails.filter((item) => isEmpty(item.gptPrompt))
        } else {
          // 不覆盖，就是全部
          bookTaskDetails = allBookTaskDetails
        }
      } else if (operateBookType == OperateBookType.BOOKTASKDETAIL) {
        let singleBookTaskDetail = await this.bookTaskDetailService.GetBookTaskDetailDataById(id, true)

        bookTask = await this.bookTaskService.GetBookTaskDataById(
          singleBookTaskDetail.bookTaskId as string,
          true
        )
        bookTaskDetails = [singleBookTaskDetail]
      } else if (operateBookType == OperateBookType.UNDERBOOKTASK) {
        let singleBookTaskDetail = await this.bookTaskDetailService.GetBookTaskDetailDataById(id, true)

        bookTask = await this.bookTaskService.GetBookTaskDataById(
          singleBookTaskDetail.bookTaskId as string,
          true
        )
        // 获取全部的分镜数据
        allBookTaskDetails = await this.bookTaskDetailService.GetBookTaskDetailDataByCondition({
          bookTaskId: bookTask.id
        })

        // 只要当前行往下的数据
        for (let i = 0; i < allBookTaskDetails.length; i++) {
          const element = allBookTaskDetails[i]
          if (bookTaskDetails == undefined) {
            bookTaskDetails = []
          }
          if (i + 1 >= (singleBookTaskDetail.no as number)) {
            bookTaskDetails.push(element)
          }
        }
      } else {
        throw new Error(t('未知操作'))
      }
      if (!allBookTaskDetails || allBookTaskDetails.length <= 0) {
        allBookTaskDetails = await this.bookTaskDetailService.GetBookTaskDetailDataByCondition({
          bookTaskId: bookTask.id
        })
      }

      if (bookTaskDetails.length <= 0) {
        throw new Error(t('没有找到要推理的分镜数据'))
      }
      let generalSettingOption = this.optionRealmService.GetOptionByKey(
        OptionKeyName.Software.GeneralSetting
      )
      let generalSetting = optionSerialization<SettingModal.GeneralSettings>(
        generalSettingOption,
        t('设置 -> 通用设置')
      )
      let tasks = [] as Array<() => Promise<any>>

      // 获取当前task的角色和场景数据
      let autoAnalyzeCharacter = bookTask.autoAnalyzeCharacter ?? '{}'
      let autoAnalyzeCharacterData =
        ValidateJsonAndParse<BookTask.BookTaskCharacterAndScene>(autoAnalyzeCharacter)
      let characterData = autoAnalyzeCharacterData[PresetCategory.Character] ?? []
      let sceneData = autoAnalyzeCharacterData[PresetCategory.Scene] ?? []
      let characterString = ''
      let sceneString = ''

      if (characterData.length > 0) {
        characterString = characterData.map((item) => item.name + '：' + item.prompt).join('\n')
        characterString = '角色设定：' + '\n' + characterString
      }
      if (sceneData.length > 0) {
        sceneString = sceneData.map((item) => item.name + '：' + item.prompt).join('\n')
        sceneString = '场景设定：' + '\n' + sceneString
      }

      // 获取自定义的提示词数据
      let optionsData: AiInferenceModelModel[] = [];
      const optionRealmService = await OptionRealmService.getInstance();
      let customInferencePreset = optionRealmService.GetOptionByKey(OptionKeyName.InferenceAI.CustomInferencePreset);
      if (customInferencePreset != null && customInferencePreset.value != null && !isEmpty(customInferencePreset.value)) {
        let customInferencePresetData = optionSerialization<PresetModel.AIPresetTemplate[]>(customInferencePreset, t('设置 -> 推理设置 -> 自定义预设'), []);
        for (let i = 0; i < customInferencePresetData.length; i++) {
          const element = customInferencePresetData[i];
          optionsData.push({
            value: element.id,
            label: element.name,
            hasExample: element.hasExample,
            mustCharacter: element.mustCharacter,
            requestBody: element.requestBody,
            allAndExampleContent: null
          })
        }
      }

      // 添加异步任务
      for (let i = 0; i < bookTaskDetails.length; i++) {
        const element = bookTaskDetails[i]
        tasks.push(async () => {
          let content = await this.aiReasonCommon.OriginalInferencePrompt(
            element,
            allBookTaskDetails,
            15, // 上下文关联行数
            characterString,
            sceneString, optionsData
          )
          console.log(element.afterGpt, content)
          // 修改推理出来的数据
          await this.bookTaskDetailService.ModifyBookTaskDetailById(element.id as string, {
            gptPrompt: content
          })
          // 每次完成，都要向前端返回信息
          SendReturnMessage(
            {
              code: 1,
              id: element.id as string,
              data: {
                content: content,
                progress: {
                  current: i,
                  total: bookTaskDetails.length
                } as GeneralResponse.ProgressResponse
              }
            },
            DEFINE_STRING.BOOK.ORIGINAL_GET_AI_PROMPT_RETURN
          )
        })
      }
      // 分批次执行异步任务
      await ExecuteConcurrently(tasks, global.am.isPro ? (generalSetting.concurrency ?? 1) : 1)
      // 执行完毕
      return successMessage(null, t("推理所有数据完成"), 'BookPrompt_OriginalGetPrompt')
    } catch (error: any) {
      // 处理错误，返回错误信息
      return errorMessage(
        t("推理所有数据失败，{error}", {
          error: error.message
        }),
        'BookPromptHandle_OriginalGetAiPrompt'
      )
    }
  }

  /**
   * AI分镜头合并
   * @param bookTaskId 小说任务ID
   * @param type 合并类型
   * @returns 操作结果
   * @throws 操作失败时抛出异常
   * @example
   * const result = await bookPromptHandle.AIStoryboardMerge("task-123", BookTask.StoryboardMergeType.MJ);
   * if (result.code === 1) {
   * // 合并成功
   * } else {
   * // 合并失败
   * }
   */
  AIStoryboardMerge = async (bookTaskId: string, type: BookTask.StoryboardMergeType) => {
    try {
      let res = await aiHandle.AIStoryboardMerge(bookTaskId, type)
      return successMessage(res, t('AI分镜头合并成功'), 'BookPromptHandle_AIStoryboardMerge')
    } catch (error: any) {
      return errorMessage(
        t('AI分镜头合并成功', {
          error: error.message
        }),
        'BookPromptHandle_AIStoryboardMerge'
      )
    }
  }

  /**
   * 合并提示词
   *
   * 该方法根据提供的合并类型(MJ或SD)，将分镜中的提示词按照规则进行合并处理。
   * 支持两种合并类型：
   * - MJ_MERGE: 使用MidJourney规则合并提示词
   * - SD_MERGE: 使用Stable Diffusion规则合并提示词
   *
   * 合并操作支持不同的处理范围，由operateBookType参数决定：
   * - BOOKTASK: 处理整个小说任务中的所有分镜
   * - BOOKTASKDETAIL: 仅处理单个指定分镜
   * - UNDERBOOKTASK: 处理指定分镜及其后续所有分镜
   *
   * @param {string} id - 根据operateBookType不同，可能是小说任务ID或分镜ID
   * @param {PromptMergeType} type - 合并类型，决定使用哪种规则合并提示词
   * @param {OperateBookType} operateBookType - 操作类型，决定处理范围
   * @returns {Promise<SuccessItem | ErrorItem>} 操作结果，成功或失败的标准化响应
   *
   * @throws {Error} 如果合并类型未知，将抛出异常
   *
   * @example
   * // 使用MidJourney规则合并单个分镜的提示词
   * const result = await bookPromptHandle.MergePrompt(
   *   "detail-123",
   *   PromptMergeType.MJ_MERGE,
   *   OperateBookType.BOOKTASKDETAIL
   * );
   */
  MergePrompt = async (
    id: string,
    type: PromptMergeType,
    operateBookType: OperateBookType
  ): Promise<GeneralResponse.SuccessItem | GeneralResponse.ErrorItem> => {
    try {
      if (type == PromptMergeType.MJ_MERGE) {
        return await this.mjServiceHandle.MergeMJPrompt(id, operateBookType)
      } else if (type == PromptMergeType.SD_MERGE) {
        return await this.sdServiceHandle.MergeSDPrompt(id, operateBookType)
      } else {
        throw new Error(t('未知的合并模式，请检查'))
      }
    } catch (error: any) {
      return errorMessage(
        t("合并提示词失败，{error}", {
          error: error.message
        }),
        'ReverseBook_MergePrompt'
      )
    }
  }

  /**
   * 自动分析书籍任务中的角色或场景
   *
   * 该方法使用AI技术分析书籍任务中的所有分镜文本，自动识别并提取角色或场景信息。
   * 处理过程包括合并所有分镜文本、调用AI分析、解析返回结果并格式化为结构化数据。
   * 分析结果会更新到书籍任务的autoAnalyzeCharacter字段中，以JSON格式存储。
   *
   * @param {string} bookTaskId - 要分析的书籍任务ID
   * @param {PresetCategory} type - 分析类型，必须是PresetCategory.Character(角色)或PresetCategory.Scene(场景)
   *
   * @returns {Promise<GeneralResponse.SuccessItem | GeneralResponse.ErrorItem>}
   *   成功时返回分析得到的角色或场景数据对象数组，失败时返回错误信息
   *
   * @throws {Error} 当分析类型无效、找不到书籍任务数据或分镜数据为空时抛出错误
   *
   * @example
   * // 分析书籍任务中的角色
   * const characterResult = await bookPromptHandle.AutoAnalyzeCharacterOrScene(
   *   "task-123",
   *   PresetCategory.Character
   * );
   *
   * // 分析书籍任务中的场景
   * const sceneResult = await bookPromptHandle.AutoAnalyzeCharacterOrScene(
   *   "task-123",
   *   PresetCategory.Scene
   * );
   */
  AutoAnalyzeCharacterOrScene = async (bookTaskId: string, type: PresetCategory): Promise<GeneralResponse.SuccessItem | GeneralResponse.ErrorItem> => {
    try {
      if (type != PresetCategory.Character && type != PresetCategory.Scene) {
        throw new Error(t('分析的类型只能是角色或场景，请检查'))
      }
      await this.InitBookBasicHandle()

      let bookTask = await this.bookTaskService.GetBookTaskDataById(bookTaskId, true)
      let bookTaskDetails = await this.bookTaskDetailService.GetBookTaskDetailDataByCondition({
        bookTaskId: bookTaskId
      })
      if (bookTaskDetails.length <= 0) {
        throw new Error(t('没有找到要分析的分镜数据，请先导入文案或者时srt！'))
      }

      let words = bookTaskDetails
        .map((item) => {
          return item.afterGpt
        })
        .join('\r\n')

      let requestData: OpenAIRequest.Request
      if (type == PresetCategory.Character) {
        requestData = cloneDeep(AICharacterAnalyseRequestData)
      } else if (type == PresetCategory.Scene) {
        requestData = cloneDeep(AISceneAnalyseRequestData)
      } else {
        throw new Error(t('未知的分析类型，请检查'))
      }

      requestData.messages = this.aiReasonCommon.replaceMessageObject(requestData.messages, {
        textContent: words
      })
      await this.aiReasonCommon.GetAISetting()
      delete requestData.model

      let content = await this.aiReasonCommon.FetchGpt(requestData.messages, requestData)

      let autoAnalyzeCharacter = bookTask.autoAnalyzeCharacter ?? '{}'
      let autoAnalyzeCharacterData =
        ValidateJsonAndParse<BookTask.BookTaskCharacterAndScene>(autoAnalyzeCharacter)
      let returnData = content.split('\n').filter((item) => !isEmpty(item))
      let newData: BookTask.BookTaskCharacterAndSceneObject[] = []
      for (let i = 0; i < returnData.length; i++) {
        const element = returnData[i]
        let splitData = element.split('：')
        if (splitData.length < 2) {
          continue
        }
        let tempData = {
          no: i + 1,
          id: crypto.randomUUID(),
          name: splitData[0],
          prompt: splitData[1]
        } as BookTask.BookTaskCharacterAndSceneObject
        newData.push(tempData)
      }

      if (type == PresetCategory.Character) {
        autoAnalyzeCharacterData[PresetCategory.Character] = newData
      } else if (type == PresetCategory.Scene) {
        // 场景数据
        autoAnalyzeCharacterData[PresetCategory.Scene] = newData
      }
      // 重新写入
      await this.bookTaskService.ModifyBookTaskDataById(bookTaskId, {
        autoAnalyzeCharacter: JSON.stringify(autoAnalyzeCharacterData)
      })

      return successMessage(
        autoAnalyzeCharacterData,
        t('自动分析角色或场景成功'),
        'ReverseBook_AutoAnalyzeCharacterOrScene'
      )
    } catch (error: any) {
      return errorMessage(
        t("自动分析角色或场景失败，{error}", {
          error: error.message
        }),
        'ReverseBook_AutoAnalyzeCharacterOrScene'
      )
    }
  }
}