"""
===========================================
设备到系统反向检索模块 - device_to_system
===========================================
功能：根据设备名称，反向检索对应的系统节点
     使用混合检索（向量+全文）匹配设备节点
     然后反向遍历2-8级找到对应的系统节点
===========================================
"""

import json
import os
import logging
import time
from typing import Any, Dict, List, Optional

from dotenv import load_dotenv
from neo4j import AsyncGraphDatabase, GraphDatabase

load_dotenv()

logger = logging.getLogger(__name__)

NEO4J_CONFIG = {
    "uri": os.getenv("NEO4J_URI", "bolt://neo4j:7687"),
    "user": os.getenv("NEO4J_USER", "neo4j"),
    "password": os.getenv("NEO4J_PASSWORD", "zdht123@"),
}


class EmbeddingWrapper:
    """嵌入模型包装器"""
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """批量获取嵌入向量（一次性处理所有文本）"""
        from modelsAPI.model_api import OpenaiAPI
        
        if not texts:
            return []
        
        try:
            result = OpenaiAPI.batch_embeddings(texts)
            if result and isinstance(result, list):
                embedding_list = [item.get('embedding') if isinstance(item, dict) else item for item in result]
                if embedding_list and len(embedding_list) == len(texts):
                    return embedding_list
                else:
                    logger.error(f"EmbeddingWrapper: 批量向量化返回数量不匹配，期望: {len(texts)}, 实际: {len(embedding_list) if embedding_list else 0}")
                    return [None] * len(texts)
            else:
                logger.error(f"EmbeddingWrapper: 批量向量化返回格式异常")
                return [None] * len(texts)
        except Exception as e:
            logger.error(f"EmbeddingWrapper: 批量向量化失败: {e}", exc_info=True)
            return [None] * len(texts)


async def _match_device_with_retrieval(
    session,
    device_name: str,
    sync_driver,
    top_k: int = 10
) -> Optional[Dict[str, Any]]:
    """
    使用 NodeRetrieval 混合检索匹配设备节点
    
    Args:
        session: Neo4j 异步会话
        device_name: 设备名称
        sync_driver: Neo4j 同步 driver（用于 NodeRetrieval）
        top_k: 检索返回的节点数量上限
        
    Returns:
        匹配到的设备节点信息，包含 id, name, labels, props, score
    """
    from graph_search.node_retrieval import NodeRetrieval
    
    try:
        embeddings = EmbeddingWrapper()
        node_retrieval = NodeRetrieval(sync_driver, embeddings)
        
        route_res = [{"label": "设备", "entity": device_name}]
        retrieved_nodes = await node_retrieval.retrieve_nodes(route_res, top_k=top_k)
        
        if not retrieved_nodes or "设备" not in retrieved_nodes:
            logger.warning(f"NodeRetrieval 未检索到设备节点: {device_name}")
            return None
        
        candidates = retrieved_nodes["设备"]
        if not candidates:
            return None
        
        logger.info(f"NodeRetrieval 检索到 {len(candidates)} 个设备节点候选")
        
        for i, node in enumerate(candidates):
            node_data = node.get("node", {})
            candidate_name = node_data.get("名称") or node_data.get("name") or node.get("name", "")
            logger.info(f"  候选 {i+1}: 名称={candidate_name}, score={node.get('score', 0.0)}")
        
        # 优先选择名称与查询精确匹配的候选（忽略大小写/空格），避免分数并列时选错
        device_key_norm = str(device_name).strip().lower()
        exact_match = None
        for node in candidates:
            node_data = node.get("node", {})
            candidate_name = node_data.get("名称") or node_data.get("name") or node.get("name", "")
            if candidate_name and str(candidate_name).strip().lower() == device_key_norm:
                exact_match = node
                logger.info(f"  命中精确名称匹配: {candidate_name}")
                break
        
        best = exact_match if exact_match is not None else candidates[0]
        node_data = best.get("node", {})
        best_name = node_data.get("名称") or node_data.get("name") or best.get("name", "")
        
        query = """
        MATCH (n:设备)
        WHERE n.名称 = $name OR n.name = $name
        RETURN elementId(n) AS id, n.名称 AS name, labels(n) AS labels, properties(n) AS props
        LIMIT 1
        """
        result = await session.run(query, name=best_name)
        record = await result.single()
        
        if record:
            return {
                "id": record.get("id"),
                "name": record.get("name"),
                "labels": record.get("labels"),
                "props": record.get("props", {}),
                "score": best.get("score", 0.0)
            }
        
        return None
    except Exception as e:
        logger.error(f"_match_device_with_retrieval 失败: {e}", exc_info=True)
        return None


async def find_system_by_device(
    device_name: str,
    min_hops: int = 2,
    max_hops: int = 8,
    top_k: int = 10,
    driver=None,
) -> Dict[str, Any]:
    """
    根据设备名称反向检索对应的系统节点
    
    Args:
        device_name: 设备名称
        min_hops: 最小跳数（默认2级）
        max_hops: 最大跳数（默认8级）
        top_k: 混合检索返回的候选节点数量
        driver: Neo4j driver 实例（可选，如果不提供则创建新连接）
        
    Returns:
        Dict: 检索结果，包含 success, systems, path 等字段
    """
    start_time = time.time()
    
    if not device_name or not str(device_name).strip():
        return {"success": False, "error": "设备名称不能为空"}
    
    uri = NEO4J_CONFIG.get("uri")
    user = NEO4J_CONFIG.get("user")
    password = NEO4J_CONFIG.get("password")
    if not uri or not user or not password:
        return {"success": False, "error": "Neo4j 配置不完整，请检查 NEO4J_CONFIG"}
    
    try:
        # 总是创建新的异步 driver，忽略外部传入的同步 driver
        driver = AsyncGraphDatabase.driver(uri, auth=(user, password))
        
        async with driver.session() as session:
            device_key = str(device_name).strip()
            
            logger.info(f"[设备到系统检索] 开始检索设备: {device_key}")
            
            sync_driver = GraphDatabase.driver(uri, auth=(user, password))
            try:
                device_node = await _match_device_with_retrieval(session, device_key, sync_driver, top_k=top_k)
            finally:
                sync_driver.close()
            
            if not device_node:
                return {"success": False, "error": "未找到匹配的设备节点"}
            
            device_id = device_node.get("id")
            device_name_found = device_node.get("name")
            logger.info(f"[设备到系统检索] 匹配到设备节点: {device_name_found}, id={device_id}, score={device_node.get('score', 0.0)}")
            
            system_query = """
            MATCH (d:设备)
            WHERE elementId(d) = $device_id
            MATCH (d)-[:包含*%d..%d]-(s:系统)
            RETURN DISTINCT elementId(s) AS id, s.名称 AS name, labels(s) AS labels, properties(s) AS props
            ORDER BY s.名称
            """ % (min_hops, max_hops)
            
            result = await session.run(system_query, device_id=device_id)
            system_records = await result.data()
            
            if not system_records:
                logger.warning(f"[设备到系统检索] 未找到设备对应的系统节点，尝试扩大搜索范围...")
                extended_query = """
                MATCH (d:设备)
                WHERE elementId(d) = $device_id
                MATCH (d)-[:包含*1..10]-(s:系统)
                RETURN DISTINCT elementId(s) AS id, s.名称 AS name, labels(s) AS labels, properties(s) AS props
                ORDER BY s.名称
                """
                result = await session.run(extended_query, device_id=device_id)
                system_records = await result.data()
                
                if not system_records:
                    return {
                        "success": False,
                        "error": "未找到设备对应的系统节点",
                        "device": {
                            "id": device_id,
                            "name": device_name_found,
                            "labels": device_node.get("labels", []),
                            "properties": device_node.get("props", {})
                        }
                    }
            
            exclude_keys = {
                "last_updated",
                "created_at",
                "knowledge_source",
                "fulltext",
                "embedding",
            }
            
            systems = []
            for record in system_records:
                props = record.get("props", {})
                props_filtered = {k: v for k, v in props.items() if k not in exclude_keys}
                
                systems.append({
                    "id": record.get("id"),
                    "name": record.get("name"),
                    "labels": record.get("labels", []),
                    "properties": props_filtered,
                })
            
            logger.info(f"[设备到系统检索] 找到 {len(systems)} 个系统节点")
            
            path_query = """
            MATCH (d:设备)
            WHERE elementId(d) = $device_id
            MATCH path = (d)-[:包含*%d..%d]-(s:系统)
            RETURN path
            LIMIT 5
            """ % (min_hops, max_hops)
            
            result = await session.run(path_query, device_id=device_id)
            path_records = await result.data()
            
            nodes_g: List[Dict[str, Any]] = []
            links_g: List[Dict[str, Any]] = []
            node_id_set = set()
            
            for path_record in path_records:
                path = path_record.get("path")
                if not path:
                    continue
                
                if hasattr(path, 'nodes'):
                    for node in path.nodes:
                        if node is None:
                            continue
                        nid = node.element_id if hasattr(node, 'element_id') else node.id
                        if nid in node_id_set:
                            continue
                        node_id_set.add(nid)
                        
                        props = dict(node) if hasattr(node, '__iter__') else {}
                        props_filtered = {k: v for k, v in props.items() if k not in exclude_keys}
                        
                        name = props.get("名称") or props.get("name") or str(nid)
                        labels = list(node.labels) if hasattr(node, 'labels') else []
                        
                        nodes_g.append({
                            "id": nid,
                            "name": name,
                            "labels": labels,
                            "properties": props_filtered,
                        })
                
                if hasattr(path, 'relationships'):
                    for rel in path.relationships:
                        if rel is None:
                            continue
                        src = rel.start_node.element_id if hasattr(rel.start_node, 'element_id') else rel.start_node.id
                        tgt = rel.end_node.element_id if hasattr(rel.end_node, 'element_id') else rel.end_node.id
                        
                        if src not in node_id_set or tgt not in node_id_set:
                            continue
                        
                        links_g.append({
                            "id": rel.element_id if hasattr(rel, 'element_id') else rel.id,
                            "label": rel.type if hasattr(rel, 'type') else str(rel.type),
                            "source": src,
                            "target": tgt,
                        })
            
            elapsed_ms = (time.time() - start_time) * 1000
            logger.info(f"[时间统计] 设备到系统检索完成，耗时: {elapsed_ms:.2f}ms")
            
            return {
                "success": True,
                "device": {
                    "id": device_id,
                    "name": device_name_found,
                    "labels": device_node.get("labels", []),
                    "properties": {k: v for k, v in device_node.get("props", {}).items() if k not in exclude_keys},
                    "match_score": device_node.get("score", 0.0)
                },
                "systems": systems,
                "graph": {
                    "nodes": nodes_g,
                    "links": links_g
                },
                "meta": {
                    "elapsed_ms": round(elapsed_ms, 2),
                    "system_count": len(systems),
                    "min_hops": min_hops,
                    "max_hops": max_hops
                }
            }
    
    except Exception as e:
        logger.error(f"[设备到系统检索] 检索失败: {e}", exc_info=True)
        return {"success": False, "error": f"检索失败: {str(e)}"}
    finally:
        if driver is not None:
            await driver.close()