"""
检索引擎业务逻辑服务
"""
from math import ceil
from typing import List, Optional, Tuple, Dict
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, func, or_
from datetime import datetime
import uuid
import random
import json
import hashlib
import math

from app.sample.models.search_engine import SearchEngine
from app.sample.schemas.search_engine import (
    SearchEngineCreate,
    SearchEngineUpdate,
    KBSearchRequest,
    KBSearchResultItem,
    KBSearchResponse
)
from app.schemas.base import PaginationSchema
from app.services.milvus_service import milvus_service
from app.utils.vector_utils import text_to_vector_algo
import logging

logger = logging.getLogger(__name__)

class SearchEngineService:
    
    async def search_kb(self, db: AsyncSession, payload: KBSearchRequest) -> KBSearchResponse:
        """
        知识库搜索 (基于算法向量)
        """
        original_kb_id = payload.kb_id 
        collection_name = original_kb_id
        
        # 0. 尝试从数据库解析 kb_id 为 collection_name (如果传的是知识库ID)
        from sqlalchemy import text
        try:
            # 简单判断是否是 UUID 格式或数字 ID，尝试查询数据库
            kb_query = text(
                "SELECT collection_name_children, collection_name_parent "
                "FROM t_samp_knowledge_base "
                "WHERE id = :kb_id "
                "   OR collection_name_children = :kb_id "
                "   OR collection_name_parent = :kb_id"
            )
            kb_res = await db.execute(kb_query, {"kb_id": original_kb_id})
            kb_row = kb_res.fetchone()
            if kb_row:
                collection_name = kb_row[0] or kb_row[1] or collection_name
                logging.info(f"Resolved kb_id {original_kb_id} to collection_name: {collection_name}")
        except Exception as db_err:
            logging.warning(f"Failed to resolve kb_id {original_kb_id} from database: {db_err}")

        # 1. 确定实际检索的集合名称 (支持父子段)
        child_col = f"{collection_name}_child"
        parent_col = f"{collection_name}_parent"
        
        is_pdr = False
        kb_id = collection_name # 默认搜索原集合
        
        # 优先检查是否存在子段集合，如果存在则进入父子段检索模式 (PDR)
        if milvus_service.has_collection(child_col):
            kb_id = child_col
            is_pdr = True
            logging.info(f"Detected PDR collection for {collection_name}, searching in {child_col}")
        elif not milvus_service.has_collection(collection_name):
            return KBSearchResponse(results=[], total=0)

        try:
            state = milvus_service.get_collection_state(kb_id)
            if state != "Loaded":
                milvus_service.set_collection_state(kb_id, "load")
        except Exception:
            pass
            
        # 1. 使用算法生成向量 (替代 Embedding 模型)
        # 尝试从 Milvus collection 获取向量维度，动态匹配维度
        # 这样相同的查询词会生成相同的向量，具备了基本的检索能力
        try:
            collection_detail = milvus_service.get_collection_detail(kb_id)
        except Exception:
            collection_detail = None

        dim = None
        if collection_detail and isinstance(collection_detail, dict):
            fields = collection_detail.get("fields", []) or []
            for f in fields:
                if not isinstance(f, dict):
                    continue
                params = f.get("params") or {}
                if params and params.get("dim"):
                    try:
                        dim = int(params.get("dim"))
                        break
                    except Exception:
                        dim = None

        # 回退默认维度：与系统 embedding 维度保持一致（避免向量维度不匹配导致检索报错）
        if not dim:
            dim = milvus_service.DENSE_DIM

        # 选择 Milvus 向量字段名（anns_field），字段名可能不是固定的 "vector"，也可能叫 'dense'/'denser' 等
        anns_field = "dense"
        if collection_detail and isinstance(collection_detail, dict):
            fields = collection_detail.get("fields", []) or []
            # 优先寻找有 params.dim 的向量字段
            for f in fields:
                if not isinstance(f, dict):
                    continue
                params = f.get("params") or {}
                if params and params.get("dim") and f.get("name"):
                    anns_field = f.get("name")
                    try:
                        dim = int(params.get("dim"))
                    except Exception:
                        pass
                    break

            # 若未找到带 dim 的字段，尝试匹配常见的向量字段名或字段类型包含 "vector"
            if anns_field == "dense":
                for f in fields:
                    if not isinstance(f, dict):
                        continue
                    fname = (f.get("name") or "")
                    ftype = str(f.get("type") or "").lower()
                    if fname and fname.lower() in ("vector", "denser", "dense", "embedding", "embeddings"):
                        anns_field = fname
                        break
                    if fname and "vector" in ftype:
                        anns_field = fname
                        break

        # 1. 向量搜索 (Dense Retrieval)
        # 默认使用 Hybrid 混合检索逻辑，但为了简化，这里先保留向量检索的核心
        # 如果 metric_type 指定为 hybrid，则可能需要结合关键词搜索等
        # 目前后端实现主要是基于 Milvus 的 ANN 搜索
        
        # 强制使用 hybrid 混合检索模式作为基础（结合关键词匹配和向量相似度）
        # 除非用户明确指定了其他度量方式（通常不会）
        requested_metric = payload.metric_type
        use_hybrid = False
        
        # 只有当 metric_type 为 None 或者特定值时才尝试混合检索
        # 或者我们可以认为只要不指定，就优先尝试混合
        if not requested_metric or requested_metric.lower() == 'hybrid':
             use_hybrid = True
        
        search_params = {
            "metric_type": "L2", # 默认内部计算用 L2
            "params": {"nprobe": 10},
        }
        
        # 如果前端指定了 metric_type (虽然前端现在默认 hybrid，但保留参数兼容性)
        if payload.metric_type and payload.metric_type.upper() != 'HYBRID':
             search_params["metric_type"] = payload.metric_type
        
        # 2. 构建过滤表达式
        expr_list = []

        metadata_type_map: Dict[str, str] = {}
        try:
            from app.sample.models.knowledge_base import KnowledgeBase
            from app.sample.models.metadata import SampleMetadata

            kb_stmt = select(KnowledgeBase).where(
                or_(
                    KnowledgeBase.collection_name_children == collection_name,
                    KnowledgeBase.collection_name_parent == collection_name
                ),
                KnowledgeBase.is_deleted == 0
            )
            kb_res = await db.execute(kb_stmt)
            kb_obj = kb_res.scalars().first()
            if kb_obj:
                meta_stmt = select(SampleMetadata.field_en_name, SampleMetadata.field_type).where(
                    SampleMetadata.knowledge_base_id == kb_obj.id
                )
                meta_res = await db.execute(meta_stmt)
                rows = meta_res.all()
                metadata_type_map = {str(r[0]): str(r[1]) for r in rows if r and r[0]}
        except Exception:
            metadata_type_map = {}

        def build_eq_expr(target_field: str, value: str, field_type: Optional[str], is_top_level: bool) -> str:
            field_type_norm = (field_type or "").strip().lower()
            looks_numeric = False
            try:
                float(value)
                looks_numeric = True
            except Exception:
                looks_numeric = False

            if field_type_norm == "num" or (looks_numeric and not is_top_level):
                try:
                    num_val = float(value)
                    num_expr = f"{target_field} == {int(num_val) if num_val.is_integer() else num_val}"
                    if is_top_level and field_type_norm == "num":
                        return num_expr
                    str_expr = f'{target_field} == "{value}"'
                    return f"({num_expr} or {str_expr})"
                except Exception:
                    return f'{target_field} == "{value}"'

            return f'{target_field} == "{value}"'
        
        # 兼容旧的单一字段过滤
        if payload.metadata_field and payload.metadata_value:
            safe_field = payload.metadata_field.replace("'", "").replace('"', "").strip()
            safe_value = payload.metadata_value.replace("'", "").replace('"', "").strip()
            
            if safe_field and safe_value:
                # [Updated] 处理 metadata 嵌套字段查询
                # 如果字段名不包含 [" 或 .，假设它是 metadata 中的一个键
                # 用户通常传 "field_en_name"，我们需要转成 metadata["field_en_name"]
                
                # 简单判断：如果不是顶层字段，且 Milvus collection 只有 metadata JSON 字段
                # 我们假设所有过滤都是针对 metadata 的
                # 为了安全，先获取字段列表
                
                is_top_level = False
                if collection_detail and isinstance(collection_detail, dict):
                    fields = collection_detail.get("fields", [])
                    for f in fields:
                        if f['name'] == safe_field:
                            is_top_level = True
                            break
                
                target_field = safe_field
                if not is_top_level:
                     target_field = f'metadata["{safe_field}"]'

                if is_top_level:
                    expr_list.append(build_eq_expr(target_field, safe_value, metadata_type_map.get(safe_field), True))
                else:
                    alt_field = f'metadata["metadata"]["{safe_field}"]'
                    expr_main = build_eq_expr(target_field, safe_value, metadata_type_map.get(safe_field), False)
                    expr_alt = build_eq_expr(alt_field, safe_value, metadata_type_map.get(safe_field), False)
                    expr_list.append(f"({expr_main} or {expr_alt})")
        
        # 处理新的多重过滤
        if payload.filters:
            for f in payload.filters:
                safe_field = f.field.replace("'", "").replace('"', "").strip()
                raw_value = (f.value or "").strip()
                safe_value = raw_value.replace("'", "").replace('"', "").strip()
                
                if safe_field and safe_value:
                    # [Special Case] 文档名称过滤 (doc_name_in)
                    # 前端传递的是 "doc_name_in", value 是 JSON 数组字符串 (e.g. '["doc1", "doc2"]')
                    if safe_field == 'doc_name_in':
                        try:
                            # doc_name_in 的 value 必须保留引号，否则不是合法 JSON
                            parse_value = raw_value
                            doc_names = json.loads(parse_value)
                            if doc_names and isinstance(doc_names, list):
                                # 构建 OR 条件: (metadata["doc_name"] == "A" || metadata["doc_name"] == "B")
                                # 注意：Milvus 字段可能是 doc_name, file_name, title, source
                                # 我们假设 metadata 中统一用 doc_name 或 file_name
                                # 更稳妥的方式是用 IN 操作符 (Milvus 支持 IN)
                                # target_field in ["A", "B"]
                                
                                # 确定目标字段名
                                target_field = 'doc_name' # 默认
                                # 尝试检测实际使用的文档名字段
                                possible_fields = ['doc_name', 'file_name', 'title', 'source']
                                detected_field = None
                                
                                # 简单逻辑：所有可能的字段都尝试匹配 (使用 OR 连接)
                                # (metadata["doc_name"] in [...]) or (metadata["file_name"] in [...]) ...
                                # 但 Milvus 的 OR 性能可能一般，且语法较长
                                # 既然我们是根据 metadata 生成的，通常字段名是固定的
                                # 这里假设使用 metadata["file_name"] 或 metadata["doc_name"]
                                # 我们可以生成一个复杂的表达式
                                
                                # 构造 list string: ["A", "B"]
                                val_list_str = json.dumps(doc_names, ensure_ascii=False)
                                
                                # 优先尝试 metadata["file_name"] 和 metadata["doc_name"]
                                sub_exprs = []
                                sub_exprs.append(f'metadata["file_name"] in {val_list_str}')
                                sub_exprs.append(f'metadata["doc_name"] in {val_list_str}')
                                sub_exprs.append(f'metadata["title"] in {val_list_str}')
                                sub_exprs.append(f'metadata["metadata"]["file_name"] in {val_list_str}')
                                sub_exprs.append(f'metadata["metadata"]["doc_name"] in {val_list_str}')
                                sub_exprs.append(f'metadata["metadata"]["title"] in {val_list_str}')
                                
                                # 组合成 (A or B or C)
                                # 注意：如果某些字段不存在，Milvus 可能会报错吗？
                                # JSON 字段内访问不存在的 key 通常返回 null/empty，不会报错
                                # 所以用 OR 连接是安全的
                                
                                doc_filter_expr = f"({' or '.join(sub_exprs)})"
                                expr_list.append(doc_filter_expr)
                                continue # 处理完特殊字段，跳过后续通用逻辑
                        except Exception as e:
                            logging.warning(f"Failed to parse doc_name_in filter: {e}")
                            continue

                    # [Updated] 同样处理嵌套字段
                    is_top_level = False
                    if collection_detail and isinstance(collection_detail, dict):
                        fields = collection_detail.get("fields", [])
                        for field_def in fields:
                            if field_def['name'] == safe_field:
                                is_top_level = True
                                break
                    
                    target_field = safe_field
                    if not is_top_level:
                         target_field = f'metadata["{safe_field}"]'

                    # [Fix] 统一将 metadata 值视为字符串查询
                    if is_top_level:
                        expr_list.append(build_eq_expr(target_field, safe_value, metadata_type_map.get(safe_field), True))
                    else:
                        alt_field = f'metadata["metadata"]["{safe_field}"]'
                        expr_main = build_eq_expr(target_field, safe_value, metadata_type_map.get(safe_field), False)
                        expr_alt = build_eq_expr(alt_field, safe_value, metadata_type_map.get(safe_field), False)
                        expr_list.append(f"({expr_main} or {expr_alt})")
        
        # 组合所有条件 (使用 AND)
        expr = " and ".join(expr_list) if expr_list else ""
        
        # 3. 确定分页参数
        page = payload.page if payload.page and payload.page > 0 else 1
        page_size = payload.page_size if payload.page_size and payload.page_size > 0 else 10
        offset = (page - 1) * page_size
        limit = page_size

        has_query = payload.query and payload.query.strip()
        
        if not has_query:
            # --- 分支 A: 纯标量查询 (无关键词) ---
            logger = logging.getLogger(__name__)
            logger.info(f"Scalar query mode for KB={kb_id}, expr={expr}")
            try:
                # 1. 获取总数
                total = 0
                count_expr = expr if expr else ""
                # 如果没有表达式，默认查所有 (需要满足 Milvus 语法)
                if not count_expr:
                    # 简单获取 stats
                    stats = milvus_service.client.get_collection_stats(collection_name=kb_id)
                    total = int(stats.get("row_count", 0)) if isinstance(stats, dict) else 0
                else:
                    # 带条件 count
                    try:
                        res_cnt = milvus_service.client.query(kb_id, filter=count_expr, output_fields=["count(*)"])
                        if res_cnt:
                            total = int(res_cnt[0].get("count(*)") or 0)
                    except Exception as e:
                        logger.warning(f"Scalar count(*) failed for KB={kb_id}, expr={count_expr}: {e}")
                        total = 0
                
                # 2. 分页查询
                # 如果没有 expr，Milvus query 需要一个 valid expression
                # 尝试用 id >= 0，前提是 id 是 int。如果是 varchar，用 id != ""
                query_expr = expr
                if not query_expr:
                    # 获取主键字段名和类型
                    pk_field = "pk" # 默认
                    is_int = True
                    try:
                        desc = milvus_service.client.describe_collection(kb_id)
                        if isinstance(desc, dict) and 'fields' in desc:
                            for f in desc['fields']:
                                if f.get('primary_key') or f.get('is_primary'):
                                    pk_field = f.get('name')
                                    # Type 5 is INT64, 21 is VARCHAR. 
                                    if f.get('type') == 21 or str(f.get('type')).upper() == 'VARCHAR':
                                        is_int = False
                                    break
                    except:
                        pass
                    
                    query_expr = f'{pk_field} >= 0' if is_int else f'{pk_field} != ""'

                res_page = milvus_service.client.query(
                    collection_name=kb_id,
                    filter=query_expr,
                    output_fields=["*"],
                    limit=limit,
                    offset=offset
                )
                
                formatted_results = []
                for item in res_page:
                    item_metadata = item.get('metadata') or {}
                    if isinstance(item_metadata, str):
                        try:
                            item_metadata = json.loads(item_metadata)
                        except Exception:
                            item_metadata = {}

                    # PDR 模式内容获取 (可选)
                    item_content = item.get('text') or item.get('content') or item.get('page_content') or ""
                    parent_id = item.get("parent_id") or item_metadata.get("parent_id") or ""

                    doc_name = (
                        item_metadata.get('doc_name')
                        or item_metadata.get('file_name')
                        or item_metadata.get('title')
                        or item_metadata.get('source')
                        or item.get('file_name')
                        or item.get('title')
                        or item.get('source')
                        or "未知文档"
                    )

                    formatted_results.append(KBSearchResultItem(
                        id=str(item.get('pk') or item.get('id')),
                        kb_name=original_kb_id,
                        doc_name=doc_name,
                        content=item_content,
                        meta_info=str(item_metadata),
                        document_id=str(item.get("document_id") or ""),
                        parent_id=str(parent_id) if parent_id is not None else None,
                        metadata=item_metadata,
                        score=0
                    ))
                
                return KBSearchResponse(results=formatted_results, total=total)

            except Exception as e:
                logging.error(f"Scalar query failed: {e}")
                raise ValueError(f"元数据过滤条件查询失败：{e}")

        # --- 分支 B: 关键词/混合检索 (有关键词) ---
        # 优先尝试用 text like 做关键词召回（比算法向量更贴合“相关性”）；无结果再回退向量检索
        query_text = (payload.query or "").strip()
        if use_hybrid and query_text:
            logger = logging.getLogger(__name__)
            safe_q = query_text.replace('"', "").replace("'", "").strip()
            like_expr = f'text like "%{safe_q}%"'
            combined_expr = f"({expr}) and ({like_expr})" if expr else like_expr

            page = payload.page if payload.page and payload.page > 0 else 1
            page_size = payload.page_size if payload.page_size and payload.page_size > 0 else 10
            offset = (page - 1) * page_size
            limit = page_size

            try:
                total = 0
                try:
                    res_cnt = milvus_service.client.query(kb_id, filter=combined_expr, output_fields=["count(*)"])
                    if res_cnt:
                        total = int(res_cnt[0].get("count(*)") or 0)
                except Exception as e:
                    logger.warning(f"Keyword count(*) failed for KB={kb_id}, expr={combined_expr}: {e}")
                    total = 0

                fetch_limit = min(500, max(limit * 3, offset + limit))
                res_page = milvus_service.client.query(
                    collection_name=kb_id,
                    filter=combined_expr,
                    output_fields=["*"],
                    limit=fetch_limit,
                    offset=0
                )

                formatted_results = []
                for item in res_page or []:
                    item_metadata = item.get('metadata') or {}
                    if isinstance(item_metadata, str):
                        try:
                            item_metadata = json.loads(item_metadata)
                        except Exception:
                            item_metadata = {}

                    content = item.get('text') or item.get('content') or item.get('page_content') or ""
                    doc_name = (
                        item_metadata.get('doc_name')
                        or item_metadata.get('file_name')
                        or item_metadata.get('title')
                        or item_metadata.get('source')
                        or item.get('file_name')
                        or item.get('title')
                        or item.get('source')
                        or "未知文档"
                    )
                    parent_id = item.get("parent_id") or item_metadata.get("parent_id") or ""
                    document_id = item.get("document_id") or item_metadata.get("document_id") or ""

                    occ = 0
                    try:
                        occ = (content or "").count(query_text)
                    except Exception:
                        occ = 0
                    score = 60.0 + min(40.0, float(min(4, occ)) * 10.0)
                    if query_text and doc_name and query_text in str(doc_name):
                        score = min(100.0, score + 10.0)

                    formatted_results.append(KBSearchResultItem(
                        id=str(item.get('pk') or item.get('id')),
                        kb_name=original_kb_id,
                        doc_name=doc_name,
                        content=content,
                        meta_info=str(item_metadata),
                        document_id=str(document_id) if document_id is not None else None,
                        parent_id=str(parent_id) if parent_id is not None else None,
                        metadata=item_metadata,
                        score=round(score, 2)
                    ))

                formatted_results.sort(key=lambda r: (r.score or 0), reverse=True)
                paged = formatted_results[offset: offset + limit]
                if total == 0:
                    total = len(formatted_results)
                return KBSearchResponse(results=paged, total=total)
            except Exception as e:
                logger.warning(f"Keyword search fallback to vector due to error: {e}")

        # 回退到向量检索（算法向量）
        query_vector = text_to_vector_algo(payload.query, dim=dim)
        
        # 检测 collection 使用的 metric (恢复这部分逻辑，因为后续 search 需要)
        metric_type = None
        # 优先从 collection_detail 检测真实 metric
        if collection_detail and isinstance(collection_detail, dict):
            indices = collection_detail.get("indices") or []
            if isinstance(indices, list) and len(indices) > 0:
                for idx in indices:
                    try:
                        mt = idx.get("metric_type") or idx.get("metric")
                        if mt:
                            metric_type = str(mt).upper()
                            break
                    except Exception:
                        continue
        
        # 尝试从 properties 中读取
        if not metric_type and collection_detail and isinstance(collection_detail, dict):
            props = collection_detail.get("properties") or {}
            if isinstance(props, dict):
                mt = props.get("metric_type") or props.get("metric")
                if mt:
                    metric_type = str(mt).upper()
        
        actual_search_metric = metric_type
        if not actual_search_metric:
             # 如果无法检测到 collection metric (如无索引)，则可以使用用户请求的或默认 L2
             actual_search_metric = requested_metric if requested_metric and requested_metric.upper() != 'HYBRID' else "L2"
        
        metric_type = actual_search_metric
        
        logger = logging.getLogger(__name__)
        logger.info(f"Search KB={kb_id} using anns_field={anns_field}, dim={dim}, metric={metric_type} (requested={requested_metric})")

        # 3. 执行 Milvus 搜索
        try:
            # 使用 collection 实际的 metric_type 作为检索度量，避免 mismatch 错误
            # metric_type 已在上面检测并存放于变量 metric_type
            search_params = {
                "metric_type": metric_type,
                "params": {"nprobe": 10}
            }

            # 如果 top_k <= 0 或未指定，解释为返回该 collection 中的所有文段
            # 优先使用 page/page_size 计算 limit 和 offset
            page = payload.page if payload.page and payload.page > 0 else 1
            page_size = payload.page_size if payload.page_size and payload.page_size > 0 else 10
            
            # 如果 payload 中有 top_k 且未传 page_size (或者 page_size 是默认值)，可以使用 top_k 覆盖 page_size
            # 但这里为了清晰，优先使用 page_size
            
            offset = (page - 1) * page_size
            limit = page_size
            
            # Milvus 对 limit + offset 有限制 (通常 16384)，这里做个简单的保护
            if offset + limit > 16384:
                # 如果超出深度分页限制，可能需要提示或截断
                # 这里暂时不做处理，让 Milvus 报错或自行截断
                pass

            # 获取集合总数用于分页显示 (total)
            collection_count = 0
            if collection_detail and isinstance(collection_detail, dict):
                collection_count = collection_detail.get("entity_count") or 0
            
            if not collection_count:
                try:
                    stats = milvus_service.client.get_collection_stats(collection_name=kb_id)
                    collection_count = int(stats.get("row_count")) if isinstance(stats, dict) and stats.get("row_count") else 0
                except Exception:
                    collection_count = 0

            # 如果是按照 top_k 逻辑 (不传 page/page_size)，保留旧逻辑 (top_k 即 limit, offset=0)
            # 但现在 Schema 默认 page=1, page_size=10，所以总是走分页逻辑
            
            try:
                # 暂时仅使用向量检索，关闭混合检索以保证相似度与查询内容强相关
                results = milvus_service.client.search(
                    collection_name=kb_id,
                    data=[query_vector],
                    anns_field=anns_field,
                    search_params=search_params,
                    limit=limit,
                    offset=offset, # 添加 offset 支持分页
                    filter=expr if expr else "",
                    output_fields=["*"] 
                )
            except Exception as milvus_err:
                # 捕获 Milvus 异常，常见原因包括 metric mismatch
                logger.error(f"Milvus search failed for collection={kb_id}, metric_requested={metric_type}, anns_field={anns_field}: {milvus_err}")
                
                # Retry Logic: 如果是因为 metric 不匹配，解析错误信息中的 expected metric 并重试
                error_msg = str(milvus_err)
                if "metric type not match" in error_msg:
                    import re
                    # 匹配 expected=COSINE 或 expected='COSINE' 等格式
                    # 支持 COSINE, L2, IP, BM25 等
                    match = re.search(r"expected\s*=\s*['\"]?([A-Za-z0-9_]+)['\"]?", error_msg)
                    if match:
                        correct_metric = match.group(1).upper()
                        logger.warning(f"Detected metric mismatch. Retrying with correct metric: {correct_metric}")
                        
                        # 更新 metric_type 并重试搜索
                        search_params["metric_type"] = correct_metric
                        # 同时也需要更新后续计算分数所用的 metric_type 变量，以便正确计算相似度
                        metric_type = correct_metric
                        
                        # 特殊处理: BM25 可能需要 sparse vector 或其他参数，但 Milvus search 接口应该是一致的
                        # 如果是 BM25，可能 anns_field 也要调整（通常 BM25 用 sparse vector）
                        # 但这里假设 anns_field 是正确的，只是 metric 不对
                        
                        results = milvus_service.client.search(
                            collection_name=kb_id,
                            data=[query_vector],
                            anns_field=anns_field,
                            search_params=search_params,
                            limit=limit,
                            offset=offset, # 同样加上 offset
                            filter=expr if expr else "",
                            output_fields=["*"] 
                        )
                    else:
                        raise
                else:
                    raise
            
            # 4. 格式化结果
            formatted_results = []
            for hits in results:
                for hit in hits:
                    # 过滤低相似度结果 (算法生成的向量相似度可能较低，阈值可适当调低)
                    # if hit.score < payload.score_threshold:
                    #     continue
                        
                    entity = hit.entity

                    content = entity.get("text") or entity.get("content") or entity.get("page_content") or ""
                    
                    # PDR 模式：从父表获取内容 (如果原逻辑有 is_pdr 标记，优先用；否则尝试通用逻辑)
                    # [Updated] 通用逻辑：无论是否 PDR 模式，只要有 parent_id 且能找到父表，就尝试获取父段内容
                    # 首先确定父表名：
                    # 1. 如果有 PDR 标记，则 parent_col 已知
                    # 2. 如果没有，需要查询 KnowledgeBase 表获取 collection_name_parent (父表)
                    
                    parent_id = entity.get("parent_id")
                    if not parent_id and isinstance(entity.get("metadata"), dict):
                        parent_id = entity.get("metadata", {}).get("parent_id")
                    
                    # 尝试从 metadata 中获取 parent_id，如果 entity 直接有也行
                    if not parent_id and "metadata" in entity and isinstance(entity["metadata"], dict):
                        parent_id = entity["metadata"].get("parent_id")
                    # 详情页需要按 parent_id 动态查询父段并支持多条展示，这里不在检索结果中拼接父段内容

                    # 获取文档名称
                    doc_name = "未知文档"
                    
                    # 1. 尝试从 metadata 中获取
                    meta_dict = entity.get("metadata", {})
                    if not meta_dict and isinstance(entity.get("metadata"), str):
                         try:
                             meta_dict = json.loads(entity.get("metadata"))
                         except:
                             meta_dict = {}

                    # 优先取 doc_name / file_name / title
                    doc_name = meta_dict.get("doc_name") or meta_dict.get("file_name") or meta_dict.get("title") or entity.get("file_name") or entity.get("title") or entity.get("source")
                    
                    # 2. 如果只有 document_id，尝试去数据库查
                    document_id = entity.get("document_id") or meta_dict.get("document_id")
                    if (not doc_name or doc_name == "未知文档") and document_id:
                        try:
                            # 尝试从 DocumentMain 表查询
                            from app.sample.models.base_info import DocumentMain
                            # document_id 可能是 int 或 str
                            stmt = None
                            if str(document_id).isdigit():
                                stmt = select(DocumentMain.title).where(DocumentMain.id == int(document_id))
                            else:
                                # 如果是字符串 ID，可能需要转 int，或者就是字符串
                                # 假设 DocumentMain.id 是 int
                                try:
                                    stmt = select(DocumentMain.title).where(DocumentMain.id == int(document_id))
                                except:
                                    pass
                            
                            if stmt is not None:
                                doc_res = await db.execute(stmt)
                                doc_title = doc_res.scalar()
                                if doc_title:
                                    doc_name = doc_title
                        except Exception as e:
                            # logging.error(f"Failed to fetch doc name for {document_id}: {e}")
                            pass

                    if not doc_name:
                        doc_name = "未知文档"
                    
                    meta_info = []
                    # 优先展示 metadata 字段内容，如果不为空则仅展示它
                    # 但为了兼容旧数据，如果 metadata 为空，还是展示其他字段
                    metadata_content = entity.get("metadata")
                    if metadata_content:
                        # 尝试格式化 JSON
                        if isinstance(metadata_content, dict):
                            meta_str = json.dumps(metadata_content, ensure_ascii=False)
                        elif isinstance(metadata_content, str):
                            meta_str = metadata_content
                        else:
                            meta_str = str(metadata_content)
                    else:
                        # Fallback: 展示扁平字段 (不推荐)
                        for k, v in entity.items():
                            if k not in [anns_field, "text", "content", "page_content", "id", "pk", "metadata"]:
                                meta_info.append(f"{k}: {v}")
                        meta_str = "; ".join(meta_info[:3])
                    
                    # 根据 collection 的 metric 动态计算相似度分数（先从原始向量距离/相似度换算到 0-100）
                    display_metric = requested_metric if requested_metric else metric_type
                    
                    similarity_pct = None
                    try:
                        raw_score = float(hit.score)
                    except Exception:
                        raw_score = None

                    if raw_score is not None:
                        if "L2" in metric_type or "EUCLIDEAN" in metric_type:
                            distance = raw_score
                            similarity_pct = round((1.0 / (1.0 + distance)) * 100.0, 2)
                        elif "COSINE" in metric_type:
                            cosine_score = raw_score
                            similarity_pct = round(max(min((cosine_score + 1.0) / 2.0, 1.0), 0.0) * 100.0, 2)
                        elif "IP" in metric_type or "INNER" in metric_type:
                             similarity_pct = round(raw_score * 100.0, 2)
                        else:
                            if "BM25" in metric_type:
                                similarity_pct = round(raw_score, 2)
                            else:
                                similarity_pct = round(raw_score * 100.0, 2)
                    else:
                        similarity_pct = 0.0

                    # 结合关键词做一次简单的“相关性校正”：若片段完全不包含检索词，则降低相似度权重
                    query_text = (payload.query or "").strip()
                    if query_text:
                        plain_content = (content or "") + " " + (doc_name or "") + " " + (meta_str or "")
                        if query_text not in plain_content:
                            # 纯向量相似但没有任何文本命中，认为是“语义可能相关，但与关键词弱相关”，适当降权
                            similarity_pct = round(similarity_pct * 0.4, 2)

                    formatted_results.append(KBSearchResultItem(
                        id=str(hit.id),
                        kb_name=original_kb_id,
                        doc_name=doc_name,
                        content=content,
                        meta_info=meta_str,
                        document_id=str(document_id) if document_id is not None else None,
                        parent_id=str(parent_id) if parent_id is not None else None,
                        metadata=meta_dict if isinstance(meta_dict, dict) else None,
                        score=similarity_pct
                    ))

            formatted_results.sort(key=lambda r: (r.score or 0), reverse=True)
            
            # [Fix] 动态计算 total 用于分页
            # 如果当前页结果不满 limit，说明是最后一页
            current_count = len(formatted_results)
            if current_count < limit:
                 final_total = offset + current_count
            else:
                 # 否则，使用 collection_count (上限 1000)
                 # 如果 collection_count 获取失败(0)，则至少允许翻页
                 base_total = collection_count if collection_count > 0 else 1000
                 final_total = min(base_total, 1000)
                 
                 # 确保 final_total 至少能覆盖当前页+1 (如果有满页结果)
                 # 这样用户能看到"下一页"按钮
                 if final_total <= offset + current_count:
                     final_total = offset + current_count + 10 # 预留一页

            return KBSearchResponse(results=formatted_results, total=final_total)
            
        except Exception as e:
            logger.exception("Search error")
            return KBSearchResponse(results=[], total=0)

    # ... (Keep existing CRUD methods below) ...
    
    async def get_list(
        self, 
        db: AsyncSession,
        page: int = 1,
        page_size: int = 10,
        keyword: Optional[str] = None,
        status: Optional[str] = None
    ) -> Tuple[List[SearchEngine], PaginationSchema]:
        """获取检索引擎列表"""
        
        query = select(SearchEngine).where(SearchEngine.is_deleted == 0)
        
        if keyword:
            query = query.where(or_(
                SearchEngine.name.like(f"%{keyword}%"),
                SearchEngine.description.like(f"%{keyword}%")
            ))
        
        if status:
            query = query.where(SearchEngine.status == status)

        # 计算总数
        count_query = select(func.count()).select_from(query.subquery())
        total = await db.scalar(count_query) or 0

        # 分页查询
        query = query.order_by(SearchEngine.created_time.desc()).offset((page - 1) * page_size).limit(page_size)
        result = await db.execute(query)
        items = result.scalars().all()

        total_pages = ceil(total / page_size) if page_size else 0
        
        meta = PaginationSchema(
            page=page,
            page_size=page_size,
            total=total,
            total_pages=total_pages,
        )
        
        return items, meta

    async def create(self, db: AsyncSession, payload: SearchEngineCreate) -> SearchEngine:
        """创建检索引擎"""
        # 1. 检查名称是否已存在
        exists = await db.execute(select(SearchEngine).where(
            SearchEngine.name == payload.name,
            SearchEngine.is_deleted == 0
        ))
        if exists.scalars().first():
            raise ValueError("检索引擎名称已存在")

        try:
            now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            new_engine = SearchEngine(
                id=str(uuid.uuid4()),
                name=payload.name,
                engine_type=payload.engine_type,
                base_url=payload.base_url,
                api_key=payload.api_key,
                description=payload.description,
                status=payload.status or "normal",
                created_time=datetime.now(),
                updated_time=datetime.now()
            )
            db.add(new_engine)
            await db.commit()
            await db.refresh(new_engine)

            return new_engine
        except Exception as e:
            await db.rollback()
            raise e

    async def update(self, db: AsyncSession, id: str, payload: SearchEngineUpdate) -> SearchEngine:
        """更新检索引擎信息"""
        result = await db.execute(select(SearchEngine).where(SearchEngine.id == id, SearchEngine.is_deleted == 0))
        engine = result.scalars().first()
        
        if not engine:
            raise ValueError("检索引擎不存在")

        try:
            if payload.name is not None:
                engine.name = payload.name
            if payload.engine_type is not None:
                engine.engine_type = payload.engine_type
            if payload.base_url is not None:
                engine.base_url = payload.base_url
            if payload.api_key is not None:
                engine.api_key = payload.api_key
            if payload.description is not None:
                engine.description = payload.description
            if payload.status is not None:
                engine.status = payload.status
            
            engine.updated_time = datetime.now()
            await db.commit()
            await db.refresh(engine)
            
            return engine
        except Exception as e:
            await db.rollback()
            raise e

    async def update_status(self, db: AsyncSession, id: str, status: str) -> SearchEngine:
        """更新检索引擎状态"""
        result = await db.execute(select(SearchEngine).where(SearchEngine.id == id, SearchEngine.is_deleted == 0))
        engine = result.scalars().first()
        
        if not engine:
            raise ValueError("检索引擎不存在")
        
        try:
            engine.status = status
            engine.updated_time = datetime.now()
            await db.commit()
            await db.refresh(engine)
            return engine
        except Exception as e:
            await db.rollback()
            raise e

    async def delete(self, db: AsyncSession, id: str) -> None:
        """删除检索引擎"""
        result = await db.execute(select(SearchEngine).where(SearchEngine.id == id))
        engine = result.scalars().first()
        
        if not engine:
            raise ValueError("检索引擎不存在")

        try:
            # 软删除
            engine.is_deleted = 1
            engine.updated_time = datetime.now()
            await db.commit()
        except Exception as e:
            await db.rollback()
            raise e

search_engine_service = SearchEngineService()
