package cn.iocoder.yudao.module.llm.service.async;

import cn.iocoder.yudao.framework.common.exception.ErrorCode;
import cn.iocoder.yudao.framework.common.util.collection.CollectionUtils;
import cn.iocoder.yudao.module.llm.dal.dataobject.knowledgedocuments.KnowledgeDocumentsDO;
import cn.iocoder.yudao.module.llm.dal.mysql.knowledgedocuments.KnowledgeDocumentsMapper;
import cn.iocoder.yudao.module.llm.enums.KnowledgeStatusEnum;
import cn.iocoder.yudao.module.llm.framework.backend.config.LLMBackendProperties;
import cn.iocoder.yudao.module.llm.service.http.RagHttpService;
import cn.iocoder.yudao.module.llm.service.http.vo.KnowledgeRagEmbedReqVO;
import cn.iocoder.yudao.module.llm.service.http.vo.RegUploadReqVO;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.List;
import java.util.Objects;

import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;

@Service
public class AsyncKnowledgeBase {


    private static final Logger log = LoggerFactory.getLogger(AsyncKnowledgeBase.class);
    @Resource
    private RagHttpService ragHttpService;
    @Resource
    private LLMBackendProperties llmBackendProperties;

    @Resource
    private KnowledgeDocumentsMapper knowledgeDocumentsMapper;


    // 向向量知识库创建文件
    public void createKnowledgeBase (List<KnowledgeDocumentsDO> knowledgeList, List<Long> ids) {
        if (!CollectionUtils.isAnyEmpty(ids)) {
            String mes = ragHttpService.ragDocumentsDel(llmBackendProperties.getRagDocumentsDel(), ids);
            log.info("delete knowledge base info {}", mes);
        }

        // 注释调试
        if (!CollectionUtils.isAnyEmpty(knowledgeList)) {
            knowledgeList.forEach(knowledge -> {
                try {
                    log.info("knowledge base begin create {}", knowledge);

                    // 修改状态为 未上传
                    updateFileState(knowledge, KnowledgeStatusEnum.NOT_UPLOADED);

                    RegUploadReqVO regUploadReqVO = new RegUploadReqVO()
                            .setUrl(llmBackendProperties.getRagEmbed())
                            .setFileId(String.valueOf(knowledge.getId()))
                            .setFileName(knowledge.getDocumentName())
                            .setFileUrl(knowledge.getFileUrl());

                    int lastIndex = knowledge.getDocumentName().lastIndexOf(".");
                    if (lastIndex != -1) {

                        String extension = knowledge.getDocumentName().substring(lastIndex + 1).toLowerCase();
                        if ("txt".equals(extension)) {
                            ragHttpService.embedUploadFile(regUploadReqVO);
                        } else {
                            knowledgeEmbed(knowledge, knowledge.getKnowledgeBaseId());
                        }
                    }
                } catch (Exception e) {
                    log.error("the creation of the knowledge base error {}", e.getMessage());
                    // 修改状态为 上传失败
                    updateFileState(knowledge, KnowledgeStatusEnum.UPLOAD_FAILED);
                    throw exception(new ErrorCode(10047, "文件上传到知识库失败！"));
                }

            });
        }

    }

    /**
     * 修改知识库文档状态
     *
     * @param documents 知识库文档
     * @param status    状态 (0:未上传，1:上传中，2:上传成功，3:上传失败)
     */
    private void updateFileState (KnowledgeDocumentsDO documents, KnowledgeStatusEnum status) {
        documents.setStatus(status.getStatus());
        knowledgeDocumentsMapper.updateById(documents);
    }


    /**
     * 知识库向量嵌入
     *
     * @param knowledge 文件
     * @param id        知识库id
     */
    public void knowledgeEmbed (KnowledgeDocumentsDO knowledge, Long id) {

        try {

            // TODO:本地调试时打开
//            String tmpUrl = "http://xhllm.xinnuojinzhi.com/admin-api/infra/file/29/get/ca3d06d24f80c127ec0300408a035176f5e0bf90ce319fda17018303226e2298.doc";
//            log.info("knowledge url {}", tmpUrl);
//            knowledge.setFileUrl(tmpUrl);

            // 创建知识向量
            KnowledgeRagEmbedReqVO ragEmbedReqVo = new KnowledgeRagEmbedReqVO()
                    .setFileId(String.valueOf(knowledge.getId()))
                    .setFileName(knowledge.getDocumentName())
                    .setFileInputStream(new ByteArrayInputStream(Objects.requireNonNull(getFileByte(knowledge.getFileUrl()))));

            ragHttpService.knowledgeEmbed(ragEmbedReqVo, id);

        } catch (Exception e) {
            log.error("the creation of the knowledge base error {}", e.getMessage(), e);
            throw e;
        }
    }

    /**
     * 获取文件字节数组
     *
     * @param fileUrl 文件地址
     * @return 文件字节数组
     */
    public static byte[] getFileByte (String fileUrl) {
        try (InputStream inputStream = new URL(fileUrl).openStream();
             ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {

            // 缓冲区大小
            byte[] buffer = new byte[1024];
            int bytesRead;

            // 读取文件内容并写入 ByteArrayOutputStream
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, bytesRead);
            }

            // 返回字节数组
            return outputStream.toByteArray();

        } catch (IOException e) {
            log.error("Failed to read remote file: {}", e.getMessage());

            throw exception(new ErrorCode(10001_001, "文件读取错误"));
        }
    }

}
