import json
from typing import Literal, Optional, Union, List

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

from common.llm import my_llm

# ======================
# 枚举定义
# ======================
EntityType = Literal["Symptom", "Disease", "Formula", "Herb", "Effect", "Source"]

RelationType = Literal[
    "TREATS_DISEASE",
    "ALLEVIATES_SYMPTOM",
    "HAS_EFFECT",
    "HAS_INGREDIENT",
    "HAS_SYMPTOM",
    "FROM_SOURCE"
]


# ======================
# 属性定义
# ======================
class FormulaAttributes(BaseModel):
    alias: Optional[str] = None
    effect: Optional[str] = None
    indication: Optional[str] = None
    taboo: Optional[str] = None
    usage: Optional[str] = None


class HerbAttributes(BaseModel):
    dosage: Optional[str] = None
    effect: Optional[str] = None
    indication: Optional[str] = None
    meridian: Optional[str] = None
    origin: Optional[str] = None
    place: Optional[str] = None
    processing: Optional[str] = None
    property_flavor: Optional[str] = None
    taboo: Optional[str] = None
    traits: Optional[str] = None


# ======================
# 实体与关系结构
# ======================
class Entity(BaseModel):
    name: str
    type: EntityType
    attributes: Optional[Union[FormulaAttributes, HerbAttributes]] = None


class Relation(BaseModel):
    subject: str
    subject_type: EntityType
    relation: RelationType
    object: str
    object_type: EntityType

#Pydantic限定了是他的输出格式
class TCMKnowledgeGraph(BaseModel):
    entities: List[Entity]
    relations: List[Relation]


# 初始化解析器
parser = JsonOutputParser(pydantic_object=TCMKnowledgeGraph)

# 定义 Prompt
prompt = PromptTemplate(
    template=(
        "你是一个中医知识图谱抽取专家。请从以下文本中提取结构化知识：\n"
        "仅当文本中存在实体之间的明确关系时（如‘某方剂治疗某疾病’、‘某药材具有某功效’、‘方剂包含药材’等），才进行抽取。\n"
        "如果文本中仅描述单个实体的信息、未涉及其他实体或关系，请不要抽取，返回空结构：\n"
        "{{\"entities\": [], \"relations\": []}}\n\n"

        "【实体类型说明】\n"
        "- Symptom：症状，如咳嗽、腹痛等\n"
        "- Disease：疾病，如感冒、肺炎、肾虚等\n"
        "- Formula：方剂，如四君子汤、桂枝汤等\n"
        "- Herb：药材，如人参、黄芪、丁香等\n"
        "- Effect：功效，如补气、活血、祛湿、止痛等\n"
        "- Source：出处，如《本草纲目》《伤寒论》等\n\n"

        "【关系类型说明】\n"
        "- TREATS_DISEASE：方剂或药材治疗某种疾病\n"
        "- ALLEVIATES_SYMPTOM：方剂或药材缓解某种症状\n"
        "- HAS_EFFECT：方剂或药材具有某种功效\n"
        "- HAS_INGREDIENT：方剂包含某种药材\n"
        "- HAS_SYMPTOM：疾病包含某种症状\n"
        "- FROM_SOURCE：方剂出自某文献或出处\n\n"

        "若文本涉及方剂或药材，请补充对应的属性字段（如功效、性味、剂量等）。\n"
        "如果文本主要是讲方剂的，请不要抽取药材的属性字段。\n"
        "如果文本主要是讲药材的，请不要抽取方剂的属性字段。\n"
        "如果值为空null，则不必显示键的值。"
        "所有输出必须严格符合以下 JSON 格式：\n"
        "{format_instructions}\n\n"
        "输入文本：{text}"
    ),
    input_variables=["text"],
    partial_variables={"format_instructions": parser.get_format_instructions()}
)


# ======================
# 主函数封装
# ======================
def extract_tcm_knowledge(text: str):
    # 构建抽取链
    chain = prompt | my_llm | parser
    return chain.invoke({"text": text})


# ======================
# Alpaca 微调数据集
# ======================
# instruction 固定不变，input 是中药详情原文，output 是抽取结果的 json 串。
# 常量和构造函数放在这里而不是放在 01_/02_ 脚本里：数字开头的模块名不是合法
# 标识符，两个脚本没法互相 import，共享的东西只能落在这个正常命名的模块。
ALPACA_INSTRUCTION = "请从以下中医文本中抽取知识图谱结构，包括实体与关系。"


def build_alpaca_record(text, entities, relations):
    """把一条抽取结果组装成 Alpaca 格式的样本。

    output 按 Alpaca 约定是字符串，不是嵌套对象，所以这里 dumps 成 json 串。
    ensure_ascii=False 让中文在数据集里保持可读；键顺序固定成
    entities 在前、relations 在后，重跑时同样的抽取结果能得到同样的字节。
    """
    return {
        "instruction": ALPACA_INSTRUCTION,
        "input": text,
        "output": json.dumps(
            {"entities": entities, "relations": relations},
            ensure_ascii=False,
        ),
    }


if __name__ == '__main__':
    text="""
    中药名：一匹绸
分类：化痰止咳平喘药
链接：https://zhongyibaike.com/wiki/一匹绸

## 名称
白面水鸡（《陆川本草》），白背丝绸（《生草药手册》），白底丝绸（《岭南草药志》），绸缎藤、银背藤（《广西药植名录》），一匹绸（《广西中草药》）。

## 一匹绸的种植和炮制

### 来源
为旋花科植物白鹤藤的茎叶。秋季采收。
【原形态】
缠绕藤本。叶卵形至椭圆形，长5～11厘米，宽3～9厘米，先端短尖或钝，基部浑圆或有时微心形，上面秃净或近秃净，下面密被银色紧贴的柔毛。花序腋生或顶生，头状或伞状；花序柄长3.5～7厘米，密被银色柔毛；花柄短，长约5毫米或不及；苞片大，萼片5，矩圆状卵形或卵形，外面的长9～10毫米，内面的较小，外密被银色柔毛；花冠漏斗状，白色，长2.7～2.8厘米，外被银色偃毛，裂片矩圆形，长1.5厘米内外；雄蕊5；子房4室，花柱线形。果近球形，径约8厘米，胭脂色，为卵形凹陷的宿萼所包围。花期秋季。
【生境分布】
生态环境：生于疏林下或路边灌丛中、河边。
资源分布：分布于广东、海南、广西等地。

### 性味
①《陆川本草》：苦辛，凉。
②《岭南草药志》：味微酸微苦，性凉。

### 性状
干燥茎叶，茎扭曲，表面暗灰棕色，有纵沟纹，断面绿白色，有明显的维管束，淡棕色。叶稍卷曲，质脆易碎，表面暗棕色至紫色，背面浅灰绿色，紧贴丝光毛。触之柔软。产广东、广西等地。

## 一匹绸的效果

### 功效
为旋花科植物白鹤藤的茎叶。用于热咳，痰喘，吐血，崩漏，白带，跌打损伤，风湿痛，疮毒。

### 主治
化痰止咳，理血祛风。治热咳，痰喘，吐血，崩漏，白带，跌打损伤，风湿痛，疮毒。
①《本草求原》：叶敷烂脚，化腐疮。
②《岭南采药录》：蛊胀，和米与黄糖捣烂煎香食之；肿胀，和苍术煎汤熏蒸之。
③《陆川本草》：叶：治肺热咳嗽痰喘。
④《广西药植名录》：消痰，润肺，止咳，止血。治肺痨，乳痈，水火烫伤。
⑤《广西中草药》：理血，祛风，除湿。治跌打损伤，内外伤出血及妇女血崩，白带。
⑥《文山中草药》：舒筋活络，祛风除湿，止咳化痰，调经止血。治跌打损伤，风湿疼痛，急慢性支气管炎。

### 用法用量
内服：煎汤，9~15克。
外用：煎水熏洗或捣敷。

## 一匹绸的药方
①治急、慢性支气管炎：白鹤藤（干品）五钱至一两，水煎服。（广州部队《常用中草药手册》）
②治内伤吐血：一匹绸叶、 虎杖 、 旱莲草 、 龙芽草 各一两，水煎服。（《广西中草药》）
③治崩漏：㈠一匹绸叶、 走马胎叶 、 龙芽草 各一两，捣烂，水煎服。（《广西中草药》）㈡紫苞银背藤叶三至四片，研末，酒送服。（《云南中草药》）
④治跌打积瘀，经络不和：一匹绸一两，水煎冲酒服。（《广西中草药》）
⑤治白带：一匹绸一两，小榕树须五钱， 鸡冠花 一两，水煎服。（《广西中草药》）
⑥治疮毒、烂脚：白鹤藤、 蒲公英 、 银花藤 煎水洗之，有消炎、化腐、去毒之效。（《岭南草药志》）
    """
    result=extract_tcm_knowledge(text)
    print(result)