"""
===========================================
操作图谱检索模块 - operation_graph_search
===========================================
功能：针对"操作"的固定图谱检索接口
     检索路径：舰艇 -> 系统 -> 子系统 -> 设备 -> 操作使用 -> 操作程序 -> 操作项目
===========================================
"""

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_node_with_retrieval(
    session,
    label: str,
    node_name: str,
    top_k: int = 10
) -> Optional[Dict[str, Any]]:
    """
    使用 NodeRetrieval 匹配节点
    """
    from graph_search.node_retrieval import NodeRetrieval
    
    uri = NEO4J_CONFIG.get("uri")
    user = NEO4J_CONFIG.get("user")
    password = NEO4J_CONFIG.get("password")
    
    sync_driver = GraphDatabase.driver(uri, auth=(user, password))
    
    try:
        embeddings = EmbeddingWrapper()
        node_retrieval = NodeRetrieval(sync_driver, embeddings)
        
        route_res = [{"label": label, "entity": node_name}]
        retrieved_nodes = await node_retrieval.retrieve_nodes(route_res, top_k=top_k)
        
        if not retrieved_nodes or label not in retrieved_nodes:
            logger.warning(f"NodeRetrieval 未检索到 {label} 节点: {node_name}")
            return None
        
        candidates = retrieved_nodes[label]
        if not candidates:
            return None
        
        logger.info(f"NodeRetrieval 检索到 {len(candidates)} 个 {label} 节点候选")
        
        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 = f"""
        MATCH (n:{label})
        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
    finally:
        sync_driver.close()


async def operation_graph_search(
    ship_number: str,
    device_name: str,
    operation_item: str,
) -> Dict[str, Any]:
    """
    操作图谱检索接口
    
    Args:
        ship_number: 舰艇舷号
        device_name: 设备名称
        operation_item: 操作项目名称
        
    Returns:
        Dict: 检索结果，包含 success, xxxx.graph 等字段
    """
    start_time = time.time()
    
    if not ship_number or not str(ship_number).strip():
        return {"success": False, "error": "舷号不能为空"}
    if not device_name or not str(device_name).strip():
        return {"success": False, "error": "设备名称不能为空"}
    if not operation_item or not str(operation_item).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:
        async with AsyncGraphDatabase.driver(uri, auth=(user, password)) as driver:
            async with driver.session() as session:
                ship_key = str(ship_number).strip()
                device_key = str(device_name).strip()
                operation_key = str(operation_item).strip()

                # ========== 1) 匹配舰艇节点 ==========
                ship_query = """
                MATCH (s:舰艇)
                WHERE s.舷号 = $ship_key OR s.舷号 CONTAINS $ship_key OR $ship_key CONTAINS s.舷号
                RETURN elementId(s) AS id, s.名称 AS name, s.舷号 AS ship_number, labels(s) AS labels, properties(s) AS props
                LIMIT 1
                """
                result = await session.run(ship_query, ship_key=ship_key)
                ship_record = await result.single()
                
                if not ship_record:
                    return {"success": False, "error": "在图谱中未找到对应舷号的舰艇节点"}
                
                ship_id = ship_record.get("id")
                ship_name = ship_record.get("name") or ship_record.get("ship_number")
                logger.info(f"匹配到舰艇节点: {ship_name}, id={ship_id}")

                # ========== 2) 匹配设备节点 ==========
                device_query = """
                MATCH (ship:舰艇)
                WHERE elementId(ship) = $ship_id
                MATCH (ship)-[:包含*0..8]->(d:设备)
                WHERE d.名称 CONTAINS $device_key OR $device_key CONTAINS d.名称
                   OR d.name CONTAINS $device_key OR $device_key CONTAINS d.name
                RETURN elementId(d) AS id, d.名称 AS name, labels(d) AS labels, properties(d) AS props
                LIMIT 1
                """
                result = await session.run(device_query, ship_id=ship_id, device_key=device_key)
                device_record = await result.single()
                
                if not device_record:
                    device_node = await _match_node_with_retrieval(session, "设备", device_key)
                    if device_node:
                        check_query = """
                        MATCH (ship:舰艇)
                        WHERE elementId(ship) = $ship_id
                        MATCH (ship)-[:包含*0..8]->(d:设备)
                        WHERE elementId(d) = $device_id
                        RETURN true AS in_range
                        LIMIT 1
                        """
                        check_result = await session.run(check_query, ship_id=ship_id, device_id=device_node["id"])
                        check_record = await check_result.single()
                        if check_record:
                            device_record = device_node
                    
                    if not device_record:
                        return {"success": False, "error": "在该舰艇范围内未找到匹配的设备节点"}
                
                device_id = device_record.get("id")
                device_name_found = device_record.get("name")
                logger.info(f"匹配到设备节点: {device_name_found}, id={device_id}")

                # ========== 3) 匹配操作项目节点 ==========
                op_item_query = """
                MATCH (d:设备)
                WHERE elementId(d) = $device_id
                MATCH (d)-[:包含*0..8]-(real_d:设备)
                MATCH (real_d)-[:设备支持操作操作使用]-(op_use:操作使用)
                      -[:使用]-(op_proc:操作程序)
                      -[:操作时使用]-(op_item:操作项目)
                WHERE op_item.名称 CONTAINS $op_key OR $op_key CONTAINS op_item.名称
                RETURN elementId(op_item) AS id, op_item.名称 AS name, labels(op_item) AS labels, properties(op_item) AS props
                LIMIT 1
                """
                result = await session.run(op_item_query, device_id=device_id, op_key=operation_key)
                op_item_record = await result.single()
                
                if not op_item_record:
                    logger.info(f"Cypher查询未找到操作项目，尝试 NodeRetrieval...")
                    op_item_node = await _match_node_with_retrieval(session, "操作项目", operation_key)
                    if op_item_node:
                        logger.info(f"NodeRetrieval 找到操作项目节点: id={op_item_node.get('id')}, name={op_item_node.get('name')}")
                        check_query = """
                        MATCH (d:设备)
                        WHERE elementId(d) = $device_id
                        MATCH (d)-[:包含*0..8]-(real_d:设备)
                        MATCH (real_d)-[:支持操作]-(op_use:操作使用)
                              -[:使用]-(op_proc:操作程序)
                              -[:操作时使用]-(op_item:操作项目)
                        WHERE elementId(op_item) = $op_item_id
                        RETURN true AS in_range
                        LIMIT 1
                        """
                        check_result = await session.run(check_query, device_id=device_id, op_item_id=op_item_node["id"])
                        check_record = await check_result.single()
                        if check_record:
                            logger.info(f"操作项目节点在设备路径下验证通过")
                            op_item_record = op_item_node
                        else:
                            logger.warning(f"操作项目节点不在设备路径下，尝试直接使用...")
                            op_item_record = op_item_node
                    
                    if not op_item_record:
                        return {"success": False, "error": "未找到匹配的操作项目节点"}
                
                op_item_id = op_item_record.get("id")
                op_item_name = op_item_record.get("name")
                logger.info(f"匹配到操作项目节点: {op_item_name}, id={op_item_id}")

                # ========== 4) 查找经过这三个节点的链路 ==========
                exclude_keys = {
                    "last_updated",
                    "created_at",
                    "knowledge_source",
                    "fulltext",
                    "embedding",
                }

                path_query = """
                MATCH (ship:舰艇)
                WHERE elementId(ship) = $ship_id
                MATCH (device:设备)
                WHERE elementId(device) = $device_id
                MATCH (op_item:操作项目)
                WHERE elementId(op_item) = $op_item_id
                
                MATCH path1 = (ship)-[:包含*0..8]->(device)
                MATCH path2 = (device)-[:包含*0..8]-(real_d:设备)
                      -[:支持操作]->(op_use:操作使用)
                      -[:使用]->(op_proc:操作程序)
                      -[:操作时使用]->(op_item)
                
                WITH path1, path2
                WITH nodes(path1) + nodes(path2) AS all_nodes,
                     relationships(path1) + relationships(path2) AS all_rels
                RETURN all_nodes, all_rels
                LIMIT 1
                """
                result = await session.run(path_query, ship_id=ship_id, device_id=device_id, op_item_id=op_item_id)
                path_record = await result.single()
                
                if not path_record:
                    return {"success": False, "error": "未找到连接这三个节点的路径"}

                # ========== 5) 格式化输出 ==========
                all_nodes = path_record.get("all_nodes", [])
                all_rels = path_record.get("all_rels", [])
                
                nodes_g: List[Dict[str, Any]] = []
                node_id_set = set()
                
                for node in all_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,
                    })
                
                links_g: List[Dict[str, Any]] = []
                for rel in all_rels:
                    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,
                    })

                # ========== 6) 生成 results 文本 ==========
                op_item_props = op_item_record.get("props", {})
                op_info = {k: v for k, v in op_item_props.items() if k not in exclude_keys and v}
                results_text = json.dumps(op_info, ensure_ascii=False) if op_info else ""

                elapsed_ms = (time.time() - start_time) * 1000
                logger.info(f"[时间统计] 操作图谱检索完成，耗时: {elapsed_ms:.2f}ms")

                graph_item = {"nodes": nodes_g, "links": links_g, "results": results_text}
                return {"success": True, "xxxx.graph": [graph_item]}

    except Exception as e:
        logger.error(f"操作图谱检索失败: {e}", exc_info=True)
        return {"success": False, "error": f"Neo4j 查询失败: {str(e)}"}