import { ChatOpenAI } from 'langchain/chat_models/openai'
import { CallbackManager } from 'langchain/callbacks'
import { ChatPromptTemplate, MessagesPlaceholder } from 'langchain/prompts'
import dedent from 'dedent'
import { ConversationChain } from 'langchain/chains'
import { BufferMemory, BufferWindowMemory } from 'langchain/memory'
import { Logger } from '@nestjs/common'
import { writeFile, mkdirSync, appendFileSync } from 'fs'
import path = require('path')
import { LLMTools } from './types'
require('dotenv').config()
export function createLLM(options?: {
    model?: 'gpt-3.5-turbo' | 'gpt-3.5-turbo-16k' | 'gpt-4' | 'gpt-4-32k'
    provider?: 'azure' | 'openai'
}): LLMTools {
    options = options || {
        model: 'gpt-3.5-turbo',
        provider: 'azure'
    }
    options.model = options.model || 'gpt-3.5-turbo'
    options.provider = options.provider || 'azure'
    const usage = { completionTokens: 0, promptTokens: 0, totalTokens: 0 }
    const commonParams = {
        maxRetries: 4,
        timeout: 1000 * 60 * 2,
        callbackManager: CallbackManager.fromHandlers({
            async handleLLMStart(llm, prompts) {
                Logger.log(`[LLM Start]LLM: ${JSON.stringify(llm)}`)
                Logger.log(`['LLM Start]Prompts: ${prompts.join('\n')}`)
            },
            async handleLLMEnd(output) {
                Logger.log(
                    `[LLM End]${output.generations
                        .reduce((acc, cur) => acc.concat(cur), [])
                        .map((i) => i.text)
                        .join('\n')}`
                )
                Logger.log(`[LLM End]${JSON.stringify(output.llmOutput)}`)
                usage.completionTokens += output.llmOutput.tokenUsage.completionTokens
                usage.promptTokens += output.llmOutput.tokenUsage.promptTokens
                usage.totalTokens += output.llmOutput.tokenUsage.totalTokens
            },
            async handleLLMError(error) {
                Logger.error(error)
            }
        }),
        onFailedAttempt(error) {
            Logger.error(error.message, error.stack, 'OpenAI API Error')
        }
    }
    let llm
    if (options.provider === 'openai') {
        llm = new ChatOpenAI({
            ...commonParams,
            openAIApiKey: process.env.OPENAI_API_KEY,
            modelName: options.model,

            // azureOpenAIApiKey: process.env.AZURE_OPENAI_KEY,
            // azureOpenAIApiVersion: process.env.AZURE_OPENAI_VERSION,
            // azureOpenAIApiInstanceName: process.env.AZURE_OPENAI_INSTANCE,
            // azureOpenAIApiDeploymentName: 'gpt-35-turbo-16k',
            // azureOpenAIBasePath:`https://gateway.ai.cloudflare.com/v1/cc11bed478e65817d30878b0a796422d/openai-gateway/azure-openai/${process.env.AZURE_OPENAI_INSTANCE}`,

            configuration: {
                baseURL: process.env.OPENAI_BASE_URL
            }
        })
    } else {
        let azureDeployment = 'gpt-35'
        switch (options.model) {
            case 'gpt-3.5-turbo':
                azureDeployment = 'gpt-35-turbo'
                break
            case 'gpt-3.5-turbo-16k':
                azureDeployment = 'gpt-35-turbo-16k'
                break
            case 'gpt-4':
                azureDeployment = 'gpt-4'
                break
            case 'gpt-4-32k':
                azureDeployment = 'gpt-4-32k'
                break
        }
        llm = new ChatOpenAI({
            ...commonParams,
            azureOpenAIApiKey: process.env.AZURE_OPENAI_KEY,
            azureOpenAIApiVersion: process.env.AZURE_OPENAI_VERSION,
            azureOpenAIApiInstanceName: process.env.AZURE_OPENAI_INSTANCE,
            azureOpenAIApiDeploymentName: azureDeployment
            // azureOpenAIBasePath: `https://gateway.ai.cloudflare.com/v1/cc11bed478e65817d30878b0a796422d/openai-gateway/azure-openai/${process.env.AZURE_OPENAI_INSTANCE}`
        })
    }

    function conversation(system, k = 4) {
        const chatPrompt = ChatPromptTemplate.fromMessages([
            ['system', system],
            new MessagesPlaceholder('history'),
            ['human', '{input}']
        ])
        const memory = new BufferWindowMemory({
            k,
            memoryKey: 'history',
            returnMessages: true
        })
        const chain = new ConversationChain({
            memory: memory,
            prompt: chatPrompt,
            llm: llm
        })
        return { memory, chain }
    }
    return { llm, usage, conversation }
}
