package com.lingyue.graph.repository;

import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.lingyue.graph.entity.VectorEmbedding;
import org.apache.ibatis.annotations.Delete;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.annotations.Param;
import org.apache.ibatis.annotations.Select;

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

/**
 * 向量嵌入数据访问层
 *
 * @author lingyue
 * @since 2026-01-15
 */
@Mapper
public interface VectorEmbeddingRepository extends BaseMapper<VectorEmbedding> {

    /**
     * 根据分块ID查询向量
     */
    @Select("SELECT * FROM vector_embeddings WHERE chunk_id = #{chunkId}")
    VectorEmbedding findByChunkId(@Param("chunkId") String chunkId);

    /**
     * 根据分块ID删除向量
     */
    @Delete("DELETE FROM vector_embeddings WHERE chunk_id = #{chunkId}")
    int deleteByChunkId(@Param("chunkId") String chunkId);

    /**
     * 向量相似度检索（按文档ID筛选）
     * 使用 pgvector 的余弦距离操作符
     *
     * @param queryEmbedding 查询向量（格式："[0.1, 0.2, ...]"）
     * @param documentId     文档ID
     * @param limit          返回数量
     * @return 相似文本块列表
     */
    @Select("""
            SELECT 
                tc.id AS chunk_id,
                tc.document_id,
                tc.content,
                tc.chunk_index,
                tc.token_count,
                1 - (ve.embedding <=> #{queryEmbedding}::vector) AS similarity
            FROM text_chunks tc
            JOIN vector_embeddings ve ON tc.id = ve.chunk_id
            WHERE tc.document_id = #{documentId}
            ORDER BY ve.embedding <=> #{queryEmbedding}::vector
            LIMIT #{limit}
            """)
    List<Map<String, Object>> searchSimilarByDocument(
            @Param("queryEmbedding") String queryEmbedding,
            @Param("documentId") String documentId,
            @Param("limit") int limit
    );

    /**
     * 全局向量相似度检索（不限制文档）
     *
     * @param queryEmbedding 查询向量
     * @param limit          返回数量
     * @return 相似文本块列表
     */
    @Select("""
            SELECT 
                tc.id AS chunk_id,
                tc.document_id,
                tc.content,
                tc.chunk_index,
                tc.token_count,
                1 - (ve.embedding <=> #{queryEmbedding}::vector) AS similarity
            FROM text_chunks tc
            JOIN vector_embeddings ve ON tc.id = ve.chunk_id
            ORDER BY ve.embedding <=> #{queryEmbedding}::vector
            LIMIT #{limit}
            """)
    List<Map<String, Object>> searchSimilarGlobal(
            @Param("queryEmbedding") String queryEmbedding,
            @Param("limit") int limit
    );

    /**
     * 检查分块是否已有向量
     */
    @Select("SELECT COUNT(*) > 0 FROM vector_embeddings WHERE chunk_id = #{chunkId}")
    boolean existsByChunkId(@Param("chunkId") String chunkId);
}
