LaiTool/src/main/Service/Book/ReverseBook.ts

import { successMessage, errorMessage } from '../../Public/generalTools'
import { BookBasic } from './BooKBasic'
import { BookService } from '../../../define/db/service/Book/bookService'
import { BookTaskService } from '../../../define/db/service/Book/bookTaskService'
import { define } from '../../../define/define'
import fs from 'fs'
import { DEFINE_STRING } from "../../../define/define_string";
import path from 'path'
import { BasicReverse } from './basicReverse'
import { BookTaskDetailService } from '../../../define/db/service/Book/bookTaskDetailService'
import { TaskScheduler } from "../task/taskScheduler"
import { Book } from '../../../model/book'
import { LoggerStatus, OtherData, ResponseMessageType } from '../../../define/enum/softwareEnum'
import { GeneralResponse } from '../../../model/generalResponse'
import { Subtitle } from '../Subtitle/subtitle'
import { Watermark } from '../watermark'
import { BookBackTaskType, BookType, OperateBookType, PromptMergeType, TagDefineType, TaskExecuteType } from '../../../define/enum/bookEnum'
import { MJOpt } from '../MJ/mj'
import { TagDefine } from '../../../define/tagDefine'
import { BookServiceBasic } from '../ServiceBasic/bookServiceBasic'
import { SDOpt } from '../SD/sd'
import { isEmpty } from 'lodash'



/**
 * 一键反推的相关操作
 */
export class ReverseBook {
  basicReverse: BasicReverse
  taskScheduler: TaskScheduler
  mjOpt: MJOpt = new MJOpt()
  sdOpt: SDOpt = new SDOpt()
  tagDefine: TagDefine
  subtitle: Subtitle
  watermark: Watermark
  bookServiceBasic: BookServiceBasic
  bookBasic: BookBasic

  constructor() {
    this.basicReverse = new BasicReverse()
    this.tagDefine = new TagDefine()
    this.subtitle = new Subtitle()
    this.watermark = new Watermark()
    this.taskScheduler = new TaskScheduler()
    this.bookServiceBasic = new BookServiceBasic()
    this.bookBasic = new BookBasic()
  }
  // 主动返回前端的消息
  sendReturnMessage(data: GeneralResponse.MessageResponse, message_name = DEFINE_STRING.BOOK.GET_COPYWRITING_RETURN) {
    global.newWindow[0].win.webContents.send(message_name, data)
  }

  //#region 小说相关操作

  /**
   * 获取当前的小说数据
   * @param {*} bookQuery
   */
  async GetBookData(bookQuery) {
    try {
      let _bookService = await BookService.getInstance()
      // 添加小说
      let res = _bookService.GetBookData(bookQuery)
      if (res.code == 0) {
        throw new Error(res.message)
      }
      return res
    } catch (error) {
      return errorMessage(error.message, 'ReverseBook_GetBookData')
    }
  }

  //#endregion

  //#region 小说批次任务相关操作

  /**
   * 获取小说的任务列表
   * @param {*} bookTaskCondition  查询任务列表的条件
   */
  async GetBookTaskData(bookTaskCondition: Book.QueryBookTaskCondition): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
    try {
      let res = await this.bookServiceBasic.GetBookTaskData(bookTaskCondition)
      // //TODO 这个后面是不是将所有的数据都是用数据库
      // // 这边加载自定义风格
      // let styleLists = await this.tagDefine.getTagDataByTypeAndProperty('dynamic', "style_tags");
      // if (styleLists.code == 0) {
      //   throw new Error('获取自定义风格失败')
      // }

      // let styleTags = styleLists.data as Book.BookStyle[];
      // // 异步操作获取风格信息
      // for (let index = 0; index < res.data.bookTasks.length; index++) {
      //   const element = res.data.bookTasks[index];
      //   let styleList = [] as Book.BookStyle[] | Book.DefineBookStyle[];
      //   // if (element.imageStyle.length > 0) { // 软件自带的风格
      //   //   let infoRes = ImageStyleDefine.getImageStyleInfomation(JSON.stringify(element.imageStyle));
      //   //   if (infoRes.code == 0) {
      //   //     throw new Error('获取风格失败');
      //   //   }
      //   //   styleList = infoRes.data;
      //   // }
      //   if (element.customizeImageStyle.length > 0) { // 自定义的风格
      //     element.customizeImageStyle.forEach((item) => {
      //       let style = styleTags.find((tag) => tag.key == item);
      //       if (style) {
      //         styleList.push(style);
      //       }
      //     });
      //   }
      //   res.data.bookTasks[index].styleList = styleList;
      // }
      return successMessage(res, '获取小说任务成功', 'ReverseBook_GetBookTaskData')
    } catch (error) {
      return errorMessage(
        '获取小说对应批次错误，错误信息入下：' + error.message,
        'ReverseBook_GetBookTaskData'
      )
    }
  }

  /**
   * 获取小说的所有的任务详情
   * @param bookTaskId 小说任务的ID
   */
  async GetBookTaskDetail(bookTaskId: string): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
    try {
      let res = await this.bookServiceBasic.GetBookTaskDetailData({ bookTaskId: bookTaskId })

      // 这边再所有的图片后面添加后缀，这样可以保证图片实时更新
      res = res.map(item => {
        return {
          ...item,
          outImagePath: isEmpty(item.outImagePath) ? item.outImagePath : item.outImagePath + '?t=' + new Date().getTime(),
          subImagePath: item.subImagePath && item.subImagePath.length > 0 ? item.subImagePath.map(it => it + '?t=' + new Date().getTime()) : item.subImagePath
        }
      })

      return successMessage(res, '获取小说任务详情成功', 'ReverseBook_GetBookTaskDetail')
    } catch (error) {
      return errorMessage(
        '获取小说对应批次错误，错误信息入下：' + error.message,
        'ReverseBook_GetBookTaskData'
      )
    }
  }

  //#endregion

  //#region 一键全自动

  /**
   * 全自动任务（这边是任务入口，都是在这边调用）
   * @param {*} value
   * @returns
   */
  async AutoAction(bookId) {
    try {
      // 在一键全自动之前，当前小说对应的批次任务中的所有的子任务都改为fail，然后再执行
      // 获取对应的小说小说数据，找到对应的小说视频地址
      // let _bookService = await BookService.getInstance()
      // let _bookTaskService = await BookTaskService.getInstance()

      // let bookData = _bookService.GetBookDataById(bookId)
      // if (bookData.data == null) {
      //   throw new Error('没有找到对应的小说数据，请检查bookId是否正确')
      // }

      // // 获取小说对应的批次任务数据，默认初始化为第一个
      // let bookTaskRes = _bookTaskService.GetBookTaskData({
      //   bookId: bookId,
      //   name: 'output_00001'
      // })
      // if (bookTaskRes.data.bookTasks.length <= 0 || bookTaskRes.data.total <= 0) {
      //   throw new Error('没有找到对应的小说批次任务数据，请检查bookId是否正确')
      // }

      // // 获取小说的视频地址
      // let book = bookData.data
      // let bookTask = bookTaskRes.data.bookTasks[0]

      // // 将当前小说对应的批次任务中的所有的子任务都改为fail
      // let updateTaskRes = _bookTaskService.UpdetedBookTaskToFail(bookId, bookTask.id)

      // // 添加分镜任务 后面就会全自动的开始执行
      // let res = await this.basicReverse.AddFrameDataTask(bookId)

      // 添加分割视频任务
      // let res = await this.basicReverse.AddCutVideoDataTask(bookId)

      // 添加音频分离任务
      // let res = await this.basicReverse.AddSplitAudioDataTask(bookId)

      // 添加图片抽帧任务
      let res = await this.basicReverse.AddGetFrameTask(bookId)

      if (res.code == 0) {
        throw new Error(res.message)
      }
    } catch (error) {
      return errorMessage(error.message, 'ReverseBook_AutoAction')
    }
  }

  /**
   * 开始分镜任务
   */
  async ComputeStoryboard(bookId: string): Promise<any> {
    try {
      let { book, bookTask } = await this.bookBasic.GetBookAndTask(bookId, 'output_00001')
      let res = await this.basicReverse.ComputeStoryboardFunc(bookId, bookTask.id);
      // 分镜成功直接返回
      return successMessage(null, res, "ReverseBook_ComputeStoryboard")

    } catch (error) {
      return errorMessage("分镜计算失败，失败信息如下 ：" + error.message, 'ReverseBook_ComputeStoryboard')
    }
  }

  //#endregion

  //#region 反推相关任务

  /**
   * 所有的反推任务的入口
   * @param id 需要反推的ID，可以是小说任务ID，也可以是小说分镜ID
   * @param operateBookType 操作的分类
   * @param type 反推的类型
   * @returns 
   */
  async AddReversePrompt(id: string, operateBookType: OperateBookType, type: BookType): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
    try {
      let bookTaskDetailIds: string[] = []
      let bookTaskDetails = [] as Book.SelectBookTaskDetail[]
      if (operateBookType == OperateBookType.BOOKTASK) {
        bookTaskDetails = await this.bookServiceBasic.GetBookTaskDetailData({
          bookTaskId: id
        })
      } else if (operateBookType == OperateBookType.BOOKTASKDETAIL) {
        let bookTaskDetail = await this.bookServiceBasic.GetBookTaskDetailDataById(id);
        bookTaskDetails = [bookTaskDetail]
      } else if (operateBookType == OperateBookType.UNDERBOOKTASK) {
        // 下生图
        let tempBooktaskDetail = await this.bookServiceBasic.GetBookTaskDetailDataById(id);
        let tempBooktaskDetails = await this.bookServiceBasic.GetBookTaskDetailData({
          bookTaskId: tempBooktaskDetail.bookTaskId
        });
        for (let i = 0; i < tempBooktaskDetails.length; i++) {
          const element = tempBooktaskDetails[i];
          if (tempBooktaskDetail.no <= element.no) {
            bookTaskDetails.push(element)
          }
        }
      } else {
        throw new Error("未知的操作模式，请检查");
      }

      for (let i = 0; i < bookTaskDetails.length; i++) {
        const element = bookTaskDetails[i];
        bookTaskDetailIds.push(element.id)
      }
      if (bookTaskDetailIds.length <= 0) {
        throw new Error("没有需要反推的数据，请检查")
      }
      await this.AddReversePromptTask(bookTaskDetailIds, type);
      return successMessage(null, "添加反推任务成功", 'ReverseBook_AddReversePrompt')
    } catch (error) {
      return errorMessage("添加反推任务失败，错误信息如下：" + error.message, "ReverseBook_AddReversePrompt")
    }
  }

  /**
   * 添加单句生图任务
   * @param bookTaskDetailId 小说单个分镜的ID
   * @param type 反推的类型
   * @returns 
   */
  async AddReversePromptTask(bookTaskDetailIds: string[], type: BookType): Promise<void> {
    try {
      for (let index = 0; index < bookTaskDetailIds.length; index++) {
        const bookTaskDetailId = bookTaskDetailIds[index];
        let bookTaskDetail = await this.bookServiceBasic.GetBookTaskDetailDataById(bookTaskDetailId)
        let book = await this.bookServiceBasic.GetBookDataById(bookTaskDetail.bookId)
        // 是不是又额外的类型，没有的话试用小说的类型
        let task_type = undefined as BookBackTaskType
        if (!type) {
          type = book.type
        }
        switch (type) {
          case BookType.MJ_REVERSE:
            task_type = BookBackTaskType.MJ_REVERSE
            break;
          case BookType.SD_REVERSE:
            task_type = BookBackTaskType.SD_REVERSE
            break;
          default:
            throw new Error("暂不支持的推理类型")
        }
        // 添加后台任务
        await this.bookServiceBasic.AddBookBackTask(book.id, task_type, TaskExecuteType.AUTO, bookTaskDetail.bookTaskId, bookTaskDetail.id, DEFINE_STRING.BOOK.REVERSE_PROMPT_RETURN
        );
        // 添加返回日志
        await this.taskScheduler.AddLogToDB(book.id, book.type, `添加 ${task_type} 反推任务成功`, bookTaskDetail.bookTaskId, LoggerStatus.SUCCESS)
      }
    } catch (error) {
      throw error
    }
  }

  // TODO SD反推，待重写
  async SDReversePrompt(task: TaskModal.Task): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
    console.log("SDReversePrompt", '')
    return successMessage(null, "", "ReverseBook_SDReversePrompt")
  }

  /**
   * 执行单句推理的任务
   * @param task 任务处理
   */
  async SingleReversePrompt(task: TaskModal.Task): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
    try {
      // 开始处理
      let res = undefined
      switch (task.type) {
        case BookBackTaskType.MJ_REVERSE:
          res = await this.mjOpt.MJImage2Text(task);
          break
        case BookBackTaskType.SD_REVERSE:
          res = await this.SDReversePrompt(task);
          break
        default:
          throw new Error("未知的任务类型")
      }
      // 正常返回信息处理
      return successMessage(null, `${task.name} _ ${task.id} 反推任务完成`, 'ReverseBook_SingleReversePrompt')
    } catch (error) {
      // 进行错误信息处理
      return errorMessage("单个反推失败，错误信息如下：" + error.message, "ReverseBook_SingleReversePrompt")
    }
  }
  //#endregion

  //#region  提示词相关操作

  /**
   * 合并提示词的入口函数
   * @param id 合并的ID，可以是小说任务ID，也可以是小说分镜ID
   * @param operateBookType 操作的类型
   * @param type 合并的类型（MJ。SD 。D3）
   * @returns 
   */
  async MergePrompt(id: string, type: PromptMergeType, operateBookType: OperateBookType): Promise<GeneralResponse.ErrorItem | GeneralResponse.SuccessItem> {
    try {
      if (type == PromptMergeType.MJ_MERGE) {
        return await this.mjOpt.MergePrompt(id, operateBookType);
      } else if (type == PromptMergeType.SD_MERGE) {
        return await this.sdOpt.MergePrompt(id, operateBookType)
      } else {
        throw new Error("未知的合并模式，请检查")
      }
    } catch (error) {
      return errorMessage("合并提示词失败，错误信息如下：" + error.message, "ReverseBook_MergePrompt")
    }
  }

  //#endregion
}