import { CoolCommException } from '@cool-midway/core';
import { BaseChatModel } from '@langchain/core/language_models/chat_models';
import { AIMessage, HumanMessage } from '@langchain/core/messages';
import { ChatPromptTemplate, PromptTemplate } from '@langchain/core/prompts';
import { IterableReadableStream } from '@langchain/core/utils/stream';
import { Inject, Provide, Scope, ScopeEnum } from '@midwayjs/core';
import { FlowContext } from '../../runner/context';
import { FlowNode } from '../../runner/node';
import { FlowStream } from '../../runner/stream';
import { FlowConfigService } from '../../service/config';
import { FlowDataService } from '../../service/data';
import { NodeLLMModel } from './model';

// 消息体
interface Message {
  // 角色
  role: 'system' | 'user' | 'assistant';
  // 内容
  content: any;
}

/**
 * LLM大模型节点
 */
@Provide()
@Scope(ScopeEnum.Prototype)
export class NodeLLM extends FlowNode {
  @Inject()
  nodeLLMModel: NodeLLMModel;

  @Inject()
  flowConfigService: FlowConfigService;

  @Inject()
  flowDataService: FlowDataService;

  /**
   * 执行
   * @param context
   */
  async run(context: FlowContext) {
    let { model, messages, history } = this.config.options;
    const config = await this.flowConfigService.getOptions(model.configId);
    const llm = await this.getModel(model.supplier, {
      ...model.params,
      ...config.comm,
    });
    const params = this.inputParams;
    const prompt = await this.getPrompt(messages, context, params, history);
    const chain = prompt.pipe(llm);
    let stream: FlowStream;
    let res;
    if (context.isStream()) {
      stream = new FlowStream();
      const _stream: IterableReadableStream<any> = await chain.stream(params);
      const content = [];
      const save = async () => {
        // 流式输出
        for await (const chunk of _stream) {
          content.push(chunk.content);
          stream.push(chunk.content);
        }
        stream.push(null);
        // 保存历史
        this.saveHistory(messages, context, history, content.join(''));
      };
      save();
    } else {
      res = await chain.invoke(params);
      // 保存历史
      await this.saveHistory(messages, context, history, res?.content);
    }
    context.set(`${this.getPrefix()}.text`, res?.content, 'output');
    context.set(`${this.getPrefix()}.stream`, stream, 'output');
    context.updateCount(
      'tokenUsage',
      res?.response_metadata.tokenUsage?.totalTokens || 0
    );
    return {
      success: true,
      stream,
      result: {
        text: res?.content,
      },
    };
  }

  /**
   * 获得提示模板
   * @param messages
   * @returns
   */
  async getPrompt(
    messages: Message[],
    context: FlowContext,
    params: any,
    history = 0
  ) {
    const type = {
      system: 'system',
      user: 'human',
      placeholder: 'placeholder',
      assistant: 'ai',
    };
    let prompt;
    if (history) {
      const objectId = context.getRequestParams().objectId;
      if (!objectId) {
        throw new CoolCommException(
          '需要保存历史信息，请求参数必须包含objectId'
        );
      }
      // 取出历史
      const historyMessages =
        (await this.flowDataService.get(this.flowId, objectId)) || [];
      prompt = ChatPromptTemplate.fromMessages([
        [type[messages[0].role], messages[0].content],
        ['placeholder', '{chat_history}'],
        // @ts-ignore
        ...messages
          .filter(item => item.role != 'system')
          .map(item => {
            return [type[item.role], item.content];
          }),
      ]);
      params['chat_history'] = historyMessages.map(item => {
        return item.role == 'user'
          ? new HumanMessage({ content: item.content })
          : new AIMessage({ content: item.content });
      });
    } else {
      // 转换格式
      prompt = ChatPromptTemplate.fromMessages(
        messages.map(item => {
          return [type[item.role], item.content];
        })
      );
    }
    return prompt;
  }

  /**
   * 保存历史
   * @param messages
   * @param context
   * @param objectId
   */
  async saveHistory(
    messages: Message[],
    context: FlowContext,
    history = 0,
    newContent: string
  ) {
    const params = this.inputParams;
    const objectId = context.getRequestParams().objectId;
    if (!objectId || !history) return;

    const historyMessages =
      (await this.flowDataService.get(this.flowId, objectId)) || [];
    const lastMessage = messages.filter(item => item.role === 'user').pop();
    const userContent = await PromptTemplate.fromTemplate(
      lastMessage.content
    ).invoke(params);
    const userMessage = {
      role: 'user',
      content: userContent.value,
    };
    const newMessages = historyMessages
      .concat(userMessage)
      .concat({ role: 'assistant', content: newContent });

    const totalLength = newMessages.length;
    if (totalLength > history) {
      newMessages.splice(0, totalLength - history);
    }

    await this.flowDataService.set(this.flowId, objectId, newMessages);
  }

  /**
   * 获得模型
   * @param name 名称
   * @param options 配置
   * @returns
   */
  async getModel(name: string, options: any): Promise<BaseChatModel> {
    const LLM = await this.nodeLLMModel.getModel(name);
    // @ts-ignore
    return new LLM(options);
  }
}
