"""
===========================================
临时图谱检索方案 - QuickGraphSearch
===========================================
功能：直接用用户问题embedding混合检索找最匹配节点
     获取该节点的一跳邻居，用rerank选择最合适的
     只输出入口节点和选中节点
===========================================
"""

import json
import os
import logging
import time
import re
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 _search_all_nodes_with_hybrid(
    session,
    query: str,
    top_k: int = 20
) -> List[Dict[str, Any]]:
    """
    直接对所有节点进行混合检索（embedding+全文），找最匹配的节点
    使用 NodeRetrieval 替代 HybridRetriever，避免兼容性问题

    Args:
        session: Neo4j 异步会话
        query: 用户查询
        top_k: 返回节点数量

    Returns:
        List[Dict]: 匹配到的节点列表
    """
    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:
        # 获取所有节点标签
        all_labels = ["设备","系统","子系统","故障现象"]
        logger.info(f"[混合检索] 所有节点标签: {all_labels}")

        # 使用 NodeRetrieval 进行检索
        embeddings = EmbeddingWrapper()
        node_retrieval = NodeRetrieval(sync_driver, embeddings)
        
        # 为每个标签构造路由结果
        route_res = [{"label": label, "entity": query} for label in all_labels]
        
        # 执行检索
        retrieved_nodes = await node_retrieval.retrieve_nodes(route_res, top_k=top_k)
        
        # 汇总结果
        all_results = []
        for label, nodes in retrieved_nodes.items():
            for node in nodes:
                all_results.append({
                    "name": node.get("name"),
                    "score": node.get("score", 0.0),
                    "node": node.get("node", {}),
                    "label": label
                })

        # 按score排序，取top_k
        all_results.sort(key=lambda x: x.get("score", 0), reverse=True)
        all_results = all_results[:top_k]

        logger.info(f"[混合检索] 找到 {len(all_results)} 个候选节点")

        # 获取完整节点信息
        final_nodes = []
        for item in all_results:
            node_data = item["node"]
            label = item["label"]
            name = item["name"]

            get_node_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(get_node_query, name=name)
            record = await result.single()

            if record:
                final_nodes.append({
                    "id": record.get("id"),
                    "name": record.get("name"),
                    "labels": record.get("labels"),
                    "props": record.get("props", {}),
                    "score": item.get("score", 0.0)
                })

        return final_nodes

    finally:
        sync_driver.close()


async def _rerank_nodes(
    query,
    nodes
):
    """
    使用 rerank 对节点进行重排序（只使用节点名称）

    Args:
        query: 用户查询
        nodes: 节点列表，每个节点需要有 name 字段

    Returns:
        List[Dict]: 重排序后的节点列表
    """
    if not nodes:
        return []

    from modelsAPI.model_api import OpenaiAPI

    # 格式化节点文本（只使用节点名称）
    documents = []
    for node in nodes:
        name = node.get("name", "")
        doc_text = name
        documents.append(doc_text)

    # Rerank（只取top 10，减少token消耗）
    try:
        rerank_results = OpenaiAPI.rerank_query(query, documents, top_n=min(len(documents), 10))

        # 按score排序
        sorted_rerank = sorted(rerank_results, key=lambda x: x.get('scores', 0), reverse=True)

        # 重新组织节点
        reranked_nodes = []
        for rr in sorted_rerank:
            idx = rr.get('index', 0)
            if 0 <= idx < len(nodes):
                node = nodes[idx].copy()
                node['rerank_score'] = rr.get('scores', 0)
                reranked_nodes.append(node)

        return reranked_nodes

    except Exception as e:
        logger.error(f"Rerank失败: {e}", exc_info=True)
        # 失败时返回原顺序
        for node in nodes:
            node['rerank_score'] = 0
        return nodes


async def _get_one_hop_neighbors(
    session,
    node_id: str
) -> List[Dict[str, Any]]:
    """
    获取指定节点的一跳邻居节点

    Args:
        session: Neo4j 异步会话
        node_id: 节点ID

    Returns:
        List[Dict]: 一跳邻居节点列表
    """
    exclude_keys = {
        "last_updated",
        "created_at",
        "knowledge_source",
        "fulltext",
        "embedding",
    }

    query = """
    MATCH (n)
    WHERE elementId(n) = $node_id
    MATCH (n)-[r]-(m)
    RETURN
        elementId(m) AS id,
        m.名称 AS name,
        labels(m) AS labels,
        properties(m) AS props,
        elementId(r) AS rel_id,
        type(r) AS rel_type,
        properties(r) AS rel_props,
        elementId(n) AS source_id,
        n.名称 AS source_name,
        labels(n) AS source_labels
    LIMIT 100
    """

    result = await session.run(query, node_id=node_id)
    records = await result.data()

    neighbors = []
    for record in records:
        props = record.get("props", {})
        props_filtered = {k: v for k, v in props.items() if k not in exclude_keys}

        rel_props = record.get("rel_props", {})
        rel_props_filtered = {k: v for k, v in rel_props.items() if k not in exclude_keys}

        neighbors.append({
            "id": record.get("id"),
            "name": record.get("name"),
            "labels": record.get("labels", []),
            "properties": props_filtered,
            "relationship": {
                "id": record.get("rel_id"),
                "type": record.get("rel_type"),
                "properties": rel_props_filtered,
                "source_id": record.get("source_id"),
                "source_name": record.get("source_name"),
                "source_labels": record.get("source_labels", [])
            }
        })

    return neighbors


def _build_graph_result(
    entry_node: Optional[Dict[str, Any]],
    best_node: Optional[Dict[str, Any]],
    query: str
) -> Dict[str, Any]:
    """
    构建与operate方法一致的输出格式（只包含入口节点和选中节点）

    Args:
        entry_node: 入口节点
        best_node: 选择的最佳节点
        query: 用户查询

    Returns:
        Dict: 格式化后的结果
    """
    exclude_keys = {
        "last_updated",
        "created_at",
        "knowledge_source",
        "fulltext",
        "embedding",
    }

    nodes_g: List[Dict[str, Any]] = []
    links_g: List[Dict[str, Any]] = []
    node_id_set = set()

    # 添加入口节点
    if entry_node:
        nid = entry_node.get("id")
        if nid and nid not in node_id_set:
            node_id_set.add(nid)

            props = entry_node.get("props", {})
            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 = entry_node.get("labels", [])

            nodes_g.append({
                "id": nid,
                "name": name,
                "labels": labels,
                "properties": props_filtered,
            })

    # 添加最佳节点和关系
    if best_node:
        # 添加最佳节点
        best_nid = best_node.get("id")
        if best_nid and best_nid not in node_id_set:
            node_id_set.add(best_nid)

            props = best_node.get("properties", {})
            name = best_node.get("name") or str(best_nid)
            labels = best_node.get("labels", [])

            nodes_g.append({
                "id": best_nid,
                "name": name,
                "labels": labels,
                "properties": props,
            })

        # 添加关系
        rel = best_node.get("relationship", {})
        if rel:
            rel_id = rel.get("id")
            src_id = rel.get("source_id")
            tgt_id = best_nid

            if src_id in node_id_set and tgt_id in node_id_set:
                links_g.append({
                    "id": rel_id,
                    "label": rel.get("type", ""),
                    "source": src_id,
                    "target": tgt_id,
                    "properties": rel.get("properties", {}),
                })

    # 生成 results 文本（优先用best_node，其次用entry_node）
    results_text = ""
    if best_node:
        best_props = best_node.get("properties", {})
        best_info = {k: v for k, v in best_props.items() if k not in exclude_keys and v}
        results_text = json.dumps(best_info, ensure_ascii=False) if best_info else ""
    elif entry_node:
        entry_props = entry_node.get("props", {})
        entry_info = {k: v for k, v in entry_props.items() if k not in exclude_keys and v}
        results_text = json.dumps(entry_info, ensure_ascii=False) if entry_info else ""

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


async def quick_graph_search(
    query: str,
    route_res: Optional[List[Dict[str, str]]] = None,
    top_k: int = 10,
) -> Dict[str, Any]:
    """
    临时图谱检索接口 - 简化版

    Args:
        query: 用户查询
        route_res: 路由结果列表（本方案不使用）
        top_k: 检索返回的节点数量上限

    Returns:
        Dict: 检索结果，包含 success, xxxx.graph 等字段
    """
    start_time = time.time()

    if not query or not str(query).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:
                logger.info(f"[QuickGraphSearch] 开始检索，查询: '{query}'")

                # 1. 混合检索找入口节点候选
                candidate_nodes = await _search_all_nodes_with_hybrid(session, query, top_k=20)

                if not candidate_nodes:
                    return {"success": False, "error": "未找到任何匹配节点"}

                logger.info(f"[QuickGraphSearch] 混合检索找到 {len(candidate_nodes)} 个候选节点")

                # 2. Rerank 选择最佳入口节点
                reranked_nodes = await _rerank_nodes(query, candidate_nodes)

                if not reranked_nodes:
                    return {"success": False, "error": "Rerank失败"}

                entry_node = reranked_nodes[0]
                entry_id = entry_node.get("id")
                logger.info(f"[QuickGraphSearch] 选择入口节点: {entry_node.get('name', '未知')}, rerank_score={entry_node.get('rerank_score', 0)}")

                # 3. 获取入口节点的一跳邻居
                neighbors = await _get_one_hop_neighbors(session, entry_id)

                logger.info(f"[QuickGraphSearch] 获取到 {len(neighbors)} 个一跳邻居")

                # 4. 对一跳邻居进行 rerank 选择最佳节点
                best_node = None
                if neighbors:
                    # 对邻居节点进行 rerank
                    reranked_neighbors = await _rerank_nodes(query, neighbors)
                    
                    if reranked_neighbors:
                        best_node = reranked_neighbors[0]
                        logger.info(f"[QuickGraphSearch] 选择最佳一跳节点: {best_node.get('name', '未知')}, rerank_score={best_node.get('rerank_score', 0)}")

                # 5. 构建结果（只包含入口节点和选中节点）
                result = _build_graph_result(entry_node, best_node, query)

                elapsed_ms = (time.time() - start_time) * 1000
                logger.info(f"[QuickGraphSearch] 完成，耗时: {elapsed_ms:.2f}ms")

                return result

    except Exception as e:
        logger.error(f"QuickGraphSearch 失败: {e}", exc_info=True)
        return {"success": False, "error": f"检索失败: {str(e)}"}

