package com.lingyue.graph.service;

import com.lingyue.graph.entity.TextChunk;
import com.lingyue.graph.repository.VectorEmbeddingRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * 向量检索服务
 * 基于 pgvector 实现相似度检索
 *
 * @author lingyue
 * @since 2026-01-15
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class VectorSearchService {

    private final VectorEmbeddingRepository vectorEmbeddingRepository;
    private final OllamaEmbeddingService ollamaEmbeddingService;

    @Value("${rag.search.top-k:3}")
    private int defaultTopK;

    /**
     * 向量相似度检索（按文档ID筛选）
     *
     * @param query      查询文本
     * @param documentId 文档ID（可选，为null时全局检索）
     * @param topK       返回前K个结果
     * @return 相似文本块列表（包含相似度分数）
     */
    public List<SearchResult> search(String query, String documentId, Integer topK) {
        if (query == null || query.trim().isEmpty()) {
            log.warn("查询文本为空");
            return new ArrayList<>();
        }

        int limit = topK != null ? topK : defaultTopK;

        try {
            // 1. 将查询文本向量化
            float[] queryVector = ollamaEmbeddingService.embed(query);
            String queryEmbedding = ollamaEmbeddingService.vectorToString(queryVector);

            // 2. 执行向量检索
            List<Map<String, Object>> results;
            if (documentId != null && !documentId.isEmpty()) {
                results = vectorEmbeddingRepository.searchSimilarByDocument(
                        queryEmbedding, documentId, limit);
            } else {
                results = vectorEmbeddingRepository.searchSimilarGlobal(
                        queryEmbedding, limit);
            }

            // 3. 转换结果
            List<SearchResult> searchResults = new ArrayList<>();
            for (Map<String, Object> row : results) {
                SearchResult result = new SearchResult();
                result.setChunkId((String) row.get("chunk_id"));
                result.setDocumentId((String) row.get("document_id"));
                result.setContent((String) row.get("content"));
                result.setChunkIndex(((Number) row.get("chunk_index")).intValue());
                result.setSimilarity(((Number) row.get("similarity")).doubleValue());

                if (row.get("token_count") != null) {
                    result.setTokenCount(((Number) row.get("token_count")).intValue());
                }

                searchResults.add(result);
            }

            log.info("向量检索完成: query='{}', documentId={}, 结果数={}",
                    query.substring(0, Math.min(50, query.length())),
                    documentId, searchResults.size());

            return searchResults;

        } catch (Exception e) {
            log.error("向量检索失败: {}", e.getMessage(), e);
            return new ArrayList<>();
        }
    }

    /**
     * 全局向量检索（不限制文档）
     */
    public List<SearchResult> searchGlobal(String query, Integer topK) {
        return search(query, null, topK);
    }

    /**
     * 检索结果
     */
    @lombok.Data
    public static class SearchResult {
        private String chunkId;
        private String documentId;
        private String content;
        private Integer chunkIndex;
        private Integer tokenCount;
        private Double similarity;

        /**
         * 转换为 TextChunk 实体
         */
        public TextChunk toTextChunk() {
            TextChunk chunk = new TextChunk();
            chunk.setId(chunkId);
            chunk.setDocumentId(documentId);
            chunk.setContent(content);
            chunk.setChunkIndex(chunkIndex);
            chunk.setTokenCount(tokenCount);
            return chunk;
        }
    }
}
