import { BadRequestException, Injectable, InternalServerErrorException, Logger } from '@nestjs/common'
import * as PdfParse from '@cyber2024/pdf-parse-fixed'
import { createHash } from 'crypto'
import { get_encoding, Tiktoken } from '@dqbd/tiktoken'
import { Configuration, OpenAIApi } from 'azure-openai'
import * as queue from 'fastq'
import { setTimeout } from 'timers/promises'
import * as dedent from 'dedent'
import { DataTypes, Sequelize } from 'sequelize'
import { ConfigService } from '@nestjs/config'
import { ChatEmbedding } from './entities/chat-embedding.entity'
import { VECTOR } from '../utils/pgvector'
import { ApiUserService } from '../api-users/api-user.service'
import { UsersService } from 'src/users/users.service'
import { Role } from 'src/model/role.enum'
import { ApiUser } from 'src/api-users/entities/api-user.entity'

function formatEmbedding(embedding: number[]) {
    return `[${embedding.join(', ')}]`
}

@Injectable()
export class ChatPdfService {
    private readonly tokenizer: Tiktoken
    private readonly openai: OpenAIApi
    private readonly sequelize: Sequelize
    constructor(
        // private readonly sysConfigService: SysConfigService,
        private readonly configService: ConfigService,
        private readonly apiUserService: ApiUserService,
        private readonly userService: UsersService
    ) {
        this.tokenizer = get_encoding('cl100k_base')
        this.openai = new OpenAIApi(
            new Configuration({
                apiKey: 'beb32e4625a94b65ba8bc0ba1688c4d2',
                // add azure info into configuration
                azure: {
                    apiKey: 'beb32e4625a94b65ba8bc0ba1688c4d2',
                    endpoint: 'https://zouma.openai.azure.com'
                }
            })
        )
        this.sequelize = new Sequelize({
            dialect: 'postgres',
            host: configService.get<string>('PG_HOST'),
            port: configService.get<number>('PG_PORT'),
            username: configService.get<string>('PG_USERNAME'),
            password: configService.get<string>('PG_PASSWORD'),
            database: configService.get<string>('PG_DATABASE'),
            // logging: (msg) => Logger.debug(msg, 'Sequelize')
            logging: false
        })
        ChatEmbedding.init(
            {
                id: {
                    primaryKey: true,
                    autoIncrement: true,
                    type: DataTypes.INTEGER
                },
                name: {
                    type: DataTypes.STRING
                },
                text: {
                    type: DataTypes.TEXT({
                        length: 'long'
                    })
                },
                num: {
                    type: DataTypes.INTEGER
                },
                embedding: {
                    type: new VECTOR(1536)
                }
            },
            { sequelize: this.sequelize }
        )
        this.sequelize.sync()
    }

    public async upload(file: Express.Multer.File, code: string) {
        const { originalname, buffer, mimetype } = file
        let md5 = this.calculateMD5(buffer)
        const res = await ChatEmbedding.findAll({
            where: {
                name: md5
            }
        })
        if (!!code) {
            md5 = code
        }
        if (res.length) {
            return {
                name: md5
            }
        }
        const pdf = await PdfParse(buffer)
        const contents = []
        let paragraph = ''
        pdf.text
            .trim()
            .split('\n')
            .forEach((line) => {
                line = line.trim()
                paragraph += line
                if (this.isFullSentence(line)) {
                    contents.push(paragraph)
                    paragraph = ''
                }
            })
        if (paragraph) {
            contents.push(paragraph)
        }

        const embeddings = await this.createEmbeddings(contents)
        Logger.log(
            `create embeddings finished, total token usage: ${embeddings.reduce((acc, cur) => acc + cur.token, 0)}`
        )
        let i = 0
        for (const item of embeddings) {
            try {
                await ChatEmbedding.create({
                    name: md5,
                    text: item.text,
                    num: i++,
                    embedding: formatEmbedding(item.embedding)
                })
            } catch (error) {
                Logger.error(error.message)
            }
        }
        return {
            name: md5
        }
    }

    public async firstUpload(file: Express.Multer.File) {
        const { originalname, buffer, mimetype } = file
        let md5 = this.calculateMD5(buffer)
        const res = await ChatEmbedding.findAll({
            where: {
                name: md5
            }
        })
        if (res.length) {
            return {
                name: md5
            }
        }
        const pdf = await PdfParse(buffer)
        const contents = []
        let paragraph = ''
        pdf.text
            .trim()
            .split('\n')
            .forEach((line) => {
                line = line.trim()
                paragraph += line
                if (this.isFullSentence(line)) {
                    contents.push(paragraph)
                    paragraph = ''
                }
            })
        if (paragraph) {
            contents.push(paragraph)
        }

        const embeddings = await this.createEmbeddings(contents)
        Logger.log(
            `create embeddings finished, total token usage: ${embeddings.reduce((acc, cur) => acc + cur.token, 0)}`
        )
        let i = 0
        for (const item of embeddings) {
            try {
                await ChatEmbedding.create({
                    name: md5,
                    text: item.text,
                    num: i++,
                    embedding: formatEmbedding(item.embedding)
                })
            } catch (error) {
                Logger.error(error.message)
            }
        }
        return {
            name: md5
        }
    }

    isFullSentence(str) {
        return /[.!？。！？…;；:：”’）】》」』〕〉》〗〞〟»"'\])}]+$/.test(str)
    }

    calculateMD5(buffer) {
        const hash = createHash('md5')
        hash.update(buffer)
        return hash.digest('hex')
    }

    async createEmbeddings(content: string[]) {
        const self = this
        const result = Array(content.length)
        async function worker(arg) {
            result[arg.index] = await self.getEmbedding(arg.text)
            Logger.log(`create embedding for ${arg.index + 1}/${content.length}`)
        }
        const q = queue.promise(worker, 32)
        content.forEach((text, index) => {
            q.push({
                text,
                index
            })
        })
        await q.drained()
        return result.filter(i=> i && i.text)
    }

    async getEmbedding(content: string, retry = 0) {
        try {
            const response = await this.openai.createEmbedding({
                model: 'embedding',
                input: content
            })
            return {
                text: content,
                embedding: response.data.data[0].embedding,
                token: response.data.usage.total_tokens
            }
        } catch (error) {
            if (retry < 3) {
                Logger.error(`fetchEmbedding error: ${error.message}, retry ${retry}`, 'fetchEmbedding')
                await setTimeout(2000)
                return await this.getEmbedding(content, retry + 1)
            }
            Logger.error(error.stack, 'fetchEmbedding')
            throw new InternalServerErrorException(error.message)
        }
    }

    async getKeywords(text: string) {
        try {
            const res = await this.openai.createChatCompletion({
                model: 'gpt35',
                messages: [
                    {
                        role: 'user',
                        content: `You need to extract keywords from the statement or question and return a series of keywords separated by commas.\ncontent: ${text}\nkeywords: `
                    }
                ]
            })
            return res.data.choices[0].message.content
        } catch (error) {
            Logger.error(error.message)
            if (error.response) {
                Logger.error(error.response.data)
            }
            throw new InternalServerErrorException(error.message)
        }
    }

    async searchEmbedding(name: string, embedding: number[]) {
        return await ChatEmbedding.findAll({
            where: {
                name
            },
            order: this.sequelize.literal(`embedding <-> '${formatEmbedding(embedding)}'`),
            limit: 100
        })
    }

    cutContext(context: string[]) {
        if (!context || !context.length) return []
        let max = 4096 - 1024
        for (let i = 0; i < context.length; i++) {
            max -= this.tokenizer.encode(context[i]).length
            if (max < 0) {
                return context.slice(0, i)
            }
        }
        return context
    }

    async getUser(userId: number) {
        return await this.userService.findById(userId)
    }

    async getSystemConfig(userId: number, q: string, name: string, code: string, sysMsg: string) {
        const user = await this.userService.findById(userId)
        if (!user.apiUserId) {
            return sysMsg
        }
        let apiUser = await this.apiUserService.findById(user.apiUserId)
        if (user.roles.includes(Role.User)) {
            code = apiUser.publicCode
        } else {
            code = apiUser.code
        }
        const keywords = await this.getKeywords(q)
        const { embedding: keywordEmbedding } = await this.getEmbedding(keywords)
        const context = this.cutContext((await this.searchEmbedding(code, keywordEmbedding)).map((item) => item.text))
        // if (!context || !context.length) {
        //     return {
        //         answer: '未找到相关内容'
        //     }
        // }
        return dedent`${apiUser.desc}'
        这是用户提出的问题:
        ${q}
        你只能根据用户的问题，以下面的内容为准结合你的身份进行回答:
        \`\`\`
        ${context.join('\n')}
        \`\`\`
        你要确保你的回答全部基于上面的问题和内容还有你的身份，如果无法找到答案，你需要根据给你设定的角色表述你不能回答这个问题。
        你只能用中文回答.`
    }
    async ask(q: string, name: string) {
        // const defSysMsg = (await this.sysConfigService.findByName('customer_system_message'))?.value
        const apiUser = await this.apiUserService.findByCode(name)
        const keywords = await this.getKeywords(q)
        const { embedding: keywordEmbedding } = await this.getEmbedding(keywords)
        const context = this.cutContext((await this.searchEmbedding(name, keywordEmbedding)).map((item) => item.text))
        if (!context || !context.length) {
            return {
                answer: '未找到相关内容'
            }
        }
        const content = dedent`${apiUser.desc}'
        根据相关性从高到低排序。 '
        这是用户提出的问题:
        ${q}
        你只能根据用户的问题，以下面的内容为准进行回答:
        \`\`\`
        ${context.join('\n')}
        \`\`\`
        你要确保你的回答全部基于上面的问题和内容，如果无法找到答案，你需要根据给你设定的角色表述你不能回答这个问题。
        You must use Chinese to respond.`
        try {
            const response = await this.openai.createChatCompletion({
                model: 'gpt35',
                messages: [
                    { role: 'user', content },
                    { role: 'user', content: q }
                ]
            })
            return { answer: response.data.choices[0].message.content }
        } catch (error) {
            Logger.error(error.message)
            if (error.response) {
                Logger.error(error.response.data)
            }
            throw new InternalServerErrorException(error.message)
        }
    }

    async customerAsk(q: string, name: string) {
        let apiUser = await this.apiUserService.findByCode(name)
        if (!apiUser) {
            throw new BadRequestException('not a enabled api user')
        }
        // const defSysMsg = (await this.sysConfigService.findByName('customer_system_message'))?.value
        const keywords = await this.getKeywords(q)
        const { embedding: keywordEmbedding } = await this.getEmbedding(keywords)
        const context = this.cutContext((await this.searchEmbedding(name, keywordEmbedding)).map((item) => item.text))
        if (!context || !context.length) {
            return {
                answer: '客服无法回答这个问题，换个试试吧'
            }
        }
        const content = dedent`${apiUser.desc}'
        根据相关性从高到低排序。 '
        这是用户提出的问题:
        ${q}
        你只能根据用户的问题，以下面的内容为准进行回答:
        \`\`\`
        ${context.join('\n')}
        \`\`\`
        你要确保你的回答全部基于上面的问题和内容，如果无法找到答案，你需要根据给你设定的角色表述你不能回答这个问题。
        You must use Chinese to respond.`
        try {
            const response = await this.openai.createChatCompletion({
                model: 'gpt35',
                messages: [
                    { role: 'user', content },
                    { role: 'user', content: q }
                ]
            })
            return { answer: response.data.choices[0].message.content }
        } catch (error) {
            Logger.error(error.message)
            if (error.response) {
                Logger.error(error.response.data)
            }
            throw new InternalServerErrorException(error.message)
        }
    }

    public async apiUpload(file: Express.Multer.File, userId: number) {
        let apiUser = await this.apiUserService.findById(userId)
        if (!apiUser) {
            throw new BadRequestException("Can't find api user")
        }
        const { originalname, buffer, mimetype } = file
        const md5 = this.calculateMD5(buffer)
        const res = await ChatEmbedding.findAll({
            where: {
                name: md5
            }
        })
        if (res.length) {
            return {
                name: md5
            }
        }
        const pdf = await PdfParse(buffer)
        const contents = []
        let paragraph = ''
        pdf.text
            .trim()
            .split('\n')
            .forEach((line) => {
                line = line.trim()
                paragraph += line
                if (this.isFullSentence(line)) {
                    contents.push(paragraph)
                    paragraph = ''
                }
            })
        if (paragraph) {
            contents.push(paragraph)
        }

        const embeddings = await this.createEmbeddings(contents)
        Logger.log(
            `create embeddings finished, total token usage: ${embeddings.reduce((acc, cur) => acc + cur.token, 0)}`
        )
        let i = 0
        for (const item of embeddings) {
            try {
                await ChatEmbedding.create({
                    name: md5,
                    text: item.text,
                    num: i++,
                    embedding: formatEmbedding(item.embedding)
                })
            } catch (error) {
                Logger.error(error.message)
            }
        }
        apiUser.code = md5
        this.apiUserService.update(userId, apiUser)
        return {
            name: md5
        }
    }
}
