import { HumanMessage } from 'langchain/schema'
import { StructuredOutputParser } from 'langchain/output_parsers'
import { z } from 'zod'
import { uploadUml } from './upload'
import { LLMTools } from './types'

async function genErNames(tools: LLMTools, title, desc) {
    const { llm, usage, conversation } = tools
    const scheme = StructuredOutputParser.fromZodSchema(
        z.array(
            z.object({
                title: z.string().describe('ER图名称')
            })
        )
    )

    const res = await llm.call([
        new HumanMessage(`${scheme.getFormatInstructions()}
------
你是一个计算机专业擅长写毕业论文的专家，你的任务是帮我的毕业设计撰写一篇论文
我的毕业设计项目是${title}
这个项目包含以下功能：
${desc}
现在我们需要开始为这个系统设计数据库，首先我们分析一下系统中有哪些实体，我们需要绘制以下五个主要ER图：`)
    ])
    return await scheme.parse(res.content)
}

async function plotER(tools, title, desc, name) {
    const { llm, usage, conversation } = tools
    const { chain } = conversation(`你是一个计算机专业擅长写毕业论文的专家，你的任务是帮我的毕业设计撰写一篇论文
我的毕业设计项目是${title}
这个项目包含以下功能：
${desc}`)
    const { response: descr } = await chain.call({
        input: `现在我们需要在论文中附上一个ER图：${name}，请你先用文字描述一下这个实体关系`
    })
    let { response: code } = await chain.call({
        input: `现在根据文字描述将它转换为plantuml语言
这是一个plantuml描述ER图的示例：
\`\`\`
@startuml 示例ER图
entity 客户  {
name
}

entity 订单  {
date
}

entity 订单详情 {
price
}

客户 -right-o{ 订单
订单 ||-right-|{ 订单详情

@enduml
\`\`\``
    })
    code = /\`\`\`(?:plantuml)?([\s\S]+)\`\`\`/.exec(code)[1]
    const url = await uploadUml(code)
    return {
        name,
        desc: descr,
        url
    }
}

export async function createER(tools: LLMTools, title, desc) {
    const pRetry = (await eval("import('p-retry')")).default
    const list = await pRetry(() => genErNames(tools, title, desc), { retries: 5 })
    return await Promise.all(list.map((item) => pRetry(() => plotER(tools, title, desc, item.title), { retries: 5 })))
}
