agent-management/backend/src/main/java/com/agent/management/service/impl/DocumentServiceImpl.java

package com.agent.management.service.impl;

import com.agent.management.common.exception.BusinessException;
import com.agent.management.config.KbProperties;
import com.agent.management.model.entity.KbChunk;
import com.agent.management.model.entity.KbDocument;
import com.agent.management.repository.KbChunkRepository;
import com.agent.management.repository.KbDocumentRepository;
import com.agent.management.service.DocumentPipeline;
import com.agent.management.service.DocumentService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageRequest;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;

@Slf4j
@Service
@RequiredArgsConstructor
@Transactional
public class DocumentServiceImpl implements DocumentService {

    private final KbDocumentRepository documentRepository;
    private final KbChunkRepository chunkRepository;
    private final EmbeddingBridgeClient bridgeClient;
    private final DocumentPipeline pipeline;
    private final KbProperties kbProps;

    @Override
    public KbDocument upload(MultipartFile file, Long categoryId) {
        if (file == null || file.isEmpty()) {
            throw new BusinessException("上传文件为空");
        }
        if (file.getSize() > kbProps.getMaxFileSize()) {
            throw new BusinessException("文件大小超过限制（" + (kbProps.getMaxFileSize() / 1024 / 1024) + "MB）");
        }
        String mime = file.getContentType();
        if (mime == null || !kbProps.getAllowedMimeTypes().contains(mime)) {
            throw new BusinessException("不支持的文件类型：" + mime);
        }

        // 落盘
        Path uploadDir = Paths.get(kbProps.getUploadDir());
        try {
            Files.createDirectories(uploadDir);
        } catch (IOException e) {
            throw new BusinessException("创建上传目录失败：" + e.getMessage());
        }
        String storedName = UUID.randomUUID() + "_" + sanitizeFileName(file.getOriginalFilename());
        Path target = uploadDir.resolve(storedName);
        try {
            Files.copy(file.getInputStream(), target, StandardCopyOption.REPLACE_EXISTING);
        } catch (IOException e) {
            throw new BusinessException("文件保存失败：" + e.getMessage());
        }

        // 创建文档元数据
        KbDocument doc = new KbDocument();
        doc.setName(file.getOriginalFilename());
        doc.setCategoryId(categoryId);
        doc.setSourcePath(storedName);
        doc.setFileSize(file.getSize());
        doc.setMimeType(mime);
        doc.setStatus("PENDING");
        doc.setChunkCount(0);
        doc.setVectorCount(0);
        doc = documentRepository.save(doc);
        log.info("文档已上传，id={}，name={}，触发异步处理", doc.getId(), doc.getName());

        // 触发异步流水线（独立 bean，避免 @Async 自调用代理失效）
        pipeline.process(doc.getId());

        return doc;
    }

    @Override
    public PageResult list(Long categoryId, String keyword, int page, int size) {
        PageRequest pageable = PageRequest.of(Math.max(0, page - 1), Math.max(1, size));
        Page<KbDocument> p;
        boolean hasKw = keyword != null && !keyword.trim().isEmpty();
        if (categoryId != null && hasKw) {
            p = documentRepository.findByCategoryIdAndNameContainingIgnoreCase(categoryId, keyword.trim(), pageable);
        } else if (categoryId != null) {
            p = documentRepository.findByCategoryId(categoryId, pageable);
        } else if (hasKw) {
            p = documentRepository.findByNameContainingIgnoreCase(keyword.trim(), pageable);
        } else {
            p = documentRepository.findAll(pageable);
        }
        return new PageResult(p.getContent(), p.getTotalElements(), page, size);
    }

    @Override
    public KbDocument get(Long id) {
        return documentRepository.findById(id)
                .orElseThrow(() -> new BusinessException(404, "文档不存在"));
    }

    @Override
    public KbDocument revectorize(Long id) {
        KbDocument doc = get(id);
        // 清理旧向量
        try {
            bridgeClient.deleteByDocument(id);
        } catch (Exception e) {
            log.warn("清理旧向量失败（忽略，继续 revectorize）: {}", e.getMessage());
        }
        chunkRepository.deleteByDocumentId(id);
        doc.setStatus("PENDING");
        doc.setErrorMessage(null);
        doc.setChunkCount(0);
        doc.setVectorCount(0);
        documentRepository.save(doc);
        pipeline.process(id);
        return doc;
    }

    @Override
    public KbDocument vectorizeDocument(Long id) {
        KbDocument doc = get(id);
        List<KbChunk> chunks = chunkRepository.findByDocumentIdOrderByChunkIndexAsc(id)
                .stream()
                .filter(c -> c.getChunkLevel() != null && c.getChunkLevel() == 3)
                .filter(c -> c.getVectorId() == null || c.getVectorId().isBlank())
                .toList();
        if (chunks.isEmpty()) {
            return doc;
        }
        List<String> ids = vectorizeChunksInternal(doc, chunks);
        refreshVectorCount(doc);
        log.info("文档向量化完成: doc={}, 新增向量={}", id, ids.size());
        return doc;
    }

    @Override
    public KbDocument vectorizeChunk(Long chunkId) {
        KbChunk chunk = chunkRepository.findById(chunkId)
                .orElseThrow(() -> new BusinessException(404, "分块不存在"));
        if (chunk.getVectorId() != null && !chunk.getVectorId().isBlank()) {
            throw new BusinessException("该分块已向量化，请勿重复操作");
        }
        if (chunk.getChunkLevel() == null || chunk.getChunkLevel() != 3) {
            throw new BusinessException("仅叶子分块（L3）支持向量化");
        }
        KbDocument doc = get(chunk.getDocumentId());
        List<String> ids = vectorizeChunksInternal(doc, List.of(chunk));
        refreshVectorCount(doc);
        log.info("分块向量化完成: chunkId={}, vectorId={}", chunkId, ids.isEmpty() ? null : ids.get(0));
        return doc;
    }

    private List<String> vectorizeChunksInternal(KbDocument doc, List<KbChunk> chunks) {
        if (chunks.isEmpty()) return List.of();
        List<Map<String, Object>> chunkMaps = new ArrayList<>(chunks.size());
        for (KbChunk c : chunks) {
            Map<String, Object> map = new HashMap<>();
            map.put("chunk_id", c.getChunkId());
            map.put("text", c.getContent());
            map.put("chunk_level", c.getChunkLevel());
            map.put("parent_chunk_id", c.getParentChunkId());
            map.put("root_chunk_id", c.getRootChunkId());
            map.put("chunk_idx", c.getChunkIndex());
            chunkMaps.add(map);
        }
        List<String> vectorIds = bridgeClient.vectorizeChunks(
                doc.getId(), doc.getName(),
                doc.getMimeType() == null ? "" : doc.getMimeType(),
                doc.getSourcePath() == null ? "" : doc.getSourcePath(),
                chunkMaps);
        for (int i = 0; i < chunks.size() && i < vectorIds.size(); i++) {
            chunks.get(i).setVectorId(vectorIds.get(i));
        }
        chunkRepository.saveAll(chunks);
        return vectorIds;
    }

    private void refreshVectorCount(KbDocument doc) {
        long count = chunkRepository.countByDocumentIdAndVectorIdIsNotNull(doc.getId());
        doc.setVectorCount((int) count);
        documentRepository.save(doc);
    }

    @Override
    public void delete(Long id) {
        KbDocument doc = get(id);
        // 删除 Milvus 向量（按 document_id）
        try {
            bridgeClient.deleteByDocument(id);
        } catch (Exception e) {
            log.warn("删除向量失败（继续删除 chunks 与文件）: {}", e.getMessage());
        }
        chunkRepository.deleteByDocumentId(id);
        // 删除源文件
        try {
            Path target = Paths.get(kbProps.getUploadDir()).resolve(doc.getSourcePath());
            Files.deleteIfExists(target);
        } catch (IOException e) {
            log.warn("删除源文件失败: {}", e.getMessage());
        }
        documentRepository.delete(doc);
    }

    @Override
    public KbDocument move(Long id, Long categoryId) {
        KbDocument doc = get(id);
        doc.setCategoryId(categoryId);
        return documentRepository.save(doc);
    }

    @Override
    public List<ChunkView> listChunks(Long documentId) {
        if (!documentRepository.existsById(documentId)) {
            throw new BusinessException(404, "文档不存在");
        }
        List<KbChunk> chunks = chunkRepository.findByDocumentIdOrderByChunkIndexAsc(documentId);
        List<ChunkView> views = new ArrayList<>(chunks.size());
        for (KbChunk c : chunks) {
            String preview = c.getContent() == null ? "" :
                    c.getContent().length() > 200 ? c.getContent().substring(0, 200) + "..." : c.getContent();
            views.add(new ChunkView(c.getId(), c.getChunkIndex(), c.getChunkLevel(), c.getCharCount(), preview, c.getContent(), c.getVectorId()));
        }
        return views;
    }

    private String sanitizeFileName(String name) {
        if (name == null) return "unknown";
        return name.replaceAll("[\\\\/]", "_").replaceAll("\\.\\.", "_");
    }
}