import loggerIns from "@egg/logger"
import { JsonOutputParser } from "@langchain/core/output_parsers"
import { PromptTemplate } from "@langchain/core/prompts"

import { adjustTimeRange, getSpecificTime, getTimeRange } from "../time"
import { getLangfuse, getModel } from "./base"

/**
 * 调用LLM模型
 * @param promptName 提示Key
 * @param variables 变量
 * @param requestId 请求ID
 * @param temperature 温度
 * @returns
 */
const invoke = async (
  promptName: string,
  variables: Record<string, any>,
  requestId: string,
  temperature = 0,
  jsonMode = false
) => {
  const logger = loggerIns.child({ requestId })
  const attemptInvoke = async () => {
    const { langfuse, langfuseHandler } = await getLangfuse("invoke", requestId)
    const prompt = await langfuse.getPrompt(promptName)

    const langchainTextPrompt = PromptTemplate.fromTemplate(
      prompt.getLangchainPrompt()
    ).withConfig({
      metadata: { langfusePrompt: prompt },
    })

    const chain = langchainTextPrompt.pipe(await getModel(temperature))

    if (jsonMode) {
      chain.pipe(new JsonOutputParser())
    }

    const { content } = await chain.invoke(variables, {
      callbacks: [langfuseHandler],
    })

    return content
  }

  let result
  let attempts = 0
  do {
    try {
      result = await attemptInvoke()
      break
    } catch (e: any) {
      logger.error(`调用LLM模型失败`, { promptName, error: e.message })
      attempts++
    }
  } while (attempts < 3)

  if (!result) {
    logger.error("多次调用LLM模型失败", { promptName, attempts: 3 })
    return ""
  }

  return result
}

/**
 * 时间解析器
 * @param userInput
 * @param requestId
 * @returns
 */
const timeParser = async (userInput: string, requestId: string) => {
  const logger = loggerIns.child({ requestId })
  const time = new Date().toLocaleString("zh-CN", { timeZone: "Asia/Shanghai" })
  const weekDay = `星期${"日一二三四五六"[new Date().getDay()]}`
  const invokeParser = async () => {
    try {
      const res = (await invoke(
        "timeParser",
        {
          time,
          weekDay,
          userInput,
          lastWeekStart: getSpecificTime("lastMonday", 0),
          lastWeekEnd: getSpecificTime("lastSunday", 24),
          yesterdayAfternoonStart: getSpecificTime("yesterday", 12),
          yesterdayAfternoonEnd: getSpecificTime("yesterday", 18),
          threeDayStart: getSpecificTime("twoDaysAgo", 0),
          threeDayEnd: getSpecificTime("today", 24),
        },
        requestId
      )) as string
      return JSON.parse(res.replaceAll("`", ""))
    } catch (e: any) {
      logger.error("时间解析失败", { userInput, error: e.message })
      // 如果解析失败，则返回空字符串
      return { s: "", e: "" }
    }
  }

  const validateResult = (result: any) => {
    const regex = /^\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}$/
    return regex.test(result.s) && regex.test(result.e)
  }

  let result
  let attempts = 0
  do {
    result = await invokeParser()
    attempts++
  } while (!validateResult(result) && attempts < 3)

  // 如果解析失败，则返回过去三天的时间范围
  if (!validateResult(result)) {
    return getTimeRange("threeDays")
  }

  // 解析成功，调整时间范围
  return adjustTimeRange(result.s, result.e)
}

const llm = {
  timeParser,
  invoke,
}

export default llm