package cn.iocoder.yudao.module.ai.service.knowledge;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
import cn.iocoder.yudao.framework.common.pojo.PageResult;
import cn.iocoder.yudao.module.ai.controller.admin.knowledge.vo.segment.*;
import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDO;
import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeDocumentMapper;
import cn.iocoder.yudao.module.ai.dal.mysql.knowledge.AiKnowledgeSegmentMapper;
import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchReqBO;
import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchRespBO;
import cn.iocoder.yudao.module.ai.service.model.AiModelService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.*;
import java.util.stream.Collectors;

import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.*;

@Slf4j
@Service
public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService {

    private static final String METADATA_KNOWLEDGE_ID = "knowledgeId";
    private static final String METADATA_DOCUMENT_ID = "documentId";
    private static final String METADATA_SEGMENT_ID = "segmentId";

    @Resource
    private AiKnowledgeSegmentMapper segmentMapper;
    @Resource
    private AiKnowledgeDocumentMapper documentMapper;
    @Resource
    private AiKnowledgeService knowledgeService;
    @Resource
    private AiModelService modelService;

    @Override
    public AiKnowledgeSegmentDO getSegment(Long id) {
        return segmentMapper.selectById(id);
    }

    @Override
    public AiKnowledgeSegmentDO validateSegmentExists(Long id) {
        AiKnowledgeSegmentDO segment = segmentMapper.selectById(id);
        if (segment == null) throw exception(KNOWLEDGE_SEGMENT_NOT_EXISTS);
        return segment;
    }

    @Override
    public PageResult<AiKnowledgeSegmentDO> getSegmentPage(AiKnowledgeSegmentPageReqVO pageReqVO) {
        return segmentMapper.selectPage(pageReqVO);
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public Long createSegment(AiKnowledgeSegmentSaveReqVO saveReqVO) {
        if (saveReqVO.getKnowledgeId() == null) throw exception(KNOWLEDGE_NOT_EXISTS);
        if (saveReqVO.getDocumentId() == null) throw exception(KNOWLEDGE_DOCUMENT_NOT_EXISTS);
        AiKnowledgeDO knowledge = knowledgeService.validateKnowledge(saveReqVO.getKnowledgeId());
        AiKnowledgeSegmentDO segment = new AiKnowledgeSegmentDO();
        segment.setKnowledgeId(saveReqVO.getKnowledgeId());
        segment.setDocumentId(saveReqVO.getDocumentId());
        segment.setContent(saveReqVO.getContent());
        segment.setContentLength(saveReqVO.getContent().length());
        segment.setTokens(estimateTokens(saveReqVO.getContent()));
        segment.setStatus(saveReqVO.getStatus() != null ? saveReqVO.getStatus()
                : CommonStatusEnum.ENABLE.getStatus());
        segment.setRetrievalCount(0);
        segment.setVectorId(AiKnowledgeSegmentDO.VECTOR_ID_EMPTY);
        segmentMapper.insert(segment);

        // 向量化
        VectorStore vectorStore = modelService.getOrCreateVectorStore(
                knowledge.getEmbeddingModelId(), buildMetadataFields());
        Map<String, Object> metadata = new HashMap<>();
        metadata.put(METADATA_KNOWLEDGE_ID, knowledge.getId());
        metadata.put(METADATA_DOCUMENT_ID, segment.getDocumentId());
        metadata.put(METADATA_SEGMENT_ID, segment.getId());
        Document doc = new Document(segment.getId().toString(), segment.getContent(), metadata);
        try {
            vectorStore.add(Collections.singletonList(doc));
            segment.setVectorId(segment.getId().toString());
            segmentMapper.updateById(segment);
        } catch (Exception e) {
            log.error("手动创建分段向量化失败，segmentId={}", segment.getId(), e);
            segment.setStatus(CommonStatusEnum.DISABLE.getStatus());
            segmentMapper.updateById(segment);
        }
        return segment.getId();
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void updateSegment(AiKnowledgeSegmentSaveReqVO saveReqVO) {
        AiKnowledgeSegmentDO segment = validateSegmentExists(saveReqVO.getId());
        AiKnowledgeDO knowledge = knowledgeService.validateKnowledge(segment.getKnowledgeId());
        // 删除旧向量
        if (StrUtil.isNotEmpty(segment.getVectorId()) && !AiKnowledgeSegmentDO.VECTOR_ID_EMPTY.equals(segment.getVectorId())) {
            VectorStore vectorStore = modelService.getOrCreateVectorStore(
                    knowledge.getEmbeddingModelId(), buildMetadataFields());
            try {
                vectorStore.delete(Collections.singletonList(segment.getVectorId()));
            } catch (Exception e) {
                log.warn("删除旧向量失败: {}", segment.getVectorId(), e);
            }
        }
        segment.setContent(saveReqVO.getContent());
        segment.setContentLength(saveReqVO.getContent().length());
        segment.setTokens(estimateTokens(saveReqVO.getContent()));
        if (saveReqVO.getStatus() != null) segment.setStatus(saveReqVO.getStatus());
        segmentMapper.updateById(segment);

        // 重新向量化
        VectorStore vectorStore = modelService.getOrCreateVectorStore(
                knowledge.getEmbeddingModelId(), buildMetadataFields());
        Map<String, Object> metadata = new HashMap<>();
        metadata.put(METADATA_KNOWLEDGE_ID, knowledge.getId());
        metadata.put(METADATA_DOCUMENT_ID, segment.getDocumentId());
        metadata.put(METADATA_SEGMENT_ID, segment.getId());
        Document doc = new Document(segment.getId().toString(), segment.getContent(), metadata);
        try {
            vectorStore.add(Collections.singletonList(doc));
            segment.setVectorId(segment.getId().toString());
            segmentMapper.updateById(segment);
        } catch (Exception e) {
            log.error("更新分段向量化失败，segmentId={}", segment.getId(), e);
            segment.setStatus(CommonStatusEnum.DISABLE.getStatus());
            segmentMapper.updateById(segment);
        }
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void deleteSegment(Long id) {
        AiKnowledgeSegmentDO segment = validateSegmentExists(id);
        AiKnowledgeDO knowledge = knowledgeService.validateKnowledge(segment.getKnowledgeId());
        if (StrUtil.isNotEmpty(segment.getVectorId()) && !AiKnowledgeSegmentDO.VECTOR_ID_EMPTY.equals(segment.getVectorId())) {
            VectorStore vectorStore = modelService.getOrCreateVectorStore(
                    knowledge.getEmbeddingModelId(), buildMetadataFields());
            try {
                vectorStore.delete(Collections.singletonList(segment.getVectorId()));
            } catch (Exception e) {
                log.warn("删除向量失败: {}", segment.getVectorId(), e);
            }
        }
        segmentMapper.deleteById(id);
    }

    @Override
    public void updateSegmentStatus(AiKnowledgeSegmentUpdateStatusReqVO updateStatusReqVO) {
        AiKnowledgeSegmentDO segment = validateSegmentExists(updateStatusReqVO.getId());
        segment.setStatus(updateStatusReqVO.getStatus());
        segmentMapper.updateById(segment);
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void deleteSegmentsByDocumentId(Long documentId) {
        List<AiKnowledgeSegmentDO> segments = segmentMapper.selectListByDocumentId(documentId);
        if (CollUtil.isEmpty(segments)) return;

        AiKnowledgeDO knowledge = knowledgeService.getKnowledge(segments.get(0).getKnowledgeId());
        VectorStore vectorStore = modelService.getOrCreateVectorStore(
                knowledge.getEmbeddingModelId(), buildMetadataFields());

        for (AiKnowledgeSegmentDO segment : segments) {
            if (StrUtil.isNotEmpty(segment.getVectorId()) && !AiKnowledgeSegmentDO.VECTOR_ID_EMPTY.equals(segment.getVectorId())) {
                try {
                    vectorStore.delete(Collections.singletonList(segment.getVectorId()));
                } catch (Exception e) {
                    log.warn("删除向量[{}]失败: {}", segment.getVectorId(), e.getMessage());
                }
            }
        }
        segmentMapper.delete(new com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper<AiKnowledgeSegmentDO>()
                .eq(AiKnowledgeSegmentDO::getDocumentId, documentId));
    }

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void saveSegments(List<AiKnowledgeSegmentDO> segments, Long knowledgeId, Long embeddingModelId) {
        if (CollUtil.isEmpty(segments)) return;

        AiKnowledgeDO knowledge = knowledgeService.validateKnowledge(knowledgeId);
        VectorStore vectorStore = modelService.getOrCreateVectorStore(embeddingModelId, buildMetadataFields());

        // 批量插入分段记录
        for (AiKnowledgeSegmentDO segment : segments) {
            segmentMapper.insert(segment);
        }

        // 转换为 Spring AI Document 并写入向量库
        List<Document> documents = new ArrayList<>();
        for (AiKnowledgeSegmentDO segment : segments) {
            Map<String, Object> metadata = new HashMap<>();
            metadata.put(METADATA_KNOWLEDGE_ID, knowledgeId);
            metadata.put(METADATA_DOCUMENT_ID, segment.getDocumentId());
            metadata.put(METADATA_SEGMENT_ID, segment.getId());
            Document doc = new Document(segment.getId().toString(), segment.getContent(), metadata);
            documents.add(doc);
        }

        try {
            vectorStore.add(documents);
            // 更新 vectorId（Milvus 返回的 ID 就是传入的 docId）
            for (AiKnowledgeSegmentDO segment : segments) {
                segment.setVectorId(segment.getId().toString());
                segmentMapper.updateById(segment);
            }
            log.info("批量写入向量库成功，知识库: {}, 分段数: {}", knowledgeId, segments.size());
        } catch (Exception e) {
            log.error("写入向量库失败，知识库: {}", knowledgeId, e);
            for (AiKnowledgeSegmentDO segment : segments) {
                segment.setStatus(CommonStatusEnum.DISABLE.getStatus());
                segmentMapper.updateById(segment);
            }
            throw new RuntimeException("向量库写入失败: " + e.getMessage(), e);
        }
    }

    @Override
    public List<AiKnowledgeSegmentSearchRespBO> searchSegments(AiKnowledgeSegmentSearchReqBO searchReqBO) {
        AiKnowledgeDO knowledge = knowledgeService.validateKnowledge(searchReqBO.getKnowledgeId());
        VectorStore vectorStore = modelService.getOrCreateVectorStore(
                knowledge.getEmbeddingModelId(), buildMetadataFields());

        int topK = searchReqBO.getTopK() != null ? searchReqBO.getTopK() : knowledge.getTopK();
        double similarityThreshold = searchReqBO.getSimilarityThreshold() != null
                ? searchReqBO.getSimilarityThreshold() : knowledge.getSimilarityThreshold();

        Filter.Expression filterExpression = new FilterExpressionBuilder()
                .eq(METADATA_KNOWLEDGE_ID, searchReqBO.getKnowledgeId())
                .build();

        SearchRequest request = SearchRequest.builder()
                .query(searchReqBO.getContent())
                .topK(topK)
                .similarityThreshold(similarityThreshold)
                .filterExpression(filterExpression)
                .build();

        log.info("向量检索: query={}, topK={}, threshold={}, filter=kid={}",
                searchReqBO.getContent(), topK, similarityThreshold, searchReqBO.getKnowledgeId());
        List<Document> results = vectorStore.similaritySearch(request);
        log.info("向量检索结果数: {}", results != null ? results.size() : 0);
        if (CollUtil.isEmpty(results)) return Collections.emptyList();

        // 更新检索次数
        List<String> vectorIds = results.stream()
                .map(Document::getId)
                .filter(StrUtil::isNotEmpty)
                .collect(Collectors.toList());
        if (CollUtil.isNotEmpty(vectorIds)) {
            List<AiKnowledgeSegmentDO> hitSegments = segmentMapper.selectListByVectorIds(vectorIds);
            if (CollUtil.isNotEmpty(hitSegments)) {
                List<Long> hitIds = hitSegments.stream().map(AiKnowledgeSegmentDO::getId).collect(Collectors.toList());
                segmentMapper.updateRetrievalCountIncrByIds(hitIds);
                Set<Long> docIds = hitSegments.stream().map(AiKnowledgeSegmentDO::getDocumentId).collect(Collectors.toSet());
                documentMapper.updateRetrievalCountIncr(docIds);
            }
        }

        // 收集文档 ID 并批量查询文档名称
        Set<Long> resultDocIds = results.stream()
                .map(doc -> convertToLong(doc.getMetadata() != null
                        ? doc.getMetadata().get(METADATA_DOCUMENT_ID) : null))
                .filter(id -> id != null)
                .collect(Collectors.toSet());
        final Map<Long, String> docNameMap;
        if (CollUtil.isNotEmpty(resultDocIds)) {
            docNameMap = documentMapper.selectBatchIds(resultDocIds).stream()
                    .collect(Collectors.toMap(
                            cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO::getId,
                            cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO::getName,
                            (a, b) -> a));
        } else {
            docNameMap = Collections.emptyMap();
        }

        return results.stream()
                .map(doc -> {
                    AiKnowledgeSegmentSearchRespBO bo = new AiKnowledgeSegmentSearchRespBO();
                    bo.setContent(doc.getText());
                    bo.setScore(doc.getScore() != null ? doc.getScore() : 0.0);
                    if (doc.getMetadata() != null) {
                        bo.setId(convertToLong(doc.getMetadata().get(METADATA_SEGMENT_ID)));
                        bo.setDocumentId(convertToLong(doc.getMetadata().get(METADATA_DOCUMENT_ID)));
                        bo.setDocumentName(docNameMap.getOrDefault(bo.getDocumentId(), "未知文档"));
                        bo.setKnowledgeId(convertToLong(doc.getMetadata().get(METADATA_KNOWLEDGE_ID)));
                        log.debug("检索结果元数据: id={}, docId={}, kid={}, metadataKeys={}",
                                bo.getId(), bo.getDocumentId(), bo.getKnowledgeId(),
                                doc.getMetadata().keySet());
                    }
                    if (bo.getContent() != null) {
                        bo.setContentLength(bo.getContent().length());
                        bo.setTokens(estimateTokens(bo.getContent()));
                    }
                    return bo;
                })
                .collect(Collectors.toList());
    }

    private Map<String, Class<?>> buildMetadataFields() {
        return MapUtil.<String, Class<?>>builder()
                .put(METADATA_KNOWLEDGE_ID, Long.class)
                .put(METADATA_DOCUMENT_ID, Long.class)
                .put(METADATA_SEGMENT_ID, Long.class)
                .build();
    }

    /**
     * 将 Milvus 返回的元数据值转为 Long。
     * Gson 反序列化 JSON 数字默认为 Double，需要兼容多种类型。
     */
    private Long convertToLong(Object value) {
        if (value == null) return null;
        if (value instanceof Long v) return v;
        if (value instanceof Integer v) return v.longValue();
        if (value instanceof Double v) return v.longValue();
        if (value instanceof Float v) return v.longValue();
        if (value instanceof Number v) return v.longValue();
        if (value instanceof String v) {
            try { return Long.valueOf(v); } catch (NumberFormatException e) { return null; }
        }
        log.warn("无法转换元数据值类型: {} = {}", value.getClass().getName(), value);
        return null;
    }

    private Integer estimateTokens(String text) {
        if (StrUtil.isEmpty(text)) return 0;
        int chineseChars = 0, englishWords = 0;
        for (char c : text.toCharArray()) {
            if (c >= 0x4E00 && c <= 0x9FA5) chineseChars++;
        }
        for (String word : text.split("\\s+")) {
            if (word.matches(".*[a-zA-Z].*")) englishWords++;
        }
        return chineseChars + (int) (englishWords * 1.3);
    }

}