"""用户输入实体抽取节点

输入: state['input_semantic_trans']  语义转写节点补全后的提问
输出: state['user_input_effects'] / _diseases / _symptoms / _formulas / _herbs / _sources

读转写结果而不是 state['input']：原话里的指代词（"它"、"那个方子"）抽不出实体，
转写节点已经把它们还原成具体名字了。

抽取的六个实体类型和 `_002__extract_information` 的 EntityType 一一对应，
抽出来的名字后面要拿去和 Neo4j 里的节点名做匹配，所以只抽名字、不抽关系，
也不做任何改写（"人参"不能变成"人参片"，否则匹配不上）。
"""

from typing import List

from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from pydantic import BaseModel, Field

from _004_langgraph_more_nodes.agent_state import AgentState
from common.llm import my_llm


# ======================
# 输出结构
# ======================
class UserInputEntities(BaseModel):
    """用户提问里出现的实体，按类型分组。没有的类型给空列表。"""

    effects: List[str] = Field(default_factory=list, description="功效，如补气、活血、祛湿、止痛")
    diseases: List[str] = Field(default_factory=list, description="疾病，如感冒、肺炎、肾虚")
    symptoms: List[str] = Field(default_factory=list, description="症状，如咳嗽、腹痛")
    formulas: List[str] = Field(default_factory=list, description="方剂，如四君子汤、桂枝汤")
    herbs: List[str] = Field(default_factory=list, description="药材，如人参、黄芪、丁香")
    sources: List[str] = Field(default_factory=list, description="出处，如《本草纲目》《伤寒论》")


parser = JsonOutputParser(pydantic_object=UserInputEntities)

prompt = PromptTemplate(
    template=(
        "你是一个中医实体识别专家。请从用户的提问中抽取出现的中医实体。\n\n"
        "【实体类型说明】\n"
        "- symptoms：症状，如咳嗽、腹痛等\n"
        "- diseases：疾病，如感冒、肺炎、肾虚等\n"
        "- formulas：方剂，如四君子汤、桂枝汤等\n"
        "- herbs：药材，如人参、黄芪、丁香等\n"
        "- effects：功效，如补气、活血、祛湿、止痛等\n"
        "- sources：出处，如《本草纲目》《伤寒论》等\n\n"
        "【抽取要求】\n"
        "1. 只抽取用户提问中真实出现的词，不要补充、联想或推理出提问里没有的实体。\n"
        "2. 保持原文写法，不要改写、不要加后缀、不要翻译成术语，出处保留书名号。\n"
        "3. 一个词只归入最贴切的一个类型，不要重复出现在多个类型里。\n"
        "4. 去掉'的'、'吗'、'怎么'这类修饰，只留实体名本身。\n"
        "5. 某个类型没有对应实体时给空列表 []，不要省略这个键。\n"
        "6. 疑问词（什么、哪些、如何）和非中医词汇一律不抽。\n\n"
        "所有输出必须严格符合以下 JSON 格式：\n"
        "{format_instructions}\n\n"
        "用户提问：{text}"
    ),
    input_variables=["text"],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

# 模型返回的键 -> 写回 state 的键
FIELD_TO_STATE_KEY = {
    "effects": "user_input_effects",
    "diseases": "user_input_diseases",
    "symptoms": "user_input_symptoms",
    "formulas": "user_input_formulas",
    "herbs": "user_input_herbs",
    "sources": "user_input_sources",
}


def _clean(values) -> List[str]:
    """规整成去空、去重、保序的字符串列表。

    JsonOutputParser 返回的是普通 dict，模型偶尔会给出 null、单个字符串或
    重复项，所以这里不能直接信任它的形状。
    """
    if values is None:
        return []
    if isinstance(values, str):
        values = [values]

    cleaned, seen = [], set()
    for value in values:
        name = str(value).strip()
        if name and name not in seen:
            seen.add(name)
            cleaned.append(name)
    return cleaned


def extract_entity_from_user_input_node(state: AgentState) -> AgentState:
    """从用户输入中抽取六类实体，分别写回 state"""
    print("实体抽取节点执行开始了")
    user_input = (state.get("input_semantic_trans") or "").strip()

    # 空输入不必调用大模型，六个键仍然要写，下游节点直接取值不用判空
    if not user_input:
        for state_key in FIELD_TO_STATE_KEY.values():
            state[state_key] = []
        return state

    chain = prompt | my_llm | parser
    result = chain.invoke({"text": user_input}) or {}

    for field, state_key in FIELD_TO_STATE_KEY.items():
        state[state_key] = _clean(result.get(field))

    total = sum(len(state[key]) for key in FIELD_TO_STATE_KEY.values())
    print(f"实体抽取节点执行完成: 共 {total} 个")
    return state


if __name__ == "__main__":
    for text in [
        "人参和黄芪都能补气吗？",
        "桂枝汤出自《伤寒论》，主治什么病？",
        "我最近咳嗽还有点腹痛，肾虚会引起这些吗？",
        "四君子汤里有哪些药材？",
        "今天天气怎么样",
        "",
    ]:
        result = extract_entity_from_user_input_node({"input_semantic_trans": text})
        print(f"\n提问: {text!r}")
        for state_key in FIELD_TO_STATE_KEY.values():
            print(f"  {state_key}: {result[state_key]}")