"""实体标准化节点：把用户嘴里的词对到 Neo4j 里的标准实体名

输入: state['user_input_effects'] / _diseases / _symptoms / _formulas / _herbs / _sources
输出: state['matched_effects']    / _diseases / _symptoms / _formulas / _herbs / _sources

标准实体名存在 `_003_create_neo4j_database/` 的两个文件里（`__003__faiss_embedding.py`
建的）：index 存向量，pkl 存 id -> 实体名 的映射，faiss 本身存不了元数据。

相似度换算：建库用的是 IndexFlatL2，向量是单位长度（bge-large-zh-v1.5 最后一层自带
Normalize），此时 d² = 2 - 2cos，而 faiss 的 IndexFlatL2 返回的就是平方距离，取值
落在 [0, 2]。所以 cos = 1 - dist / 2，阈值 70% 相似度等价于 dist <= 0.6。
"""

import os
import pickle
from typing import Dict, List, Tuple

import faiss
import numpy as np

from _004_langgraph_more_nodes.agent_state import AgentState
from common.embedding_model import embedding_model
from common.path_utils import get_file_path

INDEX_PATH = get_file_path("_003_create_neo4j_database/nero4j_embedding_faiss.index")
MAPPING_PATH = get_file_path(
    "_003_create_neo4j_database/nero4j_embedding_faiss_id2text.pkl"
)

# 需求：取 top5，相似度大于 70%
TOP_K = 5
SIMILARITY_THRESHOLD = 0.80

# Neo4j 标签 -> (输入 state 键, 输出 state 键)
TYPE_FIELDS = {
    "Effect": ("user_input_effects", "matched_effects"),
    "Disease": ("user_input_diseases", "matched_diseases"),
    "Symptom": ("user_input_symptoms", "matched_symptoms"),
    "Formula": ("user_input_formulas", "matched_formulas"),
    "Herb": ("user_input_herbs", "matched_herbs"),
    "Source": ("user_input_sources", "matched_sources"),
}

# 索引 75MB，pkl 也不小，按需加载一次后缓存，别在 import 时就吃掉这份内存
_index = None
_id2text = None


def _load_faiss():
    """加载 index 和 id2text 映射，进程内只做一次"""
    global _index, _id2text
    if _index is not None:
        return _index, _id2text

    for path in (INDEX_PATH, MAPPING_PATH):
        if not os.path.exists(path):
            raise FileNotFoundError(
                f"faiss 文件不存在: {path}\n"
                "先跑 python -X utf8 -m _003_create_neo4j_database.__003__faiss_embedding"
            )

    _index = faiss.read_index(INDEX_PATH)
    with open(MAPPING_PATH, "rb") as f:
        _id2text = pickle.load(f)
    return _index, _id2text


# 标签 -> 该标签下的实体名集合。索引里所有标签的实体混在一起（建库用的是不带
# label 的 get_all_node_names()），不按标签过滤的话，"咳嗽"可能召回一个药材名，
# 落到 matched_symptoms 里就会生成没意义的 cypher。
_label_names = None


def _load_label_names() -> Dict[str, set]:
    """从 Neo4j 读每个标签下的实体名，用来过滤召回结果。

    读不到就返回空 dict，此时退化成不按标签过滤 —— 只是召回可能串类型，
    不至于让整个节点挂掉。
    """
    global _label_names
    if _label_names is not None:
        return _label_names

    try:
        from common.neo4j_manager import neo4j_client

        _label_names = {
            label: set(neo4j_client.get_all_node_names(label)) for label in TYPE_FIELDS
        }
    except Exception as exc:
        print(f"[warn] 读取 Neo4j 标签实体名失败，跳过按标签过滤: {exc}")
        _label_names = {}
    return _label_names


def search_similar_entities(
    query: str, label: str = None, top_k: int = TOP_K, threshold: float = SIMILARITY_THRESHOLD
) -> List[Tuple[str, float]]:
    """在 faiss 里检索 query 的相似标准实体，返回 [(实体名, 相似度)]，按相似度降序。

    label 非空时只保留属于该标签的实体。为了过滤掉别的标签后仍能凑够 top_k，
    检索时多要一些候选（top_k * 6），再截断到 top_k。
    """
    query = (query or "").strip()
    if not query:
        return []

    index, id2text = _load_faiss()
    allowed = _load_label_names().get(label) if label else None
    fetch_k = min(index.ntotal, top_k * 6 if allowed else top_k)

    # 建库时用了 normalize_embeddings=True，查询保持一致
    vec = embedding_model.encode(
        [query], convert_to_numpy=True, normalize_embeddings=True
    )
    distances, indices = index.search(np.asarray(vec, dtype="float32"), fetch_k)

    hits = []
    for dist, idx in zip(distances[0], indices[0]):
        if idx < 0:  # faiss 用 -1 填充不足的位置
            continue
        name = id2text.get(int(idx))
        if name is None or (allowed is not None and name not in allowed):
            continue
        # 单位向量下 IndexFlatL2 返回平方距离，cos = 1 - dist/2
        similarity = 1.0 - float(dist) / 2.0
        if similarity >= threshold:
            hits.append((name, similarity))
        if len(hits) >= top_k:
            break
    return hits


def match_entity_from_neo4j_node(state: AgentState) -> AgentState:
    """把六类用户输入实体分别标准化成 Neo4j 里的实体名，写回 matched_*"""
    print("实体匹配节点执行开始了")
    for label, (in_key, out_key) in TYPE_FIELDS.items():
        matched, seen = [], set()
        for query in state.get(in_key) or []:
            for name, _similarity in search_similar_entities(query, label=label):
                # 多个输入词可能召回同一个实体（"拉肚子"和"腹泻"），去重保序
                if name not in seen:
                    seen.add(name)
                    matched.append(name)
        state[out_key] = matched

    total = sum(len(state[out_key]) for _in_key, out_key in TYPE_FIELDS.values())
    print(f"实体匹配节点执行完成: 共 {total} 个")
    return state


if __name__ == "__main__":
    # 单个词的召回明细：看清相似度和阈值起没起作用
    print("\n=== 单词检索明细（top5，阈值 70%）===")
    for label, query in [
        ("Symptom", "拉肚子"),
        ("Symptom", "咳嗽"),
        ("Herb", "人参"),
        ("Formula", "四君子汤"),
        ("Effect", "补气"),
        ("Source", "《伤寒论》"),
        ("Disease", "肾虚"),
        ("Herb", "阿司匹林"),  # 图里没有，预期空
    ]:
        hits = search_similar_entities(query, label=label)
        shown = " | ".join(f"{n}({s:.1%})" for n, s in hits) or "<无命中>"
        print(f"  [{label}] {query} → {shown}")

    # 整个节点：接 _002_ 抽取节点的输出形状
    print("\n=== 节点整体 ===")
    state = {
        "user_input_herbs": ["人参", "黄芪"],
        "user_input_effects": ["补气"],
        "user_input_symptoms": ["拉肚子", "睡不着觉"],
        "user_input_diseases": ["肾虚"],
        "user_input_formulas": ["四君子汤"],
        "user_input_sources": ["《伤寒论》"],
    }
    result = match_entity_from_neo4j_node(state)
    print(result)
    # for _in_key, out_key in TYPE_FIELDS.values():
    #     print(f"  {out_key}: {result[out_key]}")
    #
    # # 空输入：六个键仍要写满
    # print("\n=== 空输入 ===")
    # print(match_entity_from_neo4j_node({}))