"""
图谱检索工具模块
包含图谱搜索、故障图谱检索、操作图谱检索等功能
"""
import re
import json
import asyncio
import httpx
from typing import Optional, Dict, Any, List
from neo4j import AsyncGraphDatabase
from langchain_core.tools import tool

from config import ATLAS_CONFIG, GRAPH_SEARCH_CONFIG, NEO4J_CONFIG
from modelsAPI.model_api import OpenaiAPI
from utils.function_tracker import track_function_calls


GRAPH_SEARCH_ENDPOINT = GRAPH_SEARCH_CONFIG.get("search_endpoint")


@track_function_calls
async def graph_search(
    query: str,
    top_k: Optional[int] = 40,
    use_rerank: Optional[bool] = False,
    fulltext_top_k: Optional[int] = 15,
    rerank_top_k: Optional[int] = 15,
    hops: Optional[int] = 3,
    verbose: Optional[bool] = False
) -> Dict[str, Any]:
    """
    图谱检索工具
    """
    if not query or not query.strip():
        return {"success": False, "error": "查询文本不能为空"}

    query = query.strip()

    payload = {
        "query": query,
        "use_rerank": use_rerank,
        "top_k": top_k,
        "fulltext_top_k": fulltext_top_k,
        "rerank_top_k": rerank_top_k,
        "hops": hops,
        "verbose": verbose
    }

    headers = {
        "accept": "application/json",
        "Content-Type": "application/json"
    }

    try:
        async with httpx.AsyncClient(timeout=180.0, verify=False) as client:
            response = await client.post(GRAPH_SEARCH_ENDPOINT, json=payload, headers=headers)
    except (httpx.ReadTimeout, httpx.ConnectTimeout, httpx.RequestError) as e:
        return {"success": False, "error": f"请求超时: {str(e)}"}

    if response.status_code != 200:
        return {"success": False, "error": f"HTTP {response.status_code}"}

    result = response.json()
    if result.get("code") != 200:
        msg = result.get("message", "未知错误")
        return {"success": False, "error": msg}

    data = result.get("data", {}) or {}

    try:
        graph_data = []
        results_text = ""

        if isinstance(data, dict):
            raw_nodes = data.get("nodes", [])
            _filter_searchable_labels(raw_nodes)
            # 过滤 nodes 的 properties 字段，仅保留"名称"属性
            for node in raw_nodes:
                if isinstance(node, dict):
                    props = node.get("properties", {}) or node.get("props", {})
                    name_value = props.get("名称", props.get("name", ""))
                    node.pop("properties", None)
                    node.pop("props", None)
                    if name_value:
                        node["properties"] = {"名称": name_value}
            graph_item = {
                "nodes": raw_nodes,
                "links": data.get("links", []),
            }
            graph_item = apply_styles(graph_item)
            graph_data.append(graph_item)
            results_text = data.get("results", "")
        elif isinstance(data, list):
            parts = []
            for item in data:
                if isinstance(item, dict):
                    raw_nodes = item.get("nodes", [])
                    _filter_searchable_labels(raw_nodes)
                    # 过滤 nodes 的 properties 字段，仅保留"名称"属性
                    for node in raw_nodes:
                        if isinstance(node, dict):
                            props = node.get("properties", {}) or node.get("props", {})
                            name_value = props.get("名称", props.get("name", ""))
                            node.pop("properties", None)
                            node.pop("props", None)
                            if name_value:
                                node["properties"] = {"名称": name_value}
                    graph_item = {
                        "nodes": raw_nodes,
                        "links": item.get("links", []),
                    }
                    graph_item = apply_styles(graph_item)
                    graph_data.append(graph_item)
                    if item.get("results"):
                        parts.append(item["results"])
            results_text = "\n".join(parts)

        for graph_item in graph_data:
            links = graph_item.get("links", [])
            for link in links:
                link.pop("properties", None)

        return {
            "success": True,
            "xxxx.graph": graph_data,
            "results": results_text,
        }

    except Exception as e:
        return {"success": False, "error": str(e)}


def _string_similarity(a: str, b: str) -> float:
    """
    简单字符串相似度：综合长度差、包含关系和公共字符比例。
    0~1 之间，越大越相似。
    """
    if not a or not b:
        return 0.0
    a = str(a).strip().lower()
    b = str(b).strip().lower()
    if not a or not b:
        return 0.0
    if a == b:
        return 1.0
    if a in b or b in a:
        return 0.9

    set_a, set_b = set(a), set(b)
    inter = len(set_a & set_b)
    union = len(set_a | set_b) or 1
    jaccard = inter / union

    len_diff = abs(len(a) - len(b))
    len_penalty = max(0.0, 1.0 - len_diff / max(len(a), len(b), 1))

    return 0.5 * jaccard + 0.5 * len_penalty


def _pick_best_by_name(candidates: List[Dict[str, Any]], target_name: str) -> Optional[Dict[str, Any]]:
    """
    在候选节点中，根据名称相似度选择最优节点。
    约定每个节点为：{id, labels, props, display_name}
    """
    if not candidates or not target_name:
        return None

    best = None
    best_score = 0.0
    for item in candidates:
        name = (item.get("display_name") or "").strip()
        if not name:
            props = item.get("props") or {}
            name = props.get("name") or props.get("名称") or props.get("设备名称") or ""
        score = _string_similarity(target_name, name)
        if score > best_score:
            best_score = score
            best = item

    return best if best_score >= 0.3 else None


def _extract_mt_codes(text: str) -> List[str]:
    """
    从文本中抽取形如 MT007 的维修项目编号（MT + 3位数字）。
    """
    if not text:
        return []
    codes = re.findall(r"\bMT\d{3}\b", str(text))
    seen = set()
    ordered: List[str] = []
    for c in codes:
        if c not in seen:
            seen.add(c)
            ordered.append(c)
    return ordered


def _node_record_to_dict(record: Any) -> Dict[str, Any]:
    """
    将 Cypher 返回的 id/labels/props 结构统一为字典。
    """
    if hasattr(record, "data"):
        data = record.data()
    else:
        data = record or {}
    result = {
        "id": data.get("id"),
        "labels": data.get("labels") or [],
        "props": data.get("props") or {},
        "display_name": data.get("display_name")
        if "display_name" in data
        else None,
    }
    if "embedding" in data:
        result["embedding"] = data.get("embedding")
    if "real_device_id" in data:
        result["real_device_id"] = data.get("real_device_id")
    if "real_device_name" in data:
        result["real_device_name"] = data.get("real_device_name")
    return result


def _to_graph_display_node(node_dict: Dict[str, Any]) -> Dict[str, Any]:
    """将内部节点格式转为 graph_search 展示用节点：id、name、label。"""
    nid = node_dict.get("id")
    props = node_dict.get("props") or {}
    name = (
        node_dict.get("display_name")
        or props.get("名称")
        or props.get("name")
        or str(nid)
    )
    labels = [l for l in (node_dict.get("labels") or []) if l != "Searchable"]
    label = labels[0] if labels else ""
    return {"id": nid, "name": name, "label": label}


def _build_fault_graph_display(
    ship_node: Optional[Dict[str, Any]],
    device_node: Optional[Dict[str, Any]],
    fault_mode_node: Optional[Dict[str, Any]],
    repair_projects: List[Dict[str, Any]],
) -> tuple:
    """
    构建与 graph_search 一致的图谱展示结构：nodes 与 links。
    """
    nodes: List[Dict[str, Any]] = []
    links: List[Dict[str, Any]] = []
    seen_ids = set()

    def add_node(nd: Optional[Dict[str, Any]]):
        if not nd or nd.get("id") in seen_ids:
            return
        seen_ids.add(nd["id"])
        nodes.append(_to_graph_display_node(nd))

    def add_link(source_id: Any, target_id: Any, rel_type: str = ""):
        links.append({"source": source_id, "target": target_id, "type": rel_type})

    if ship_node:
        add_node(ship_node)
    if device_node:
        add_node(device_node)
        if ship_node:
            add_link(ship_node["id"], device_node["id"], "包含")
    if fault_mode_node:
        add_node(fault_mode_node)
        if device_node:
            add_link(device_node["id"], fault_mode_node["id"], "存在")
    for p in repair_projects or []:
        add_node(p)
        if device_node:
            add_link(device_node["id"], p["id"], "维修时使用")

    return nodes, links


def _build_fault_graph_results_text(
    ship_node: Optional[Dict[str, Any]],
    device_node: Optional[Dict[str, Any]],
    fault_mode_node: Optional[Dict[str, Any]],
    fault_reason: str,
    repair_plan: str,
    mt_codes: List[str],
    repair_projects: List[Dict[str, Any]],
    ship_number: str,
    device_name: str,
    fault_symptom: str,
) -> str:
    """生成与 graph_rag_search 类似的一段话，供展示与下游使用。"""
    ship_name = ""
    if ship_node:
        ship_name = (ship_node.get("display_name") or (ship_node.get("props") or {}).get("名称") or "").strip()
    device_name_d = ""
    if device_node:
        device_name_d = (device_node.get("display_name") or (device_node.get("props") or {}).get("名称") or "").strip()
    fault_name = ""
    if fault_mode_node:
        fault_name = (fault_mode_node.get("display_name") or (fault_mode_node.get("props") or {}).get("名称") or "").strip()

    parts = [
        f"根据舷号「{ship_number}」、设备「{device_name}」、故障现象「{fault_symptom}」进行图谱检索：",
        f"定位到舰艇「{ship_name or '未知'}」、设备节点「{device_name_d or '未知'}」、故障模式「{fault_name or '未知'}」。",
    ]
    if fault_reason:
        parts.append(f"故障原因：{fault_reason}")
    if repair_plan:
        parts.append(f"维修方案摘要：{repair_plan[:500]}{'…' if len(repair_plan) > 500 else ''}")
    if mt_codes:
        parts.append(f"涉及维修项目编号：{'、'.join(mt_codes)}。")
    if repair_projects:
        proj_names = []
        for p in repair_projects:
            nm = (p.get("display_name") or (p.get("props") or {}).get("名称") or (p.get("props") or {}).get("维修项目编号") or "").strip()
            if nm:
                proj_names.append(nm)
        if proj_names:
            parts.append(f"匹配到的维修项目：{'、'.join(proj_names)}。")
    return "\n".join(parts)


@track_function_calls
async def fault_graph_reason_and_projects(
        ship_number: Optional[str] = None,
        device_name: str = "",
        fault_symptom: str = "",
) -> Dict[str, Any]:
    """
    图谱检索工具
    """

    if not device_name or not str(device_name).strip():
        return {"success": False, "error": "设备名称不能为空"}
    if not fault_symptom or not str(fault_symptom).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:
                device_key = str(device_name).strip()
                ship_id = None
                ship_node = None

                # 有舷号时，先定位舰艇节点
                if ship_number and str(ship_number).strip():
                    ship_key = str(ship_number).strip()
                    ship_query = """
                    MATCH (s:舰艇)
                    WHERE s.舷号 IN [$ship_key]
                    RETURN DISTINCT
                           id(s) AS id,
                           labels(s) AS labels,
                           properties(s) AS props,
                           coalesce(s.名称, s.name, s.舷号) AS display_name
                    LIMIT 10
                    """
                    ship_records = await session.run(ship_query, ship_key=ship_key)
                    ship_candidates = [_node_record_to_dict(r) async for r in ship_records]
                    if ship_candidates:
                        ship_node = ship_candidates[0]
                        ship_id = ship_node["id"]

                if ship_id is not None:
                    # 有舰艇节点：在该舰艇范围内搜索设备
                    device_query = """
                    MATCH (ship:舰艇)
                    WHERE id(ship) = $ship_id
                    MATCH (ship)-[:包含*0..8]->(n:设备)
                    WITH n, coalesce(n.名称, n.name, n.设备名称) AS nm
                    WHERE nm IS NOT NULL
                    RETURN DISTINCT
                           id(n) AS id,
                           labels(n) AS labels,
                           properties(n) AS props,
                           nm AS display_name,
                           n.embedding AS embedding
                    LIMIT 100
                    """
                    dev_records = await session.run(
                        device_query, ship_id=ship_id
                    )
                else:
                    # 无舷号：全局搜索设备节点
                    device_query = """
                    MATCH (n:设备)
                    WITH n, coalesce(n.名称, n.name, n.设备名称) AS nm
                    WHERE nm IS NOT NULL
                    RETURN DISTINCT
                           id(n) AS id,
                           labels(n) AS labels,
                           properties(n) AS props,
                           nm AS display_name,
                           n.embedding AS embedding
                    LIMIT 200
                    """
                    dev_records = await session.run(device_query)

                device_candidates = [_node_record_to_dict(r) async for r in dev_records]
                if not device_candidates:
                    return {"success": False, "error": "未找到匹配的设备节点"}

                device_node, device_score = await OpenaiAPI.hybrid_match_with_embeddings(
                    query_text=device_key,
                    candidates=device_candidates,
                    name_key="display_name",
                    embedding_key="embedding",
                    text_weight=0.3,
                    semantic_weight=0.7
                )
                if not device_node:
                    return {"success": False, "error": f"未找到匹配的设备节点（最高得分: {device_score:.3f}）"}
                device_id = device_node["id"]

                fault_mode_query = """
                MATCH (start_e:设备)
                WHERE id(start_e) = $device_id
                MATCH (start_e)-[:包含*0..3]-(real_e:设备)
                MATCH (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:表征]-(fm:故障模式)
                RETURN DISTINCT
                       id(fm) AS id,
                       labels(fm) AS labels,
                       properties(fm) AS props,
                       fm.名称 AS display_name,
                       fm.embedding AS embedding,
                       id(real_e) AS real_device_id,
                       real_e.名称 AS real_device_name
                """
                fm_records = await session.run(fault_mode_query, device_id=device_id)
                fault_modes = [_node_record_to_dict(r) async for r in fm_records]
                if not fault_modes:
                    return {"success": False, "error": "在该设备节点下未找到故障模式节点"}

                fault_mode_node, fault_score = await OpenaiAPI.hybrid_match_with_embeddings(
                    query_text=fault_symptom,
                    candidates=fault_modes,
                    name_key="display_name",
                    embedding_key="embedding",
                    text_weight=0.3,
                    semantic_weight=0.7
                )
                if not fault_mode_node:
                    return {"success": False, "error": f"未能在故障模式节点中找到与故障现象足够匹配的节点（最高得分: {fault_score:.3f}）"}

                fault_mode_id = fault_mode_node.get("id")
                if fault_mode_id is None:
                    return {"success": False, "error": "故障模式节点缺少 id，无法继续检索维修项目"}

                real_device_id = fault_mode_node.get("real_device_id")

                fm_props = fault_mode_node.get("props") or {}
                fault_reason = fm_props.get("故障原因") or ""
                repair_plan = fm_props.get("维修方案") or ""
                mt_codes = _extract_mt_codes(repair_plan)

                repair_project_query = """
                MATCH (fm:故障模式)
                WHERE id(fm) = $fault_mode_id
                OPTIONAL MATCH (fm)<-[:修理]-(p:维修项目)
                RETURN DISTINCT
                       id(p) AS id,
                       labels(p) AS labels,
                       properties(p) AS props,
                       p.名称 AS display_name
                """
                rp_records = await session.run(
                    repair_project_query, fault_mode_id=fault_mode_id
                )
                all_projects = [_node_record_to_dict(r) async for r in rp_records]

                repair_projects: List[Dict[str, Any]] = []
                if mt_codes:
                    for proj in all_projects:
                        props = proj.get("props") or {}
                        proj_code = (props.get("维修项目编号") or "").strip()
                        if proj_code and (
                                proj_code in mt_codes or any(c in proj_code for c in mt_codes)
                        ):
                            repair_projects.append(proj)
                        elif not proj_code:
                            values_text = " ".join(
                                [str(v) for v in props.values() if v is not None]
                            )
                            if any(code in values_text for code in mt_codes):
                                repair_projects.append(proj)
                else:
                    repair_projects = all_projects

                repair_ids = [
                    p.get("id") for p in repair_projects if p.get("id") is not None
                ]

                exclude_keys = {
                    "last_updated",
                    "created_at",
                    "knowledge_source",
                    "fulltext",
                    "embedding",
                }

                if ship_id is not None:
                    graph_nodes_query = """
                    MATCH (ship:舰艇)
                    WHERE id(ship) = $ship_id
                    MATCH (real_e:设备)
                    WHERE id(real_e) = $real_device_id
                    MATCH path0 = (ship)-[:包含*0..8]->(real_e)
                    MATCH path1 = (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:表征]-(fm:故障模式)
                    WHERE id(fm) = $fault_mode_id
                    OPTIONAL MATCH (fm)<-[:修理]-(rp:维修项目)
                    WHERE id(rp) IN $repair_ids
                    WITH path0, path1, collect(DISTINCT rp) AS rps
                    WITH nodes(path0) + nodes(path1) + rps AS ns
                    UNWIND ns AS n
                    WITH DISTINCT n
                    WHERE n IS NOT NULL
                    RETURN elementId(n) AS id,
                           labels(n) AS labels,
                           coalesce(n.名称, n.name) AS name,
                           properties(n) AS props
                    """
                    node_records = await session.run(
                        graph_nodes_query,
                        ship_id=ship_id,
                        real_device_id=real_device_id,
                        fault_mode_id=fault_mode_id,
                        repair_ids=repair_ids or [-1],
                    )
                else:
                    # 无舷号：不查舰艇路径，只查设备→故障模式→维修项目
                    graph_nodes_query = """
                    MATCH (real_e:设备)
                    WHERE id(real_e) = $real_device_id
                    MATCH path1 = (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:表征]-(fm:故障模式)
                    WHERE id(fm) = $fault_mode_id
                    OPTIONAL MATCH (fm)<-[:修理]-(rp:维修项目)
                    WHERE id(rp) IN $repair_ids
                    WITH path1, collect(DISTINCT rp) AS rps
                    WITH nodes(path1) + rps AS ns
                    UNWIND ns AS n
                    WITH DISTINCT n
                    WHERE n IS NOT NULL
                    RETURN elementId(n) AS id,
                           labels(n) AS labels,
                           coalesce(n.名称, n.name) AS name,
                           properties(n) AS props
                    """
                    node_records = await session.run(
                        graph_nodes_query,
                        real_device_id=real_device_id,
                        fault_mode_id=fault_mode_id,
                        repair_ids=repair_ids or [-1],
                    )

                nodes_g: List[Dict[str, Any]] = []
                async for r in node_records:
                    data_r = r.data()
                    nid = data_r.get("id")
                    if nid is None:
                        continue
                    labels = data_r.get("labels") or []
                    labels = [l for l in labels if l != "Searchable"]
                    name = data_r.get("name") or str(nid)
                    props = data_r.get("props") or {}
                    # 仅保留"名称"属性
                    name_value = props.get("名称", props.get("name", ""))
                    properties = {"名称": name_value} if name_value else {}
                    nodes_g.append(
                        {
                            "id": nid,
                            "name": name,
                            "labels": labels,
                            "properties": properties,
                        }
                    )

                if ship_id is not None:
                    graph_links_query = """
                    MATCH (ship:舰艇)
                    WHERE id(ship) = $ship_id
                    MATCH (real_e:设备)
                    WHERE id(real_e) = $real_device_id
                    MATCH path0 = (ship)-[:包含*0..8]->(real_e)
                    MATCH path1 = (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:表征]-(fm:故障模式)
                    WHERE id(fm) = $fault_mode_id
                    WITH relationships(path0) + relationships(path1) AS rs
                    UNWIND rs AS r
                    RETURN elementId(r) AS id,
                           type(r) AS label,
                           elementId(startNode(r)) AS source,
                           elementId(endNode(r)) AS target,
                           properties(r) AS props
                    UNION
                    MATCH (fm:故障模式)
                    WHERE id(fm) = $fault_mode_id
                    MATCH (fm)<-[r:修理]-(rp:维修项目)
                    WHERE id(rp) IN $repair_ids
                    RETURN elementId(r) AS id,
                           type(r) AS label,
                           elementId(startNode(r)) AS source,
                           elementId(endNode(r)) AS target,
                           properties(r) AS props
                    """
                    link_records = await session.run(
                        graph_links_query,
                        ship_id=ship_id,
                        real_device_id=real_device_id,
                        fault_mode_id=fault_mode_id,
                        repair_ids=repair_ids or [-1],
                    )
                else:
                    # 无舷号：只查设备→故障模式的边
                    graph_links_query = """
                    MATCH (real_e:设备)
                    WHERE id(real_e) = $real_device_id
                    MATCH path1 = (real_e)-[:发生故障使用维修工作]-(w:维修工作)-[:表征]-(fm:故障模式)
                    WHERE id(fm) = $fault_mode_id
                    WITH relationships(path1) AS rs
                    UNWIND rs AS r
                    RETURN elementId(r) AS id,
                           type(r) AS label,
                           elementId(startNode(r)) AS source,
                           elementId(endNode(r)) AS target,
                           properties(r) AS props
                    UNION
                    MATCH (fm:故障模式)
                    WHERE id(fm) = $fault_mode_id
                    MATCH (fm)<-[r:修理]-(rp:维修项目)
                    WHERE id(rp) IN $repair_ids
                    RETURN elementId(r) AS id,
                           type(r) AS label,
                           elementId(startNode(r)) AS source,
                           elementId(endNode(r)) AS target,
                           properties(r) AS props
                    """
                    link_records = await session.run(
                        graph_links_query,
                        real_device_id=real_device_id,
                        fault_mode_id=fault_mode_id,
                        repair_ids=repair_ids or [-1],
                    )

                links_g: List[Dict[str, Any]] = []
                async for r in link_records:
                    data_r = r.data()
                    src = data_r.get("source")
                    tgt = data_r.get("target")
                    if src is None or tgt is None:
                        continue
                    links_g.append(
                        {
                            "id": data_r.get("id"),
                            "label": data_r.get("label", ""),
                            "source": src,
                            "target": tgt,
                        }
                    )

                node_ids = {n["id"] for n in nodes_g}
                links_g = [
                    l
                    for l in links_g
                    if l.get("source") in node_ids and l.get("target") in node_ids
                ]

                repair_plans_data: List[Dict[str, Any]] = []
                for p in repair_projects:
                    props = p.get("props") or {}
                    filtered = {k: v for k, v in props.items() if k not in exclude_keys}
                    if filtered:
                        repair_plans_data.append(filtered)

                parts = []

                if fault_reason and str(fault_reason).strip():
                    parts.append(str(fault_reason).strip())

                if repair_plans_data:
                    for plan in repair_plans_data:
                        if plan is not None:
                            plan_str = str(plan).strip()
                            if plan_str:
                                parts.append(plan_str)

                if parts:
                    combined_text = "\n".join(parts)
                else:
                    combined_text = "暂无故障信息或修复方案"

                results_obj = combined_text

                results_text = json.dumps(results_obj, ensure_ascii=False)

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

    except Exception as e:
        return {"success": False, "error": f"Neo4j 查询失败: {str(e)}"}


def add_node_style(node):
    node_type = node.get("labels", [])
    node_type = [l for l in node_type if l != "Searchable"]
    if node_type and len(node_type) > 0:
        node_type = node_type[0]
    else:
        node_type = ""
    
    if node_type == "舰艇":
        node["color"] = "#F89A64"
        node["size"] = 86
    elif node_type == "系统":
        node["color"] = "#C25C21"
        node["size"] = 86
    elif node_type == "子系统":
        node["color"] = "#7D2D00"
        node["size"] = 86
    elif node_type == "设备":
        node["color"] = "#540802"
        node["size"] = 86
    elif node_type in ["功能", "环境条件", "接口", "安全警告", "调试", "器材保障", "图册", "设计单位", "维修项目", "操作项目"]:
        node["color"] = "#2261DF"
        node["size"] = 52
    else:
        node["color"] = "#68C3FF"
        node["size"] = 53
    
    return node


def add_link_style(link, nodes_dict):
    source_id = link.get("source")
    target_id = link.get("target")
    
    source_type = ""
    target_type = ""
    
    if source_id and source_id in nodes_dict:
        source_labels = [l for l in nodes_dict[source_id].get("labels", []) if l != "Searchable"]
        if source_labels and len(source_labels) > 0:
            source_type = source_labels[0]
    
    if target_id and target_id in nodes_dict:
        target_labels = [l for l in nodes_dict[target_id].get("labels", []) if l != "Searchable"]
        if target_labels and len(target_labels) > 0:
            target_type = target_labels[0]
    
    hierarchy_types = ["舰艇", "系统", "子系统", "设备"]
    if source_type in hierarchy_types and target_type in hierarchy_types:
        link["color"] = "#F59E0B"
    else:
        link["color"] = "#FFFFFF"
    
    return link


def _filter_searchable_labels(nodes: List[Dict[str, Any]]):
    for node in nodes:
        if isinstance(node, dict) and "labels" in node:
            node["labels"] = [l for l in node["labels"] if l != "Searchable"]


def apply_styles(graph_data):
    nodes = graph_data.get("nodes", [])
    links = graph_data.get("links", [])
    
    nodes_dict = {n.get("id"): n for n in nodes if n.get("id")}
    
    for node in nodes:
        add_node_style(node)
    
    for link in links:
        add_link_style(link, nodes_dict)
    
    return graph_data


@tool
async def fault_graph_rag_search(
    ship_number: Optional[str] = None,
    device_name: str = "",
    fault_symptom: str = "",
) -> str:
    """
    故障诊断图谱检索工具（供工作流调用）：仅返回一段话摘要，与 graph_rag_search 一致。
    根据设备名称、故障现象在图谱中定位设备→故障模式→维修项目，并生成文字摘要。舷号为可选参数，提供时可缩小搜索范围。
    工作流中用于与 RAG 结果一起作为上下文，不返回节点/关系（展示用 fault_graph_reason_and_projects）。
    """
    result = await fault_graph_reason_and_projects(
        ship_number=ship_number,
        device_name=device_name,
        fault_symptom=fault_symptom,
    )
    if not result.get("success", False):
        return result.get("error", "检索失败")
    return result.get("results", "")


@track_function_calls
async def operation_graph_search_tool(
        ship_number: str,
        device_name: str,
        operation_item: str,
) -> Dict[str, Any]:
    """
    图谱检索工具
    """
    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": "操作项目不能为空"}

    endpoint = GRAPH_SEARCH_CONFIG.get("operation_search_endpoint")
    timeout = GRAPH_SEARCH_CONFIG.get("graph_timeout")

    if not endpoint:
        return {"success": False, "error": "操作图谱检索服务配置不完整"}

    payload = {
        "ship_number": str(ship_number).strip(),
        "device_name": str(device_name).strip(),
        "operation_item": str(operation_item).strip(),
    }

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(
                endpoint,
                json=payload,
                headers={"Content-Type": "application/json"},
            )
            response.raise_for_status()
            result = response.json()

            if not result.get("success", False):
                return {"success": False, "error": result.get("error", "检索失败")}

            graph_list = result.get("xxxx.graph") or []
            results_parts = []
            if graph_list:
                for i, graph_item in enumerate(graph_list):
                    if isinstance(graph_item, dict):
                        _filter_searchable_labels(graph_item.get("nodes", []))
                        if "results" in graph_item:
                            r = graph_item.pop("results")
                            if r:
                                results_parts.append(r)
                        graph_list[i] = apply_styles(graph_item)
                result["xxxx.graph"] = graph_list

            if results_parts:
                result["results"] = "\n".join(results_parts)
            elif "results" not in result:
                result["results"] = ""

            return result

    except httpx.ReadTimeout:
        return {"success": False, "error": "操作图谱检索服务超时"}
    except httpx.ConnectError:
        return {"success": False, "error": "无法连接操作图谱检索服务"}
    except httpx.HTTPStatusError as e:
        return {"success": False, "error": f"操作图谱检索服务返回错误: {e.response.status_code}"}
    except Exception as e:
        return {"success": False, "error": f"操作图谱检索失败: {str(e)}"}


@tool
async def operation_graph_rag_search(
        ship_number: str,
        device_name: str,
        operation_item: str,
) -> str:
    """
    操作图谱检索工具（供工作流调用）：仅返回一段话摘要，与 graph_rag_search 一致。
    根据舷号、设备名称、操作项目在图谱中定位舰艇→设备→操作项目，并生成文字摘要。
    工作流中用于与 RAG 结果一起作为上下文，不返回节点/关系（展示用 operation_graph_search_tool）。
    """
    result = await operation_graph_search_tool(
        ship_number=ship_number,
        device_name=device_name,
        operation_item=operation_item,
    )
    if not result.get("success", False):
        return result.get("error", "检索失败")
    return result.get("results", "")


@tool
async def graph_rag_search(query: str, entity_type: Optional[str] = None, top_k: int = 40) -> str:
    """
    GraphRAG检索工具（异步版本）
    在知识图谱中检索与query相关的实体和关系

    Args:
        query: 检索查询
        entity_type: 实体类型过滤（如：设备、故障、维修项目）- 注意：此参数在当前实现中暂未使用
        top_k: 返回结果数量（保留参数以兼容接口）

    Returns:
        图谱检索结果文本字符串（旧格式，保持兼容）
    """
    print(f"graph_rag_query+++++++++++{query}")
    result = await graph_search(query)

    if not result.get("success", False):
        return ""

    results_text = result.get("results", "")
    if not results_text:
        return ""
    if len(results_text) > 1000:
        results_text = results_text[:1000] + "..."
    return results_text



async def atlas_retrieval(
    node_names: List[str],
    top_k: int = 10,
    timeout: int = 30
) -> Dict[str, Any]:
    """
    图册检索工具：从图谱中找到节点并关联图册PDF

    Args:
        node_names: 节点名称列表，如 ["主发动机", "压缩机"]
        top_k: 返回的结果数量
        timeout: 请求超时时间（秒）

    Returns:
        包含success和data字段的字典
        data格式示例：
        {
            "船舶主发动机": {
                "船舶主发动机的原理图": [
                    "http://192.168.0.46:59085/files/163-06A0015-B01001_船舶主发动机原理图.pdf"
                ]
            }
        }
    """
    api_url = ATLAS_CONFIG["endpoint"]

    if not node_names:
        return {
            "success": False,
            "error": "node_names不能为空",
            "data": {}
        }

    payload = {
        "node_names": node_names,
        "top_k": top_k
    }

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(
                api_url,
                json=payload,
                headers={
                    "accept": "application/json",
                    "Content-Type": "application/json"
                }
            )

        if response.status_code == 200:
            result = response.json()
            if result.get("success", False):
                return {
                    "success": True,
                    "data": result.get("data", {}),
                    "error": None
                }
            else:
                return {
                    "success": False,
                    "error": result.get("message", "API返回失败"),
                    "data": {}
                }
        else:
            return {
                "success": False,
                "error": f"HTTP {response.status_code}: {response.text}",
                "data": {}
            }

    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "data": {}
        }


def format_atlas_results(atlas_data: Dict[str, Any]) -> str:
    """
    格式化图册检索结果为文本，便于在提示词中使用

    Args:
        atlas_data: atlas_retrieval返回的data字段

    Returns:
        格式化后的文本
    """
    if not atlas_data:
        return ""

    parts = []
    for node_name, node_data in atlas_data.items():
        if isinstance(node_data, dict):
            for title, urls in node_data.items():
                if isinstance(urls, list) and urls:
                    parts.append(f"\n【{node_name} - {title}】")
                    for url in urls:
                        # 确保URL后面有明确的换行，避免和其他文字连在一起
                        parts.append(f"- {url}\n")

    return "\n".join(parts) if parts else ""


async def extract_entity_names_from_history(
    history_messages: List[Dict[str, Any]],
    device_name: str = "",
    operation_item: str = "",
    fault: str = "",
    last_generated_scheme: str = ""
) -> List[str]:
    """
    从历史记录中抽取部件/设备/零部件名称

    Args:
        history_messages: 历史对话记录
        device_name: 设备名称
        operation_item: 操作项目名称
        fault: 故障现象
        last_generated_scheme: 最后一次生成的方案（用于重点抽取）

    Returns:
        抽取的实体名称列表
    """
    entity_names = []

    # 添加明确提供的名称
    if device_name:
        entity_names.append(device_name)
    if operation_item:
        entity_names.append(operation_item)
    if fault:
        entity_names.append(fault)

    # 去重
    entity_names = list(dict.fromkeys(entity_names))
    
    # 调用大模型进一步抽取和补充
    if entity_names or last_generated_scheme:
        try:
            extraction_prompt = f"""从以下信息中抽取尽可能多的设备、部件、零部件、机构、系统名称（用于图册检索）。
注意：
- 请尽可能多抽取，宁多勿少，准不准没关系
- 包括所有提及的机械部件、电子设备、机构系统等
- 即使是重复提到的也可以保留，我们会去重
- 重点从【生成的方案】中抽取，这里信息最全面

【示例 - 输入】
船舶主发动机（舷号：101）故障分析与维修方案
1. 故障概况
设备名称：船舶主发动机
舷号：101
故障现象：主机运行中出现异响
2. 故障原因分析
主轴承或连杆大端轴承磨损：轴承间隙过大或磨损严重会导致金属撞击声。
曲轴主轴颈跳动超差：曲轴旋转不平衡或跳动量超出标准。
飞轮连接螺栓松动：飞轮与曲轴连接处的螺栓松动。
活塞头烧蚀或活塞环断裂：活塞环断裂或活塞头部损坏可能导致敲缸声。
排气阀密封面磨损或阀座烧蚀：气门间隙异常或密封失效。
涡轮增压器轴承损坏：增压器内部轴承损坏。
3. 维修方案
执行项目：MT003 检测曲轴组件（含飞轮紧固状态检查）。
执行项目：MTO03 检测曲轴主轴颈跳动量。

【示例 - 输出】
{{
    "entity_names": ["船舶主发动机", "主轴承", "连杆大端轴承", "曲轴", "飞轮", "飞轮连接螺栓", "活塞头", "活塞环", "排气阀", "排气阀座", "涡轮增压器", "涡轮增压器轴承", "曲轴组件"]
}}

【设备/操作】
- 设备名: {device_name or '无'}
- 操作项: {operation_item or '无'}
- 故障: {fault or '无'}

【生成的方案】
{last_generated_scheme or '无'}

请输出JSON格式，格式如下：
{{
    "entity_names": ["名称1", "名称2", "名称3", "..."]
}}

仅输出JSON，不要其他内容。"""

            response = await OpenaiAPI.open_api_chat_without_thinking(
                query=extraction_prompt,
                json_output=True,
                system_prompt="你是一个专业的实体抽取助手，负责从文本中抽取尽可能多的设备、部件、零部件、机构、系统名称。请尽量多抽取，宁多勿少。重点从生成的维修方案中抽取。"
            )

            parsed = _safe_json_extract(response)
            if isinstance(parsed, dict):
                extracted = parsed.get("entity_names", [])
                if extracted:
                    entity_names.extend(extracted)
                    entity_names = list(dict.fromkeys(entity_names))
        except Exception as e:
            print(f"[图册检索] 实体抽取失败: {str(e)}")

    # 过滤空值和太短的名称
    entity_names = [name.strip() for name in entity_names if name and len(name.strip()) >= 2]
    return entity_names[:30]  # 最多30个，尽量多找图册


def _safe_json_extract(text: str) -> Any:
    """安全地从文本中提取JSON"""
    try:
        # 尝试直接解析
        return json.loads(text)
    except json.JSONDecodeError:
        # 尝试用正则提取
        match = re.search(r'\{[\s\S]*\}', text)
        if match:
            try:
                return json.loads(match.group(0))
            except json.JSONDecodeError:
                pass
        return None


