import { OptionRealmService } from '@/define/db/service/optionService'
import { OptionKeyName } from '@/define/enum/option'
import { SettingModal } from '@/define/model/setting'
import { TaskModal } from '@/define/model/task'
import { optionSerialization } from '../option/optionSerialization'
import {
  CheckFileOrDirExist,
  CheckFolderExistsOrCreate,
  CopyFileOrFolder
} from '@/define/Tools/file'
import fs from 'fs'
import { ValidateJson } from '@/define/Tools/validate'
import axios from 'axios'
import { isEmpty } from 'lodash'
import { BookBackTaskStatus, BookTaskStatus, OperateBookType } from '@/define/enum/bookEnum'
import { SendReturnMessage } from '@/public/generalTools'
import { Book } from '@/define/model/book/book'
import { MJAction } from '@/define/enum/mjEnum'
import { ImageGenerateMode } from '@/define/data/mjData'
import path from 'path'
import { getProjectPath } from '../option/optionCommonService'
import { SDServiceHandle } from './sdServiceHandle'

export class ComfyUIServiceHandle extends SDServiceHandle {
  constructor() {
    super()
  }

  ComfyUIImageGenerate = async (task: TaskModal.Task) => {
    try {
      await this.InitSDBasic()
      let comfyUISettingCollection = await this.GetComfyUISetting()
      let bookTaskDetail = await this.bookTaskDetailService.GetBookTaskDetailDataById(
        task.bookTaskDetailId as string
      )
      if (bookTaskDetail == null) {
        throw new Error('未找到对应的小说分镜')
      }
      let book = await this.bookService.GetBookDataById(bookTaskDetail.bookId as string)
      if (book == null) {
        throw new Error('未找到对应的小说')
      }
      let bookTask = await this.bookTaskService.GetBookTaskDataById(
        bookTaskDetail.bookTaskId as string
      )
      if (bookTask == null) {
        throw new Error('未找到对应的小说任务')
      }

      // 调用方法合并提示词
      let mergeRes = await this.MergeSDPrompt(
        task.bookTaskDetailId as string,
        OperateBookType.BOOKTASKDETAIL
      )
      if (mergeRes.code == 0) {
        throw new Error(mergeRes.message)
      }
      // 获取提示词
      bookTaskDetail.prompt = mergeRes.data[0].prompt

      let prompt = bookTaskDetail.prompt
      let negativePrompt = comfyUISettingCollection.comfyuiSimpleSetting.negativePrompt

      // 开始组合请求体
      let body = await this.GetComfyUIAPIBody(
        prompt ?? '',
        negativePrompt ?? '',
        comfyUISettingCollection.comfyuiSelectedWorkflow.workflowPath
      )

      // 开始发送请求
      let resData = await this.SubmitComfyUIImagine(body, comfyUISettingCollection)

      // 修改任务状态
      await this.bookTaskDetailService.ModifyBookTaskDetailById(task.bookTaskDetailId as string, {
        status: BookTaskStatus.IMAGE
      })
      this.taskListService.UpdateTaskStatus({
        id: task.id as string,
        status: BookBackTaskStatus.RUNNING
      })
      SendReturnMessage(
        {
          code: 1,
          message: '任务已提交',
          id: task.bookTaskDetailId as string,
          data: {
            status: 'submited',
            message: '任务已提交',
            id: task.bookTaskDetailId as string
          } as any
        },
        task.messageName as string
      )

      await this.FetchImageTask(
        task,
        resData.prompt_id,
        book,
        bookTask,
        bookTaskDetail,
        comfyUISettingCollection
      )
    } catch (error: any) {
      let errorMsg = 'ComfyUI 生图失败，失败信息如下：' + error.toString()
      this.taskListService.UpdateTaskStatus({
        id: task.id as string,
        status: BookBackTaskStatus.FAIL,
        errorMessage: errorMsg
      })
      await this.bookTaskDetailService.UpdateBookTaskDetailMjMessage(
        task.bookTaskDetailId as string,
        {
          mjApiUrl: '',
          progress: 0,
          category: ImageGenerateMode.ComfyUI,
          imageClick: '',
          imageShow: '',
          messageId: '',
          action: MJAction.IMAGINE,
          status: 'error',
          message: errorMsg
        }
      )

      SendReturnMessage(
        {
          code: 0,
          message: errorMsg,
          id: task.bookTaskDetailId as string,
          data: {
            status: 'error',
            message: errorMsg,
            id: task.bookTaskDetailId
          }
        },
        task.messageName as string
      )
      throw error
    }
  }

  //#region  获取ComfyUI的设置
  /**
   * 获取ComfyUI的设置
   * @returns
   */
  private async GetComfyUISetting(): Promise<SettingModal.ComfyUISettingCollection> {
    let result = {} as SettingModal.ComfyUISettingCollection
    let optionRealmService = await OptionRealmService.getInstance()
    let comfyuiSimpleSettingOption = optionRealmService.GetOptionByKey(
      OptionKeyName.SD.ComfyUISimpleSetting
    )
    result['comfyuiSimpleSetting'] = optionSerialization<SettingModal.ComfyUISimpleSettingModel>(
      comfyuiSimpleSettingOption,
      '设置 -> ComfyUI 设置'
    )

    let comfyuiWorkFlowSettingOption = optionRealmService.GetOptionByKey(
      OptionKeyName.SD.ComfyUIWorkFlowSetting
    )

    let comfyuiWorkFlowList = optionSerialization<SettingModal.ComfyUIWorkFlowSettingModel[]>(
      comfyuiWorkFlowSettingOption,
      '设置 -> ComfyUI 设置'
    )
    result['comfyuiWorkFlowSetting'] = comfyuiWorkFlowList

    if (comfyuiWorkFlowList.length <= 0) {
      throw new Error('ComfyUI的工作流设置为空，请检查是否正确设置！！')
    }

    // 获取选中的工作流
    let selectedWorkflow = comfyuiWorkFlowList.find(
      (item) => item.id == result.comfyuiSimpleSetting.selectedWorkflow
    )
    if (selectedWorkflow == null) {
      throw new Error('未找到选中的工作流，请检查是否正确设置！！')
    }

    // 判断工作流对应的文件是不是存在
    if (!(await CheckFileOrDirExist(selectedWorkflow.workflowPath))) {
      throw new Error('本地未找到选中的工作流文件地址，请检查是否正确设置！！')
    }
    result['comfyuiSelectedWorkflow'] = selectedWorkflow

    return result
  }

  //#endregion

  //#region 组合ComfyUI的请求体
  /**
   * 组合ComfyUI的请求体
   * @param prompt 正向提示词
   * @param negativePrompt 反向提示词
   * @param workflowPath 工作流地址
   */
  private async GetComfyUIAPIBody(
    prompt: string,
    negativePrompt: string,
    workflowPath: string
  ): Promise<string> {
    let jsonContentString = await fs.promises.readFile(workflowPath, 'utf-8')

    if (!ValidateJson(jsonContentString)) {
      throw new Error('工作流文件内容不是有效的JSON格式，请检查是否正确设置！！')
    }

    let jsonContent = JSON.parse(jsonContentString)
    // 判断是否是对象
    if (jsonContent !== null && typeof jsonContent === 'object' && !Array.isArray(jsonContent)) {
      // 遍历对象属性
      for (const key in jsonContent) {
        let element = jsonContent[key]
        if (element && element.class_type === 'CLIPTextEncode') {
          if (element._meta?.title === '正向提示词') {
            jsonContent[key].inputs.text = prompt
          }
          if (element._meta?.title === '反向提示词') {
            jsonContent[key].inputs.text = negativePrompt
          }
        }
        if (element && element.class_type === 'KSampler') {
          const crypto = require('crypto')
          const buffer = crypto.randomBytes(8)
          let seed = BigInt('0x' + buffer.toString('hex'))

          jsonContent[key].inputs.seed = seed.toString()
        } else if (element && element.class_type === 'KSamplerAdvanced') {
          const crypto = require('crypto')
          const buffer = crypto.randomBytes(8)
          let seed = BigInt('0x' + buffer.toString('hex'))

          jsonContent[key].inputs.noise_seed = seed.toString()
        }
      }
    } else {
      throw new Error('工作流文件内容不是有效的JSON对象格式，请检查是否正确设置！！')
    }
    let result = JSON.stringify({
      prompt: jsonContent
    })
    return result
  }

  //#endregion

  //#region 提交ComfyUI生成图片任务

  private async SubmitComfyUIImagine(
    body: string,
    comfyUISettingCollection: SettingModal.ComfyUISettingCollection
  ): Promise<any> {
    let url = comfyUISettingCollection.comfyuiSimpleSetting.requestUrl?.replace(
      'localhost',
      '127.0.0.1'
    )
    if (url.endsWith('/')) {
      url = url + 'api/prompt'
    } else {
      url = url + '/api/prompt'
    }
    var config = {
      method: 'post',
      url: url,
      headers: {
        'User-Agent': 'Apifox/1.0.0 (https://apifox.com)',
        'Content-Type': 'application/json'
      },
      data: body
    }

    let res = await axios(config)
    let resData = res.data
    // 判断是不是失败
    if (resData.error) {
      let errorNode = ''
      if (resData.node_errors) {
        for (const key in resData.node_errors) {
          errorNode += key + ', '
        }
      }
      let msg = '错误信息：' + resData.error.message + '错误节点：' + errorNode
      throw new Error(msg)
    }
    // 没有错误 判断是不是成功
    if (resData.prompt_id && !isEmpty(resData.prompt_id)) {
      // 成功
      return resData
    } else {
      throw new Error('未知错误，未获取到请求ID，请检查是否正确设置！！')
    }
  }

  //#endregion

  //#region 获取出图任务

  async FetchImageTask(
    task: TaskModal.Task,
    promptId: string,
    book: Book.SelectBook,
    bookTask: Book.SelectBookTask,
    bookTaskDetail: Book.SelectBookTaskDetail,
    comfyUISettingCollection: SettingModal.ComfyUISettingCollection
  ) {
    while (true) {
      try {
        let resData = await this.GetComfyUIImageTask(promptId, comfyUISettingCollection)

        // 判断他的状态是不是成功
        if (resData.status == 'error') {
          // 生图失败
          await this.bookTaskDetailService.ModifyBookTaskDetailById(
            task.bookTaskDetailId as string,
            {
              status: BookTaskStatus.IMAGE_FAIL
            }
          )
          let errorMsg = `MJ生成图片失败，失败信息如下：${resData.message}`
          this.bookTaskDetailService.UpdateBookTaskDetailMjMessage(
            task.bookTaskDetailId as string,
            {
              mjApiUrl: comfyUISettingCollection.comfyuiSimpleSetting.requestUrl,
              progress: 100,
              category: ImageGenerateMode.ComfyUI,
              imageClick: '',
              imageShow: '',
              messageId: promptId,
              action: MJAction.IMAGINE,
              status: 'error',
              message: errorMsg
            }
          )
          this.taskListService.UpdateTaskStatus({
            id: task.id as string,
            status: BookBackTaskStatus.FAIL,
            errorMessage: errorMsg
          })
          SendReturnMessage(
            {
              code: 0,
              message: errorMsg,
              id: task.bookTaskDetailId as string,
              data: {
                status: 'error',
                message: errorMsg,
                id: task.bookTaskDetailId
              }
            },
            task.messageName as string
          )
          return
        } else if (resData.status == 'in_progress') {
          // 生图中
          this.bookTaskDetailService.UpdateBookTaskDetailMjMessage(
            task.bookTaskDetailId as string,
            {
              mjApiUrl: comfyUISettingCollection.comfyuiSimpleSetting.requestUrl,
              progress: 0,
              category: ImageGenerateMode.ComfyUI,
              imageClick: '',
              imageShow: '',
              messageId: promptId,
              action: MJAction.IMAGINE,
              status: 'running',
              message: '任务正在执行中'
            }
          )

          SendReturnMessage(
            {
              code: 1,
              message: 'running',
              id: task.bookTaskDetailId as string,
              data: {
                status: 'running',
                message: '任务正在执行中',
                id: task.bookTaskDetailId
              }
            },
            task.messageName as string
          )
        } else {
          let res = await this.DownloadFileUrl(
            resData.imageNames,
            comfyUISettingCollection,
            book,
            bookTask,
            bookTaskDetail
          )
          let projectPath = await getProjectPath()
          // 修改数据库数据
          // 修改数据库
          await this.bookTaskDetailService.ModifyBookTaskDetailById(bookTaskDetail.id as string, {
            outImagePath: path.relative(projectPath, res.outImagePath),
            subImagePath: res.subImagePath.map((item) => path.relative(projectPath, item))
          })

          this.taskListService.UpdateTaskStatus({
            id: task.id as string,
            status: BookBackTaskStatus.DONE
          })

          this.bookTaskDetailService.UpdateBookTaskDetailMjMessage(
            task.bookTaskDetailId as string,
            {
              mjApiUrl: comfyUISettingCollection.comfyuiSimpleSetting.requestUrl,
              progress: 100,
              category: ImageGenerateMode.ComfyUI,
              imageClick: '',
              imageShow: '',
              messageId: promptId,
              action: MJAction.IMAGINE,
              status: 'success',
              message: 'ComfyUI 生成图片成功'
            }
          )

          SendReturnMessage(
            {
              code: 1,
              message: 'ComfyUI 生成图片成功',
              id: task.bookTaskDetailId as string,
              data: {
                status: 'success',
                message: 'ComfyUI 生成图片成功',
                id: task.bookTaskDetailId,
                outImagePath: res.outImagePath + '?t=' + new Date().getTime(),
                subImagePath: res.subImagePath.map((item) => item + '?t=' + new Date().getTime())
              }
            },
            task.messageName as string
          )
          break
        }
        await new Promise((resolve) => setTimeout(resolve, 3000))
      } catch (error) {
        throw error
      }
    }
  }

  //#endregion

  //#region 获取comfyui出图任务

  /**
   * 获取ComfyUI出图任务
   * @param promptId
   * @param comfyUISettingCollection
   */
  private async GetComfyUIImageTask(
    promptId: string,
    comfyUISettingCollection: SettingModal.ComfyUISettingCollection
  ): Promise<any> {
    if (isEmpty(promptId)) {
      throw new Error('未获取到请求ID，请检查是否正确设置！！')
    }
    if (isEmpty(comfyUISettingCollection.comfyuiSimpleSetting.requestUrl)) {
      throw new Error('未获取到ComfyUI的请求地址，请检查是否正确设置！！')
    }

    let url = comfyUISettingCollection.comfyuiSimpleSetting.requestUrl?.replace(
      'localhost',
      '127.0.0.1'
    )
    if (url.endsWith('/')) {
      url = url + 'api/history'
    } else {
      url = url + '/api/history'
    }

    var config = {
      method: 'get',
      url: `${url}/${promptId}`,
      headers: {
        'User-Agent': 'Apifox/1.0.0 (https://apifox.com)'
      }
    }

    let res = await axios.request(config)
    let resData = res.data
    // 判断状态是失败还是成功
    let data = resData[promptId]
    if (data == null) {
      // 还在执行中 或者是任务不存在
      return {
        progress: 0,
        status: 'in_progress',
        message: '任务正在执行中'
      }
    }
    let completed = data.status?.completed
    let outputs = data.outputs
    if (completed && outputs) {
      let imageNames: string[] = []
      for (const key in outputs) {
        let outputNode = outputs[key]
        if (outputNode && outputNode?.images && outputNode?.images.length > 0) {
          for (let i = 0; i < outputNode?.images.length; i++) {
            const element = outputNode?.images[i]
            imageNames.push(element.filename as string)
          }
        }
      }
      return {
        progress: 100,
        status: 'success',
        imageNames: imageNames
      }
    } else {
      return {
        progress: 0,
        status: 'error',
        message: '生图失败，详细失败信息看启动器控制台'
      }
    }
  }

  //#endregion

  //#region 请求下载图片

  /**
   * 请求下载对应的图片
   * @param url
   * @param path
   */
  private async DownloadFileUrl(
    imageNames: string[],
    comfyUISettingCollection: SettingModal.ComfyUISettingCollection,
    book: Book.SelectBook,
    bookTask: Book.SelectBookTask,
    bookTaskDetail: Book.SelectBookTaskDetail
  ): Promise<{
    outImagePath: string
    subImagePath: string[]
  }> {
    let url = comfyUISettingCollection.comfyuiSimpleSetting.requestUrl?.replace(
      'localhost',
      '127.0.0.1'
    )
    if (url.endsWith('/')) {
      url = url + 'api/view'
    } else {
      url = url + '/api/view'
    }
    let outImagePath = ''
    let subImagePath: string[] = []
    for (let i = 0; i < imageNames.length; i++) {
      const imageName = imageNames[i]

      var config = {
        method: 'get',
        url: `${url}?filename=${imageName}&nocache=${Date.now()}`,
        headers: {
          'User-Agent': 'Apifox/1.0.0 (https://apifox.com)'
        },
        responseType: 'arraybuffer' as 'arraybuffer' // 明确指定类型
      }

      let res = await axios.request(config)

      // 检查响应状态和类型
      console.log(`图片下载状态: ${res.status}, 内容类型: ${res.headers['content-type']}`)

      // 确保得到的是图片数据
      if (!res.headers['content-type']?.includes('image/')) {
        console.error(`响应不是图片: ${res.headers['content-type']}`)
        continue
      }

      let resData = res.data
      console.log(resData)

      let subImageFolderPath = path.join(
        bookTask.imageFolder as string,
        `subImage/${bookTaskDetail.name}`
      )
      await CheckFolderExistsOrCreate(subImageFolderPath)
      let outputFolder = bookTask.imageFolder as string
      await CheckFolderExistsOrCreate(outputFolder)
      let inputFolder = path.join(book.bookFolderPath as string, 'tmp/input')
      await CheckFolderExistsOrCreate(inputFolder)

      // 包含info信息的图片地址
      let infoImgPath = path.join(subImageFolderPath, `${new Date().getTime()}_${i}.png`)

      // 直接将二进制数据写入文件
      await fs.promises.writeFile(infoImgPath, Buffer.from(resData))

      if (i == 0) {
        // 复制到对应的文件夹里面
        let outPath = path.join(outputFolder, `${bookTaskDetail.name}.png`)
        await CopyFileOrFolder(infoImgPath, outPath)
        outImagePath = outPath
      }
      subImagePath.push(infoImgPath as string)
    }
    console.log(outImagePath)
    console.log(subImagePath)

    // 将获取的数据返回
    return {
      outImagePath: outImagePath,
      subImagePath: subImagePath
    }
  }

  //#endregion
}