import {
    BadRequestException,
    ConflictException,
    Injectable,
    InternalServerErrorException,
    Logger,
    OnModuleInit
} from '@nestjs/common'
import { InjectRepository } from '@nestjs/typeorm'
import { KnowledgeBase } from './entities/knowledge-base.entity'
import { In, Repository } from 'typeorm'
import { Tiktoken, get_encoding } from '@dqbd/tiktoken'
import { Configuration, OpenAIApi } from 'azure-openai'
import { DataTypes, Sequelize } from 'sequelize'
import { ConfigService } from '@nestjs/config'
import { KnowledgeEmbedding } from './entities/knowledge-embedding.entity'
import { Documents } from './entities/documents.entity'
import { VECTOR } from '../utils/pgvector'
import * as queue from 'fastq'
import { setTimeout } from 'timers/promises'
import * as PdfParse from '@cyber2024/pdf-parse-fixed'
import { createHash } from 'crypto'
import { PageRequest } from '../common/dto/page-request'
import { Pagination, paginate } from 'nestjs-typeorm-paginate'
import { KnowledgeFile } from './entities/knowledge-file.entity'
import { FileService } from 'src/file/file.service'
import { FileStatus } from './enums/file-status.enum'
import xlsx from 'node-xlsx'
import * as mime from 'mime'
import { OpenAIEmbeddings } from 'langchain/embeddings/openai'
import { ChatOpenAI } from 'langchain/chat_models/openai'
import { HumanMessage } from 'langchain/schema'
import { TypeORMVectorStore } from 'langchain/vectorstores/typeorm'
import { UnstructuredLoader } from 'langchain/document_loaders/fs/unstructured'
import { mkdtempSync, unlinkSync, rmdirSync, writeFileSync, rmSync } from 'fs'
import { join } from 'path'
import { OpenAIParallelEmbeddings } from './embedding'
import { ConversationalRetrievalQAChain } from 'langchain/chains'
import { HNSWLib } from 'langchain/vectorstores/hnswlib'
import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter'
import { BufferMemory } from 'langchain/memory'
import { PromptTemplate } from 'langchain/prompts'
function formatEmbedding(embedding: number[]) {
    return `[${embedding.join(', ')}]`
}

@Injectable()
export class KnowledgeBaseService implements OnModuleInit {
    private readonly tokenizer: Tiktoken
    private embeddingModel: OpenAIParallelEmbeddings
    private chatModel: ChatOpenAI
    private readonly sequelize: Sequelize
    private vectorStore: TypeORMVectorStore
    constructor(
        @InjectRepository(KnowledgeBase)
        private readonly knowledgeBaseRepository: Repository<KnowledgeBase>,
        @InjectRepository(KnowledgeFile)
        private readonly knowledgeFileRepository: Repository<KnowledgeFile>,
        private readonly configService: ConfigService,
        private readonly fileService: FileService
    ) {
        this.tokenizer = get_encoding('cl100k_base')
        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
        })
    }

    async onModuleInit() {
        KnowledgeEmbedding.init(
            {
                id: {
                    primaryKey: true,
                    autoIncrement: true,
                    type: DataTypes.INTEGER
                },
                orgId: {
                    type: DataTypes.INTEGER
                },
                knowledgeId: {
                    type: DataTypes.INTEGER
                },
                fileId: {
                    type: DataTypes.INTEGER
                },
                fileHash: {
                    type: DataTypes.STRING
                },
                text: {
                    type: DataTypes.TEXT({
                        length: 'long'
                    })
                },
                embedding: {
                    type: new VECTOR(1536)
                },
                index: {
                    type: DataTypes.INTEGER
                }
            },
            { sequelize: this.sequelize }
        )
        Documents.init(
            {
                id: {
                    primaryKey: true,
                    type: DataTypes.UUIDV4,
                    defaultValue: 'uuid_generate_v4()'
                },
                pageContent: {
                    type: DataTypes.TEXT
                },
                metadata: {
                    type: DataTypes.JSONB
                },
                embedding: {
                    type: new VECTOR(1536)
                }
            },
            { sequelize: this.sequelize, tableName: 'documents' }
        )
        this.sequelize.sync()

        this.embeddingModel = new OpenAIParallelEmbeddings({
            azureOpenAIApiKey: process.env.AZURE_EMBEDDING_KEY,
            azureOpenAIApiInstanceName: process.env.AZURE_EMBEDDING_INSTANCE,
            azureOpenAIApiDeploymentName: process.env.AZURE_EMBEDDING_DEPLOYMENT,
            azureOpenAIApiVersion: process.env.AZURE_EMBEDDING_VERSION,
            azureOpenAIApiEmbeddingsDeploymentName: process.env.AZURE_EMBEDDING_DEPLOYMENT,
            maxRetries: 5,
            timeout: 15000
        })
        this.chatModel = new ChatOpenAI({
            azureOpenAIApiKey: process.env.AZURE_OPENAI_KEY,
            azureOpenAIApiInstanceName: process.env.AZURE_OPENAI_INSTANCE,
            azureOpenAIApiDeploymentName: process.env.AZURE_OPENAI_DEPLOYMENT,
            azureOpenAIApiVersion: process.env.AZURE_OPENAI_VERSION,
            maxRetries: 5,
            timeout: 15000,
            temperature: 0.1
        })
        this.vectorStore = await TypeORMVectorStore.fromDataSource(this.embeddingModel, {
            postgresConnectionOptions: {
                type: 'postgres',
                host: process.env.PG_HOST,
                port: Number(process.env.PG_PORT),
                username: process.env.PG_USERNAME,
                password: process.env.PG_PASSWORD,
                database: process.env.PG_DATABASE
            },
            verbose: true
        })
        await this.vectorStore.ensureTableInDatabase()
    }

    async findAllKnowledgeBase(req: PageRequest<KnowledgeBase>): Promise<Pagination<KnowledgeBase>> {
        return await paginate<KnowledgeBase>(this.knowledgeBaseRepository, req.page, req.search)
    }

    async createKnowledgeBase(knowledgeBase: Partial<KnowledgeBase>): Promise<KnowledgeBase> {
        return await this.knowledgeBaseRepository.save(knowledgeBase)
    }

    async updateKnowledgeBase(knowledgeBase: Partial<KnowledgeBase>): Promise<KnowledgeBase> {
        return await this.knowledgeBaseRepository.save(knowledgeBase)
    }

    async deleteKnowledgeBase(knowledgeBaseId: number): Promise<void> {
        await this.knowledgeBaseRepository.delete(knowledgeBaseId)
        await this.knowledgeFileRepository.delete({ knowledgeId: knowledgeBaseId })
        await KnowledgeEmbedding.destroy({
            where: {
                knowledgeId: knowledgeBaseId
            }
        })
        await Documents.destroy({
            where: {
                metadata: {
                    knowledgeId: knowledgeBaseId
                }
            }
        })
    }

    async getKnowledgeBaseById(knowledgeBaseId: number): Promise<KnowledgeBase> {
        return await this.knowledgeBaseRepository.findOneOrFail({ where: { id: knowledgeBaseId } })
    }

    async fileAllKnowledgeFile(req: PageRequest<KnowledgeFile>): Promise<Pagination<KnowledgeFile>> {
        return await paginate<KnowledgeFile>(this.knowledgeFileRepository, req.page, req.search)
    }

    async updateKnowledgeFile(knowledgeFile: Partial<KnowledgeFile>): Promise<KnowledgeFile> {
        return await this.knowledgeFileRepository.save(knowledgeFile)
    }

    async deleteKnowledgeFile(knowledgeFileId: number): Promise<void> {
        await this.knowledgeFileRepository.delete(knowledgeFileId)
        await KnowledgeEmbedding.destroy({
            where: {
                fileId: knowledgeFileId
            }
        })
        await Documents.destroy({
            where: {
                metadata: {
                    fileId: knowledgeFileId
                }
            }
        })
    }

    public async uploadKnowledgeFile(file: Express.Multer.File, knowledgeId: number) {
        const knowledgeBase = await this.getKnowledgeBaseById(knowledgeId)
        const { originalname, buffer, mimetype, size } = file
        const fileName = Buffer.from(originalname, 'latin1').toString('utf8')
        let fileHash = this.calculateMD5(buffer)
        let knowledgeFile = await this.knowledgeFileRepository.findOneBy({
            fileHash
        })
        if (knowledgeFile) {
            throw new ConflictException(`File ${fileName} already exists`)
        }
        const { url: fileUrl } = await this.fileService.uploadBuffer(
            buffer,
            mimetype.split('/')[1],
            mime.getExtension(mimetype)
        )
        knowledgeFile = new KnowledgeFile()
        knowledgeFile.orgId = knowledgeBase.orgId
        knowledgeFile.knowledgeId = knowledgeId
        knowledgeFile.fileHash = fileHash
        knowledgeFile.fileType = mimetype
        knowledgeFile.fileName = fileName
        knowledgeFile.size = size
        knowledgeFile.fileUrl = fileUrl
        await this.knowledgeFileRepository.save(knowledgeFile)
        switch (mimetype) {
            case 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet':
                this.processExcelKnowledgeFile(knowledgeFile, buffer)
                break
            case 'application/pdf':
                this.processPdfKnowledgeFile(knowledgeFile, buffer)
                // this.processFile(knowledgeFile, buffer)
                break
        }
        return knowledgeFile
    }

    public async processPdfKnowledgeFile(knowledgeFile: KnowledgeFile, buffer: Buffer) {
        knowledgeFile.status = FileStatus.PROCESSING
        try {
            await this.knowledgeFileRepository.save(knowledgeFile)
            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.map((i) => {
                    return { text: i }
                })
            )
            Logger.log(
                `create embeddings finished, total token usage: ${embeddings.reduce((acc, cur) => acc + cur.token, 0)}`
            )
            await KnowledgeEmbedding.destroy({
                where: {
                    fileHash: knowledgeFile.fileHash
                }
            })
            let i = 0
            for (const item of embeddings) {
                try {
                    await KnowledgeEmbedding.create({
                        orgId: knowledgeFile.orgId,
                        knowledgeId: knowledgeFile.knowledgeId,
                        fileId: knowledgeFile.id,
                        fileHash: knowledgeFile.fileHash,
                        text: item.text,
                        embedding: formatEmbedding(item.embedding),
                        index: i++
                    })
                } catch (error) {
                    Logger.error(error.message)
                }
            }
            knowledgeFile.status = FileStatus.DONE
            await this.knowledgeFileRepository.save(knowledgeFile)
        } catch (e) {
            knowledgeFile.status = FileStatus.FAILED
            knowledgeFile.error = e.message
            await this.knowledgeFileRepository.save(knowledgeFile)
        }
    }

    public async processFile(knowledgeFile: KnowledgeFile, buffer: Buffer) {
        knowledgeFile.status = FileStatus.PROCESSING
        try {
            await this.knowledgeFileRepository.save(knowledgeFile)
            const tmpDir = mkdtempSync('file')
            const tmpFile = join(tmpDir, knowledgeFile.fileName)
            writeFileSync(tmpFile, buffer)
            const loader = new UnstructuredLoader(tmpFile, {
                apiUrl: process.env.UNSTRUCTURED_API,
                strategy: 'fast'
            })
            const docs = await loader.load()
            const splitter = new RecursiveCharacterTextSplitter({
                chunkSize: 300,
                chunkOverlap: 100
            })
            let output: any[] = await splitter.createDocuments(
                [docs.map((i) => i.pageContent).join('\n')],
                [
                    {
                        orgId: knowledgeFile.orgId,
                        knowledgeId: knowledgeFile.knowledgeId,
                        fileId: knowledgeFile.id,
                        fileHash: knowledgeFile.fileHash,
                        fileName: knowledgeFile.fileName
                    }
                ],
                {
                    chunkHeader: `DOCUMENT NAME: [${knowledgeFile.fileName}](${knowledgeFile.fileUrl})\n\n---\n\n`,
                    appendChunkOverlapHeader: true
                }
            )
            Logger.log(`load ${output.length} docs`)
            rmSync(tmpDir, { recursive: true })
            await this.vectorStore.addDocuments(output)
            knowledgeFile.status = FileStatus.DONE
            await this.knowledgeFileRepository.save(knowledgeFile)
        } catch (e) {
            Logger.error(e)
            Logger.error(e.stack)
            knowledgeFile.status = FileStatus.FAILED
            knowledgeFile.error = e.message
            await this.knowledgeFileRepository.save(knowledgeFile)
        }
    }

    public async processExcelKnowledgeFile(knowledgeFile: KnowledgeFile, buffer: Buffer) {
        knowledgeFile.status = FileStatus.PROCESSING
        try {
            await this.knowledgeFileRepository.save(knowledgeFile)
            const sheets = xlsx.parse(buffer)
            for (let sheet of sheets) {
                const data = sheet.data.filter((i) => i[0].length && i[1].length)
                const contents = data.map((i) => {
                    return {
                        text: i[0],
                        detail: i[1]
                    }
                })
                const embeddings = await this.createEmbeddings(contents)
                Logger.log(
                    `create embeddings finished, total token usage: ${embeddings.reduce(
                        (acc, cur) => acc + cur.token,
                        0
                    )}`
                )
                await KnowledgeEmbedding.destroy({
                    where: {
                        fileHash: knowledgeFile.fileHash
                    }
                })
                let i = 0
                for (const item of embeddings) {
                    try {
                        await KnowledgeEmbedding.create({
                            orgId: knowledgeFile.orgId,
                            knowledgeId: knowledgeFile.knowledgeId,
                            fileId: knowledgeFile.id,
                            fileHash: knowledgeFile.fileHash,
                            text: item.text + '\n' + item.detail,
                            embedding: formatEmbedding(item.embedding),
                            index: i++
                        })
                    } catch (error) {
                        Logger.error(error.message)
                    }
                }
                knowledgeFile.status = FileStatus.DONE
                await this.knowledgeFileRepository.save(knowledgeFile)
            }
        } catch (e) {
            knowledgeFile.status = FileStatus.FAILED
            knowledgeFile.error = e.message
            await this.knowledgeFileRepository.save(knowledgeFile)
        }
    }

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

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

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

    async getEmbedding(content: string) {
        try {
            const response = await this.embeddingModel.embedQuery(content)
            return {
                text: content,
                embedding: response,
                token: this.tokenizer.encode(content).length
            }
        } catch (error) {
            Logger.error(error.stack, 'fetchEmbedding')
            throw new InternalServerErrorException(error.message)
        }
    }

    async getKeywords(text: string) {
        return (
            await this.chatModel.call([
                new HumanMessage(
                    `You need to extract keywords from the statement or question and return a series of keywords separated by commas.\ncontent: ${text}\nkeywords: `
                )
            ])
        ).content
    }

    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 askKnowledge(question: string, orgId: number, knowledgeId?: number, fileId?: number) {
        const keywords = await this.getKeywords(question)
        const { embedding: keywordEmbedding } = await this.getEmbedding(keywords)
        const where = { orgId, knowledgeId, fileId }
        Object.keys(where).forEach((key) => (where[key] === undefined ? delete where[key] : {}))
        const relatedEmbeddings = await KnowledgeEmbedding.findAll({
            where,
            order: this.sequelize.literal(`embedding <-> '${formatEmbedding(keywordEmbedding)}'`),
            limit: 100
        })
        const files = await this.knowledgeFileRepository.findBy({
            fileHash: In([...new Set(relatedEmbeddings.map((item) => item.fileHash))])
        })
        const context = await this.cutContext(
            relatedEmbeddings.map((item) => {
                const file = files.find((i) => i.fileHash === item.fileHash)
                if (file) {
                    return `***\n[${file.fileName}](${file.fileUrl}):\n\n${item.text}`
                }
                return item.text
            })
        )
        return context
    }

    async askKnowledge1(options: { question: string; orgId: number; knowledgeId?: number; fileId?: number }) {
        /* Create the chain */
        const chain = ConversationalRetrievalQAChain.fromLLM(
            this.chatModel,
            this.vectorStore.asRetriever(100, {
                orgId: options.orgId
            }),
            {
                memory: new BufferMemory({
                    memoryKey: 'chat_history' // Must be set to "chat_history"
                }),
                qaChainOptions: {
                    prompt: PromptTemplate.fromTemplate(`(以下内容是相关的本地知识库信息，按照相关程度从高到低排序，第一行是文件名和链接)
\`\`\`
{context}
\`\`\`
请仔细思考，优先使用以上提供的内容进行回答，并告诉我相关的文件和链接。如果无法从内容中找到答案，你需要根据给你自己掌握的知识回答问题。
请确保使用中文进行回答。
请在末尾附上相关的文件链接。`),
                    questionPrompt: PromptTemplate.fromTemplate(`{question}`)
                },
                questionGeneratorChainOptions: {
                    template: `Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question.

Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:`
                }
            }
        )
        /* Ask it a question */
        const res = await chain.call({ question: options.question })
        console.log(res)
    }
}
