package com.ard.agent.service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.FileSystemResource;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.LocalDateTime;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
public class KnowledgeBaseService {

    private final VectorStore vectorStore;
    private final TokenTextSplitter textSplitter;

    public KnowledgeBaseService(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
        this.textSplitter = new TokenTextSplitter(800, 150, 5, 10000, true);
    }

    public Map<String, Object> uploadDocument(MultipartFile file) {
        Map<String, Object> result = new HashMap<>();
        String fileName = file.getOriginalFilename();
        result.put("fileName", fileName);

        Path tempFile = null;
        try {
            log.info("开始处理文档: {}", fileName);

            String contentType = file.getContentType();
            if (!isSupportedContentType(contentType)) {
                result.put("success", false);
                result.put("message", "不支持的文件类型: " + contentType);
                return result;
            }

            // 创建临时文件
            tempFile = Files.createTempFile("upload_", "_" + fileName);
            file.transferTo(tempFile.toFile());
            log.info("临时文件创建成功: {}", tempFile.toString());

            // 关键修改：使用 FileSystemResource 而不是字符串路径
            FileSystemResource resource = new FileSystemResource(tempFile.toFile());
            TikaDocumentReader reader = new TikaDocumentReader(resource);
            List<Document> documents = reader.get();

            if (documents == null || documents.isEmpty()) {
                result.put("success", false);
                result.put("message", "文档解析结果为空");
                return result;
            }

            List<Document> splitDocuments = textSplitter.apply(documents);
            log.info("文本分割完成，共 {} 个分片", splitDocuments.size());

            for (Document doc : splitDocuments) {
                doc.getMetadata().put("fileName", fileName);
                doc.getMetadata().put("contentType", contentType);
                doc.getMetadata().put("uploadTime", LocalDateTime.now().toString());
                doc.getMetadata().put("fileSize", file.getSize());
            }

            vectorStore.add(splitDocuments);

            log.info("文档入库完成: {}，共 {} 个向量分片", fileName, splitDocuments.size());
            result.put("success", true);
            result.put("message", "入库成功");
            result.put("chunksCount", splitDocuments.size());

        } catch (Exception e) {
            log.error("文档入库失败: {}", fileName, e);
            result.put("success", false);
            result.put("message", "入库失败：" + e.getMessage());
        } finally {
            if (tempFile != null) {
                try {
                    Files.deleteIfExists(tempFile);
                } catch (IOException e) {
                    log.warn("删除临时文件失败: {}", tempFile, e);
                }
            }
        }
        return result;
    }

    private boolean isSupportedContentType(String contentType) {
        if (contentType == null) return false;
        String[] supportedTypes = {
                "text/plain",
                "application/pdf",
                "application/msword",
                "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                "application/vnd.ms-excel",
                "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                "text/markdown"
        };
        return Arrays.asList(supportedTypes).contains(contentType);
    }


}