"""
===========================================
节点检索模块 - NodeRetrieval
===========================================
功能：根据路由结果，使用混合检索获取候选入口节点
     使用HybridRetriever混合检索（向量+全文）
===========================================
"""

import re
import asyncio
import logging
import time
from typing import List, Dict, Any, Optional

logger = logging.getLogger(__name__)

# ========== ========== 导入依赖 ========== ==========
try:
    from neo4j_graphrag.retrievers import HybridRetriever
    import jieba
    # 过滤 jieba 缓存写入权限错误的日志（不影响功能，只是无法写入缓存）
    jieba_logger = logging.getLogger('jieba')
    jieba_logger.setLevel(logging.CRITICAL)  # 只显示 CRITICAL 级别，过滤所有 ERROR/WARNING/INFO/DEBUG
    NODE_RETRIEVAL_AVAILABLE = True
except ImportError:
    NODE_RETRIEVAL_AVAILABLE = False


class NodeRetrieval:
    """
    节点检索类
    
    功能：根据路由结果，检索候选入口节点
    """
    
    def __init__(self, driver, embeddings):
        """
        初始化节点检索器
        
        Args:
            driver: Neo4j driver实例
            embeddings: 嵌入模型实例（需要有embed_documents方法）
        """
        self.driver = driver
        self.embeddings = embeddings
    
    async def retrieve_nodes(
        self,
        route_res: List[Dict[str, str]],  # [{"label": "Drug", "entity": "xxx"}]
        top_k: int = 10
    ) -> Dict[str, List[Dict[str, Any]]]:
        """
        节点检索：根据标签和实体，检索入口节点
        
        功能：
        使用混合检索（向量+全文）进行检索
        
        Args:
            route_res: 路由结果，包含标签和实体信息
            top_k: 检索返回的节点数量上限
            
        Returns:
            Dict[str, List[Dict]]: 检索到的节点结果，以标签为键
        """
        if not NODE_RETRIEVAL_AVAILABLE:
            logger.error("节点检索依赖未安装")
            return {}
        
        pairs = []  # 用于存储需要检索的标签-实体对
        retrieved_nodes = {}  # 用于存储检索到的节点结果，以标签为键
        
        for item in route_res:
            label = item.get("label", "")
            entity = item.get("entity", "")
            
            if not entity:  # 如果实体为空则跳过
                continue
            
            # 将标签和实体作为一个元组添加到pairs列表中
            pairs.append((label, entity))
        
        if not pairs:  # 如果没有需要检索的标签-实体对，直接返回
            return retrieved_nodes
        
        # 将标签-实体对分离成两个独立的列表
        labels, entities = zip(*pairs)
        labels, entities = list(labels), list(entities)
        
        # ========== 对每个实体进行中文分词处理，生成全文检索查询文本 ==========
        query_texts = []
        for entity in entities:
            try:
                words = jieba.lcut(entity)
                filtered_words = [
                    word.strip()
                    for word in words
                    if word.strip() and re.fullmatch(r"[a-zA-Z0-9\u4e00-\u9fa5]+", word.strip())
                ]
                # 如果过滤后没有词，使用原始实体
                if not filtered_words:
                    filtered_words = [entity] if entity.strip() else []
                query_text = " OR ".join(filtered_words) if filtered_words else entity
                query_texts.append(query_text)
            except Exception as e:
                logger.warning(f"实体 '{entity}' 分词失败: {e}，使用原始实体作为查询文本")
                query_texts.append(entity)
        
        # ========== 对实体进行向量化处理 ==========
        try:
            query_vectors = self.embeddings.embed_documents(entities)
        except Exception as e:
            logger.error(f"向量化失败: {e}", exc_info=True)
            query_vectors = [None] * len(entities)
        
        # ========== 为每个标签创建混合检索任务 ==========
        tasks = []
        for label, entity, query_text, query_vector in zip(labels, entities, query_texts, query_vectors):
            try:
                # vector_index_name = label.lower() + "_vector"
                # fulltext_index_name = label.lower() + "_fulltext"
                vector_index_name = "global_searchable_embedding"
                fulltext_index_name = "global_searchable_fulltext_search"
                retriever = HybridRetriever(
                    self.driver,
                    vector_index_name=vector_index_name,
                    fulltext_index_name=fulltext_index_name
                )
                tasks.append(
                    asyncio.to_thread(
                        retriever.get_search_results,
                        query_text,
                        query_vector,
                        top_k,
                        effective_search_ratio=2,
                    )
                )
            except Exception as e:
                logger.error(f"创建HybridRetriever失败 ({label}): {e}", exc_info=True)
                tasks.append(None)
        
        # ========== 并发执行所有检索任务 ==========
        valid_tasks = [task for task in tasks if task is not None]
        if valid_tasks:
            results = await asyncio.gather(*valid_tasks, return_exceptions=True)
        else:
            logger.warning("没有有效的检索任务")
            results = []
        
        # ========== 处理检索结果 ==========
        valid_results = []
        exception_results = []
        
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                label = labels[i] if i < len(labels) else "unknown"
                logger.error(f"检索任务 {i} (标签: {label}) 执行异常: {result}", exc_info=result)
                exception_results.append((i, result))
            else:
                valid_results.append((i, result))
        
        # ========== 按标签分组，准备进行RRF融合 ==========
        # 结构: {label: [(entity, result_index, nodes_list), ...]}
        label_results_map = {}
        
        for i, result in valid_results:
            label = labels[i] if i < len(labels) else "unknown"
            entity = entities[i] if i < len(entities) else "unknown"
            
            if result is None:
                logger.warning(f"标签 '{label}' (实体: '{entity}') 的检索结果为 None")
                continue
            
            if not hasattr(result, 'records'):
                logger.warning(f"标签 '{label}' (实体: '{entity}') 的结果没有 'records' 属性，类型: {type(result)}")
                continue
            
            
            # 提取节点信息（优先使用"名称"属性，兼容"name"）
            nodes_list = [
                {
                    "name": record["node"].get("名称") or record["node"].get("name") or str(record["node"]),
                    "score": record.get("score", 0.0),
                    "node": dict(record["node"])
                }
                for record in result.records
            ]
            
            if nodes_list:
                if label not in label_results_map:
                    label_results_map[label] = []
                label_results_map[label].append((entity, i, nodes_list))
        
        # ========== 对每个标签的多个实体结果进行RRF融合 ==========
        # 策略：只对同一标签下相同实体名的结果进行RRF融合
        rrf_fusion_start = time.time()
        for label, entity_results in label_results_map.items():
            # ========== 按实体名分组 ==========
            entity_name_groups = {}  # {entity_name: [(entity, result_index, nodes_list), ...]}
            
            for entity, result_index, nodes_list in entity_results:
                if entity not in entity_name_groups:
                    entity_name_groups[entity] = []
                entity_name_groups[entity].append((entity, result_index, nodes_list))
            
            # ========== 对每个实体名的结果进行RRF融合，然后合并 ==========
            all_fused_nodes = []
            
            for entity_name, group_results in entity_name_groups.items():
                if len(group_results) == 1:
                    # 只有一个结果，不需要RRF融合，直接使用
                    _, _, nodes_list = group_results[0]
                    all_fused_nodes.extend(nodes_list)
                    logger.debug(f"标签 '{label}' 实体 '{entity_name}' 只有一个结果，跳过RRF融合")
                else:
                    # 多个结果（相同实体名），使用RRF融合
                    fused_nodes = self._rrf_fusion(group_results, top_k=top_k)
                    all_fused_nodes.extend(fused_nodes)
            
            # ========== 对同一标签内的所有节点按score排序 ==========
            all_fused_nodes.sort(key=lambda x: x.get('score', 0), reverse=True)
            
            # ========== 限制每个标签的数量 ==========
            retrieved_nodes[label] = all_fused_nodes[:top_k]
        
        # 统计去重后的结果
        dedup_stats = {k: len(v) for k, v in retrieved_nodes.items()}
        logger.info(f"检索结果统计（RRF融合后）: {list(dedup_stats.items())}")
        return retrieved_nodes
    
    def _rrf_fusion(
        self, 
        entity_results: List[tuple], 
        top_k: int = 10,
        k: int = 60
    ) -> List[Dict[str, Any]]:
        """
        使用RRF (Reciprocal Rank Fusion) 融合多个实体的检索结果
        
        RRF公式: RRF(d) = Σ 1/(k + rank_i(d))
        其中:
        - d 是文档/节点
        - rank_i(d) 是文档在第i个排序结果中的排名（从1开始）
        - k 是常数，通常为60
        
        Args:
            entity_results: [(entity, result_index, nodes_list), ...]
            top_k: 返回的节点数量上限
            k: RRF常数，默认60
            
        Returns:
            List[Dict]: 融合后的节点列表，按RRF分数降序排序
        """
        if not entity_results:
            return []
        
        # ========== 1. 为每个节点计算RRF分数 ==========
        node_rrf_scores = {}  # {node_name: {"rrf_score": float, "node_data": dict, "ranks": [rank1, rank2, ...]}}
        
        for entity, result_index, nodes_list in entity_results:
            for rank, node in enumerate(nodes_list, start=1):  # rank从1开始
                node_name = node.get("name")
                if not node_name:
                    continue
                
                if node_name not in node_rrf_scores:
                    node_rrf_scores[node_name] = {
                        "rrf_score": 0.0,
                        "node_data": node,  # 保留第一个出现的节点数据
                        "ranks": []
                    }
                
                # 累加RRF分数: 1/(k + rank)
                rrf_contribution = 1.0 / (k + rank)
                node_rrf_scores[node_name]["rrf_score"] += rrf_contribution
                node_rrf_scores[node_name]["ranks"].append(rank)
        
        # ========== 2. 转换为列表并按RRF分数排序 ==========
        fused_nodes = []
        for node_name, data in node_rrf_scores.items():
            node_data = data["node_data"].copy()
            node_data["score"] = data["rrf_score"]  # 使用RRF分数替换原始score
            node_data["_rrf_ranks"] = data["ranks"]  # 保留排名信息用于调试
            fused_nodes.append(node_data)
        
        # 按RRF分数降序排序
        fused_nodes.sort(key=lambda x: x["score"], reverse=True)
        
        # ========== 3. 限制返回数量 ==========
        fused_nodes = fused_nodes[:top_k]
        
        logger.debug(
            f"RRF融合: {len(entity_results)} 个实体, "
            f"融合前总节点数: {sum(len(nodes) for _, _, nodes in entity_results)}, "
            f"融合后: {len(fused_nodes)} 个节点"
        )
        
        return fused_nodes

