import type { AbortManager } from "#Source/abort/index.ts"
import { fromWithAbortSignal } from "#Source/abort/index.ts"
import * as Aio from "#Source/aio/index.ts"
import { scheduleMacroTask } from "#Source/basic/index.ts"
import type { LoggerFriendly, LoggerFriendlyOptions } from "#Source/log/index.ts"
import { Logger } from "#Source/log/index.ts"
import { Dispatcher } from "#Source/orchestration/index.ts"
import { controllerFromEitherType, eitherToTuple } from "#Source/result/index.ts"
import { connectTube, Tube } from "#Source/tube/index.ts"

import type {
  BaseChatCompletion,
  ChatCompletionOptions,
  ChatCompletionResult,
  Completion,
} from "./chat-completion.ts"

export interface ChatCompletionAiOptions extends LoggerFriendlyOptions {
  chatCompletionInstanceList: BaseChatCompletion[]
  /**
   * @description The timeout for the first chunk of data to be received, in milliseconds.
   *
   * @default timeoutPerChunk | 30_000
   */
  timeoutFirstChunk?: number | undefined
  /**
   * @description The timeout for each chunk of data to be received, in milliseconds.
   *
   * @default 30_000
   */
  timeoutPerChunk?: number | undefined
  /**
   * @description The maximum number of tries for the request.
   *
   * @default 3
   */
  maxTries?: number | undefined
}
interface ResolvedChatCompletionAiOptions {
  chatCompletionInstanceList: BaseChatCompletion[]
  timeoutFirstChunk: number
  timeoutPerChunk: number
  maxTries: number
}
export class ChatCompletionAi implements LoggerFriendly {
  protected readonly options: ResolvedChatCompletionAiOptions

  readonly logger: Logger
  private dispatcher: Dispatcher<BaseChatCompletion>

  constructor(options: ChatCompletionAiOptions) {
    this.options = {
      chatCompletionInstanceList: options.chatCompletionInstanceList,
      timeoutFirstChunk: options.timeoutFirstChunk ?? options.timeoutPerChunk ?? 30_000,
      timeoutPerChunk: options.timeoutPerChunk ?? 30_000,
      maxTries: options.maxTries ?? 3,
    }

    this.logger = Logger.fromOptions(options).setDefaultName("ChatCompletionAi")

    this.dispatcher = new Dispatcher({
      itemList: this.options.chatCompletionInstanceList,
      logger: this.logger,
    })
  }

  async chatCompletion(options: ChatCompletionOptions): Promise<ChatCompletionResult> {
    const controller = controllerFromEitherType<ChatCompletionResult>()

    const completionTube = new Tube<Completion>({
      historyCount: Infinity,
      replayHistory: true,
    })

    scheduleMacroTask({
      task: async () => {
        await this.request(options, completionTube)
      },
    })

    return await controller.returnRight({
      completionTube,
    })
  }

  private async request(
    options: ChatCompletionOptions,
    targetCompletionTube: Tube<Completion>,
  ): Promise<void> {
    const { selector, destroy } = this.dispatcher.getSelector({
      filter: (item) => {
        return item.isSupportModel(options.model) === true
      },
    })
    targetCompletionTube.subscribeEndEvent({
      subscriber: () => {
        destroy()
      },
    })

    const latestCompletion: Completion = {
      content: {
        deltaList: [],
        total: "",
      },
      token: {
        deltaList: [],
        total: 0,
      },
    }

    const maxTries = options.maxTries ?? this.options.maxTries

    interface HandleTryErrorOptions {
      tryIndex: number
      abortManager: AbortManager
    }
    const handleTryError = async (options: HandleTryErrorOptions): Promise<void> => {
      const { tryIndex, abortManager } = options

      this.logger.log(
        `Handle try error, current failure count: ${tryIndex + 1}, max tries: ${maxTries}`,
      )

      if (abortManager.isAborted()) {
        this.logger.log("Chat completion aborted by invoker.")
        await targetCompletionTube.pushError(new Error("Chat completion aborted by invoker."))
        abortManager.abort("Chat completion aborted by invoker.")
        return
      } else {
        this.logger.log("Chat completion not aborted by invoker.")
      }
      // 清理资源
      abortManager.abort("Failed request")

      // 重置数据
      latestCompletion.content = Aio.applyDeltaToTextContent(latestCompletion.content, {
        type: "reset",
      })
      latestCompletion.token = Aio.applyDeltaToNumberContent(latestCompletion.token, {
        type: "reset",
      })
      await targetCompletionTube.pushData(structuredClone(latestCompletion))

      // 如果尝试次数未达到上限，则继续尝试
      const nextTryIndex = tryIndex + 1
      if (nextTryIndex < maxTries) {
        this.logger.log(`Not exceed maximum tries (${maxTries}) for chat completion, try again`)
        scheduleMacroTask({
          task: async () => {
            await startTry(nextTryIndex)
          },
          timeout: 500,
        })
      } else {
        this.logger.log(`Exceeded maximum tries (${maxTries}) for chat completion`)
        await targetCompletionTube.pushError(new Error("Exceeded maximum tries."))
      }
    }
    const startTry = async (tryIndex: number): Promise<void> => {
      this.logger.log(`Start try, tries ${tryIndex + 1} of ${maxTries}`)

      const abortManager = fromWithAbortSignal(options)
      const [getItemLeft, getItemRight] = eitherToTuple(await selector.getItem())
      if (getItemLeft !== undefined) {
        await handleTryError({ tryIndex, abortManager })
        return
      }
      const { item, markUnavailable } = getItemRight
      this.logger.log(`${item.logger.getName()} is selected from pool`)
      const [leftResult, rightResult] = eitherToTuple(
        await this.requestAndConstrainSpeed(
          {
            ...options,
            abortSignal: abortManager.abortSignal,
          },
          item,
        ),
      )
      if (leftResult !== undefined) {
        await handleTryError({ tryIndex, abortManager })
        return
      }
      const { completionTube: sourceCompletionTube } = rightResult
      sourceCompletionTube.subscribeEndEvent({
        subscriber: async (): Promise<void> => {
          if (sourceCompletionTube.isError()) {
            this.logger.log("Source completion tube ended with error")
            markUnavailable()
            await handleTryError({ tryIndex, abortManager })
          } else {
            this.logger.log("Source completion tube ended normally")
            await targetCompletionTube.end()
          }
        },
      })
      sourceCompletionTube.subscribeData({
        subscriber: async (data: Completion): Promise<void> => {
          await targetCompletionTube.pushData(data)
        },
      })
    }

    await startTry(0)
  }

  /**
   * @description 发起请求并对返回结果施加时间约束，当发生以下情况时，将被视为请求失败：
   * 1. 首个数据块超时；
   * 2. 相邻数据块超时；
   */
  private async requestAndConstrainSpeed(
    options: ChatCompletionOptions,
    instance: BaseChatCompletion,
  ): Promise<ChatCompletionResult> {
    const controller = controllerFromEitherType<ChatCompletionResult>()

    const [chatCompletionLeft, chatCompletionRight] = eitherToTuple(
      await instance.chatCompletion(options),
    )
    if (chatCompletionLeft !== undefined) {
      return await controller.returnLeft(chatCompletionLeft)
    }
    const { completionTube } = chatCompletionRight

    const newCompletionTube = new Tube<Completion>({ historyCount: Infinity, replayHistory: true })
    // 无异常的情况下，上游数据直接传递给下游即可
    connectTube(completionTube, newCompletionTube)

    const timeoutFirstChunk = options.timeoutFirstChunk ?? this.options.timeoutFirstChunk
    const timeoutPerChunk = options.timeoutPerChunk ?? this.options.timeoutPerChunk

    // 定义首个数据块超时的处理逻辑
    let hasReceivedFirstChunk = false
    const firstChunkTimeoutTimer: NodeJS.Timeout = setTimeout(() => {
      if (hasReceivedFirstChunk === false) {
        this.logger.addOnceTags(["ConstrainSpeed"]).log("First chunk timeout")
        void newCompletionTube.pushError(new Error("FirstChunkTimeout"))
      }
    }, timeoutFirstChunk)

    // 定义相邻数据块超时的处理逻辑
    let perChunkTimeoutTimer: NodeJS.Timeout | undefined
    completionTube.subscribeData({
      subscriber: () => {
        if (hasReceivedFirstChunk === false) {
          this.logger
            .addOnceTags(["ConstrainSpeed"])
            .log("Received first chunk, clear first chunk timeout timer")
          hasReceivedFirstChunk = true
          // Clear the timeout timer for the first chunk
          clearTimeout(firstChunkTimeoutTimer)
        } else {
          // Reset the timeout timer for each chunk
          clearTimeout(perChunkTimeoutTimer)
        }
        // Set a new timeout timer for the next chunk
        perChunkTimeoutTimer = setTimeout(() => {
          this.logger.addOnceTags(["ConstrainSpeed"]).log("Per chunk timeout")
          void newCompletionTube.pushError(new Error("PerChunkTimeout"))
        }, timeoutPerChunk)
      },
    })
    completionTube.subscribeEndEvent({
      subscriber: () => {
        this.logger.addOnceTags(["ConstrainSpeed"]).log("Source completion tube ended")
        clearTimeout(firstChunkTimeoutTimer)
        clearTimeout(perChunkTimeoutTimer)
      },
    })

    return await controller.returnRight({
      completionTube: newCompletionTube,
    })
  }
}