"""
===========================================
设备到系统反向检索模块 - 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://192.168.0.46:57687"),
    "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)}")
        
        best = 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 _do_find_system_by_device(
    session,
    device_name: str,
    uri: str,
    user: str,
    password: str,
    min_hops: int,
    max_hops: int,
    top_k: int,
    start_time: float
) -> Dict[str, Any]:
    """
    执行设备到系统的检索逻辑（内部函数）
    """
    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
        }
    }


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"}
    
    external_driver = driver is not None
    
    try:
        if external_driver:
            async with driver.session() as session:
                return await _do_find_system_by_device(
                    session, device_name, uri, user, password, min_hops, max_hops, top_k, start_time
                )
        else:
            async with AsyncGraphDatabase.driver(uri, auth=(user, password)) as driver:
                async with driver.session() as session:
                    return await _do_find_system_by_device(
                        session, device_name, uri, user, password, min_hops, max_hops, top_k, start_time
                    )
    
    except Exception as e:
        logger.error(f"[设备到系统检索] 检索失败: {e}", exc_info=True)
        return {"success": False, "error": f"检索失败: {str(e)}"}

