import os
from pydantic import BaseModel
from typing import Dict, List, Optional


class NodeSchema(BaseModel):
    label: str
    properties: Dict[str, str]  # 属性名: 类型


class RelationshipSchema(BaseModel):
    type: str
    from_node: str              # 起始节点标签
    to_node: str                # 目标节点标签
    properties: Dict[str, str]  # 属性名: 类型


class GraphSchema(BaseModel):
    nodes: List[NodeSchema]
    relationships: List[RelationshipSchema]


# 示例图模式 (按照neo4j数据库中的定义schema来填充)
EXAMPLE_SCHEMA = GraphSchema(
    # 节点的名称一定要严格保持跟neo4j一致
    nodes=[
        # --- PDF 原有的 4 个节点 ---
        NodeSchema(label="Disease", properties={"name": "string", "desc": "string", "cause": "string", "prevent": "string", "cure_lasttime": "string", "cure_department": "string", "cure_way": "string", "cure_prob": "string", "easy_get": "string"}),
        NodeSchema(label="Drug", properties={"name": "string"}),
        NodeSchema(label="Food", properties={"name": "string"}),
        NodeSchema(label="Symptom", properties={"name": "string"}),
        # --- 基于 Neo4j 截图新增的 3 个节点 ---
        NodeSchema(label="Check", properties={"name": "string"}),
        NodeSchema(label="Department", properties={"name": "string"}),
        NodeSchema(label="Producer", properties={"name": "string"}),
    ],
    # 关系的相关字段一定要严格保持跟neo4j一致, 大小写都不能错
    relationships=[
        # --- PDF 原有的 3 个关系 ---
        RelationshipSchema(
            type="has_symptom",
            from_node="Disease",
            to_node="Symptom",
            properties={}
        ),
        RelationshipSchema(
            type="recommand_drug",
            from_node="Disease",
            to_node="Drug",
            properties={}
        ),
        RelationshipSchema(
            type="recommand_eat",
            from_node="Disease",
            to_node="Food",
            properties={}
        ),
        # --- 基于 Neo4j 截图新增的关系 ---
        # Disease 需要做的检查项目
        RelationshipSchema(
            type="need_check",
            from_node="Disease",
            to_node="Check",
            properties={}
        ),
        # Disease 所属的科室
        RelationshipSchema(
            type="belongs_to",
            from_node="Disease",
            to_node="Department",
            properties={}
        ),
        # Disease 的常用药物
        RelationshipSchema(
            type="common_drug",
            from_node="Disease",
            to_node="Drug",
            properties={}
        ),
        # Disease 宜吃的食物
        RelationshipSchema(
            type="do_eat",
            from_node="Disease",
            to_node="Food",
            properties={}
        ),
        # Disease 忌吃的食物
        RelationshipSchema(
            type="no_eat",
            from_node="Disease",
            to_node="Food",
            properties={}
        ),
        # Disease 的并发症
        RelationshipSchema(
            type="acompany_with",
            from_node="Disease",
            to_node="Disease",
            properties={}
        ),
        # Drug 的生产商
        RelationshipSchema(
            type="drugs_of",
            from_node="Producer",
            to_node="Drug",
            properties={}
        ),
    ]
)


if __name__ == '__main__':
    res = str(EXAMPLE_SCHEMA.model_dump())
    print(res)