import { WritableStreamBuffer } from 'stream-buffers'
import { createLLM } from './llm'
import { uploadDoc } from './upload'
import { genChapters } from './chapter'
import { Logger } from '@nestjs/common'
import { HumanMessage, SystemMessage } from 'langchain/schema'
import * as NumberToChinese from '@siakhooi/number-to-chinese-words'
import { createAbstract } from './abstract'

function chapters2md(chapters, sub = []) {
    let md = ''
    const list: any[] = []
    function _update(_data: any, level: number) {
        _data.forEach((e: any) => {
            list.push({
                title: e.title,
                level
            })
            if (e.children) _update(e.children, level + 1)
        })
    }
    _update(chapters, 0)
    const levels = [0, 0, 0, 0]
    for (const child of list) {
        const level = child.level
        levels[level] = levels[level] + 1
        for (let i = level + 1; i < levels.length; i++) {
            levels[i] = 0
        }
        let title

        if (sub.length) {
            title = sub.concat(levels.slice(0, level + 1)).join('.') + ' '
        } else {
            if (level === 0 && !sub.length) {
                title = ''
            } else {
                title = levels.slice(0, level + 1).join('.') + ' '
            }
        }

        md += `${'    '.repeat(level)}- ${title}${child.title}\n`
    }
    return md
}

export async function genGeneralPaper1(major: string, title: string, desc: string, chapters) {
    const tools = createLLM('gpt-3.5-turbo')
    const { llm, usage, conversation } = tools
    const pRetry = (await eval("import('p-retry')")).default

    chapters = chapters || (await genChapters(tools, major, title, desc))
    Logger.log(JSON.stringify(chapters, null, 4))

    let paper = new WritableStreamBuffer()
    const startTime = Date.now()

    paper.write(`${await createAbstract(tools, title, desc)}\n\n`)

    for (let i = 0; i < chapters.length; i++) {
        const chapter = chapters[i]
        const sysPrompt = `你是一个擅长写${major}专业毕业论文的专家。
你的任务是写一篇标题为"${title}"的论文。
全文按照以下章节组织：
${chapters2md(chapters)}

你的语气应正式，内容应适合和吸引一般读者。
你要严格按照markdown格式返回内容。
不要擅作主张输出你对该内容的总结或者介绍等其他文字，不需要出现介绍下一章或其他承上启下的文字。
你要严格遵守我的指令，按照我的指定的章节和小节完成论文的写作。
请直接开始写作，不要对内容作任何解释。`
        paper.write('\n\n# ' + chapter.title)
        if (chapter.children && chapter.children.length) {
            // for (let j = 0; j < chapter.children.length; j++) {
            //     const section = chapter.children[j]
            //     paper.write(`\n\n## ${i + 1}.${j + 1} ${section.title}`)
            //     let prompt = `第${i + 1}章: # ${chapter.title}`
            //     prompt += `\n## ${i + 1}.${j + 1} ${section.title} (请注意只需要编写这一小节的内容):`
            //     const { content } = await llm.call([new SystemMessage(sysPrompt), new HumanMessage(prompt)])
            //     paper.write('\n\n' + content)
            // }

            const content = (
                await Promise.all(
                    chapter.children.map(async (section, j) => {
                        let prompt
                        if (section.children?.length) {
                            prompt = `${chapters2md(section.children, [i + 1, j + 1])}
请按照以上结构编写第${i + 1}章: # ${chapter.title}
## ${i + 1}.${j + 1} ${section.title} (请注意只需要编写这一小节的内容):`
                        } else {
                            prompt = `第${i + 1}章: # ${chapter.title}
## ${i + 1}.${j + 1} ${section.title} (请注意只需要编写这一小节的内容):`
                        }
                        const { content } = await llm.call([new SystemMessage(sysPrompt), new HumanMessage(prompt)])
                        return `## ${i + 1}.${j + 1} ${section.title}\n\n` + content
                    })
                )
            ).join('\n\n')
            paper.write('\n\n' + content)
        } else {
            let prompt = `第${i + 1}章: # ${chapter.title}`
            if (chapter.title.includes('参考文献')) {
                prompt += `（请给出一些文献来支持我的论点和方法，请给出文献的标题，作者，年份，摘要和关键词)`
            }
            const { content } = await llm.call([new SystemMessage(sysPrompt), new HumanMessage(prompt)])
            // if (!content.trim().split('\n')[0].includes(chapter.title)) {
            //     paper.write(`\n\n# ${chapter.title}`)
            // }
            paper.write('\n\n' + content)
        }
    }

    let content = paper.getContentsAsString('utf8').split('\n')
    for (let i = 0; i < content.length; i++) {
        const match = content[i].match(/(^#+)([^ #]+)/)
        if (match) {
            content[i] = match[1] + ' ' + match[2]
        }
        if (/^#+.*$/.test(content[i])) {
            if (content[i - 1] && content[i - 1].length) {
                content[i] = '\n' + content[i]
            }
            if (content[i + 1] && content[i + 1].length) {
                content[i] += '\n'
            }
        }
    }
    content = content.join('\n')

    const duration = (Date.now() - startTime) / 1000
    const fileUrl = await uploadDoc(title, content)
    return {
        content,
        duration,
        tokenUsage: usage,
        fileUrl
    }
}
