import { Inject, Provide } from '@midwayjs/decorator';
import { BaseService } from '@cool-midway/core';
import { SearchOptions } from '../interface';
import { VectorStore } from '@langchain/core/vectorstores';
import { Config, IMidwayContext } from '@midwayjs/core';
import {} from '../store/base';
import { KnowStore } from '../store';
import { KnowDataTypeService } from './data/type';
import { DocumentInterface } from '@langchain/core/documents';
import { KnowConfigService } from './config';
import { RerankModel } from '../rerank';
import { KnowRerankBase } from '../rerank/base';

/**
 * 检索服务
 */
@Provide()
export class KnowRetrieverService extends BaseService {
  @Inject()
  ctx: IMidwayContext;

  @Inject()
  knowStore: KnowStore;

  @Config('module.know.prefix')
  prefix: string;

  @Inject()
  knowDataTypeService: KnowDataTypeService;

  @Inject()
  knowConfigService: KnowConfigService;

  /**
   * 调用
   * @param knowId
   * @param text
   */
  async invoke(
    knowId: number,
    text: string,
    options?: SearchOptions
  ): Promise<[DocumentInterface, number][]> {
    const store = await this.getStore(knowId);
    const result = await store.similaritySearchWithScore(
      text,
      options?.size || 10
    );
    return await this.rerank(knowId, text, result);
  }

  /**
   * 重新排序调用rerank
   * @param knowId
   * @param result
   * @returns
   */
  async rerank(
    knowId,
    text: string,
    result: [DocumentInterface, number][],
    options?: SearchOptions
  ) {
    const know = await this.knowDataTypeService.getKnow(knowId);
    if (know.enable == 0) {
      return [];
    }
    // 处理rerank
    if (!know.enableRerank) {
      return result;
    }
    const config = await this.knowConfigService.info(know.rerankConfigId);
    const rerank: KnowRerankBase = new RerankModel[config.type]();
    rerank.config({
      ...config.options.comm,
      ...know.rerankOptions,
    });
    const res = await rerank.rerank(
      result.map(item => item[0]),
      text,
      options?.size || 10
    );
    // 对比 重新返回新的结果
    const newResult = [];
    for (const item of res) {
      const arr = [];
      arr.push(result[item.index]);
      arr.push(item.relevanceScore);
      newResult.push(arr);
    }
    return newResult;
  }

  /**
   * 搜索
   * @param knowIds 知识库ID 多个
   * @param text 文本
   * @param options 配置
   */
  async search(knowIds: number[], text: string, options: SearchOptions) {
    const results: [DocumentInterface, number][][] = await Promise.all(
      knowIds.map(knowId => {
        return this.invoke(knowId, text, options);
      })
    );
    // 合并结果，按照得分排序 number是得分
    const result = results.reduce((prev, curr) => {
      return prev.concat(curr);
    }, []);
    result.sort((a, b) => b[1] - a[1]);
    // 只返回文档结果，并取前N个
    return result.map(item => item[0]).slice(0, options.size || 10);
  }

  /**
   * 获得存储器
   * @param knowIds
   * @returns
   */
  async getStores(knowIds: number[]): Promise<VectorStore[]> {
    return await Promise.all(
      knowIds.map(knowId => {
        return this.getStore(knowId);
      })
    );
  }

  /**
   * 获得存储器
   * @param knowId
   * @returns
   */
  async getStore(knowId: number): Promise<VectorStore> {
    const store = await this.knowStore.get(knowId);
    return await store.getStore(`${this.prefix}${knowId}`);
  }
}
