package com.lanyuanxiaoyao.service.ai.knowledge.service.node;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.ai.knowledge.entity.EmbeddingContext;
import com.yomahub.liteflow.annotation.LiteflowComponent;
import com.yomahub.liteflow.annotation.LiteflowMethod;
import com.yomahub.liteflow.core.NodeComponent;
import com.yomahub.liteflow.enums.LiteFlowMethodEnum;
import com.yomahub.liteflow.enums.NodeTypeEnum;
import io.qdrant.client.QdrantClient;
import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.core.io.PathResource;

/**
 * @author lanyuanxiaoyao
 * @version 20250523
 */
@LiteflowComponent
public class EmbeddingNodes {
  private static final Logger logger = LoggerFactory.getLogger(EmbeddingNodes.class);

  private final ChatClient chatClient;
  private final QdrantClient qdrantClient;
  private final EmbeddingModel embeddingModel;

  public EmbeddingNodes(ChatClient.Builder builder, VectorStore vectorStore, EmbeddingModel embeddingModel) {
    this.chatClient = builder.build();
    this.qdrantClient = (QdrantClient) vectorStore.getNativeClient().orElseThrow();
    this.embeddingModel = embeddingModel;
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS_BOOLEAN, nodeId = "embedding_check_if_file_needed", nodeName = "判断是否需要读取文件", nodeType = NodeTypeEnum.BOOLEAN)
  public boolean checkIfFileReadNeeded(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    Assert.notNull(context, "EmbeddingContext is null");
    if (StrUtil.isNotBlank(context.getFile())) {
      Assert.isTrue(FileUtil.exist(context.getFile()), "File [{}] not exist", context.getFile());
      return true;
    }
    Assert.notBlank(context.getContent(), "Contents is empty");
    return false;
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS_SWITCH, nodeId = "file_reader_switch", nodeName = "判断文件格式", nodeType = NodeTypeEnum.SWITCH)
  public String fileReaderSwitch(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    String extName = FileUtil.extName(context.getFileFormat());
    return switch (extName.toLowerCase()) {
      case "txt", "md", "markdown" -> "txt_file_reader";
      case "pdf" -> "pdf_file_reader";
      case "doc", "docx", "xls", "xlsx", "ppt", "pptx", "html", "xml", "wps", "et", "dpt" -> "any_file_reader";
      default -> throw new IllegalStateException("Unsupported ext: " + extName);
    };
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS, nodeId = "txt_file_reader", nodeName = "读取文本文件", nodeType = NodeTypeEnum.COMMON)
  public void txtFileReader(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    context.setContent(FileUtil.readString(context.getFile(), Charset.defaultCharset()));
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS, nodeId = "pdf_file_reader", nodeName = "读取pdf文件", nodeType = NodeTypeEnum.COMMON)
  public void pdfFileReader(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    PagePdfDocumentReader reader = new PagePdfDocumentReader(
        new PathResource(context.getFile()),
        PdfDocumentReaderConfig.builder()
            .withPageTopMargin(0)
            .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                .withNumberOfTopTextLinesToDelete(0)
                .build())
            .build());
    context.setContent(readBySpringAiReader(reader));
  }

  /**
   * <a href="https://tika.apache.org/3.1.0/formats.html#Database_formats">Tika支持的格式</a>
   */
  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS, nodeId = "any_file_reader", nodeName = "使用Tika尝试读取文件", nodeType = NodeTypeEnum.COMMON)
  public void anyFileReader(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    context.setContent(readBySpringAiReader(new TikaDocumentReader(new PathResource(context.getFile()))));
  }

  private String readBySpringAiReader(DocumentReader reader) {
    return reader.get()
        .stream()
        .map(Document::getText)
        .collect(Collectors.joining("\n"))
        .trim();
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS_SWITCH, nodeId = "split_switch", nodeName = "判断使用什么分段方法", nodeType = NodeTypeEnum.SWITCH)
  public String splitSwitch(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    return switch (context.getConfig().getSplitStrategy()) {
      case NORMAL -> "normal_split";
      case LLM -> "llm_split";
      case QA -> "qa_split";
    };
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS, nodeId = "normal_split", nodeName = "使用普通分段", nodeType = NodeTypeEnum.COMMON)
  public void normalSplit(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    TokenTextSplitter splitter = new TokenTextSplitter(200, 100, 5, 200, true);
    Document document = Document.builder()
        .text(context.getContent())
        .build();
    context.setDocuments(splitter.split(document));
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS, nodeId = "llm_split", nodeName = "使用大模型分段", nodeType = NodeTypeEnum.COMMON)
  public void llmSplit(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    context.getDocuments().addAll(llmSplit(
        """
            请你将用户输入的文本进行语义切分，生成用于知识库检索的文本段。
            每个文本段要尽可能多地覆盖用户输入文本的各方面知识和细节，包括但不限于主题、概念、关键信息等。对于关键的数字、理论、细节等，要严格遵循原文，不能进行任何虚构和捏造不存在的知识，确保输出内容准确、真实且全面。
            输出格式为纯文本段，分段之间使用“---”作为分割，方便后续使用代码进行切分。
            输出文本避免添加markdown格式，保持文本格式紧凑。
            切分过程中，要注重保持文本的完整性和逻辑性，确保每个文本段都能独立地表达出清晰、准确的信息，以便更好地进行知识库检索。
            """,
        context.getContent(),
        context.getMetadata()
    ));
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS, nodeId = "qa_split", nodeName = "使用Q/A格式分段", nodeType = NodeTypeEnum.COMMON)
  public void qaSplit(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    context.getDocuments().addAll(llmSplit(
        """
            对用户输入的文本，生成多组高质量的问答对。请遵循以下指南：
            1. 问题部分：
              为同一个主题创建尽可能多的不同表述的问题，确保问题的多样性。
              每个问题应考虑用户可能的多种问法，例如：
              直接询问（如“什么是...？”）
              请求确认（如“是否可以说...？”）
              寻求解释（如“请解释一下...的含义。”）
              假设性问题（如“如果...会怎样？”）
              例子请求（如“能否举个例子说明...？”）
              问题应涵盖文本中的关键信息、主要概念和细节，确保不遗漏重要内容。
            2. 答案部分：
              提供一个全面、信息丰富的答案，涵盖问题的所有可能角度，确保逻辑连贯。
              答案应直接基于给定文本，确保准确性和一致性。
              包含相关的细节，如日期、名称、职位等具体信息，必要时提供背景信息以增强理解。
            3. 格式：
              使用"问："标记问题集合的开始，所有问题应在一个段落内，问题之间用空格分隔。
              使用"答："标记答案的开始，答案应清晰分段，便于阅读。
              问答对之间用“---”分隔，以提高可读性。
            4. 内容要求：
              确保问答对紧密围绕文本主题，避免偏离主题。
              避免添加文本中未提及的信息，确保信息的真实性。
              一个问题搭配一个答案，避免一组问答对中同时涉及多个问题。
              如果文本信息不足以回答某个方面，可以在答案中说明 "根据给定信息无法确定"，并尽量提供相关的上下文。
            格式样例：
            问：苹果通常是什么颜色的？
            答：红色。
            ---
            问：苹果长在树上还是地上？
            答：苹果长在树上。
            """,
        context.getContent(),
        context.getMetadata()
    ));
  }

  private List<Document> llmSplit(String prompt, String content, Map<String, Object> metadata) {
    String response = chatClient.prompt()
        .system(prompt)
        .user(content)
        .call()
        .content();
    Assert.notBlank(response, "LLM response is empty");
    // noinspection DataFlowIssue
    return Arrays.stream(StrUtil.trim(response).split("---"))
        .map(text -> text.replaceAll("(?!^.+) +$", ""))
        .map(StrUtil::trim)
        .map(text -> Document.builder()
            .text(text)
            .metadata(Optional.ofNullable(metadata).orElse(new HashMap<>()))
            .build())
        .toList();
  }

  @LiteflowMethod(value = LiteFlowMethodEnum.PROCESS, nodeId = "import_vector_source", nodeName = "导入向量库", nodeType = NodeTypeEnum.COMMON)
  public void importVectorSource(NodeComponent node) {
    EmbeddingContext context = node.getContextBean(EmbeddingContext.class);
    if (ObjectUtil.isNotEmpty(context.getDocuments())) {
      VectorStore vs = QdrantVectorStore.builder(qdrantClient, embeddingModel)
          .collectionName(String.valueOf(context.getVectorSourceId()))
          .build();
      for (Document document : context.getDocuments()) {
        Map<String, Object> metadata = document.getMetadata();
        metadata.put("filename", context.getFileFormat());
        metadata.put("filepath", context.getFile());
        metadata.put("group_id", String.valueOf(context.getGroupId()));
        metadata.put("vector_source_id", String.valueOf(context.getVectorSourceId()));
      }
      vs.add(context.getDocuments());
    }
  }
}
