from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation


def get_trip_ontology():
    # Attraction + Restaurant + City + Country Ontology
    trip_data_ontology = Ontology()

    trip_data_ontology.add_entity(
        Entity(
            label="Country",
            attributes=[
                Attribute(
                    name="name",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
            ],
        )
    )
    trip_data_ontology.add_entity(
        Entity(
            label="City",
            attributes=[
                Attribute(
                    name="name",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
                Attribute(
                    name="weather",
                    attr_type=AttributeType.STRING,
                    required=False,
                    unique=False,
                ),
                Attribute(
                    name="population",
                    attr_type=AttributeType.NUMBER,
                    required=False,
                    unique=False,
                ),
            ],
        )
    )
    trip_data_ontology.add_entity(
        Entity(
            label="Restaurant",
            attributes=[
                Attribute(
                    name="name",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
                Attribute(
                    name="description",
                    attr_type=AttributeType.STRING,
                    required=False,
                    unique=False,
                ),
                Attribute(
                    name="rating",
                    attr_type=AttributeType.NUMBER,
                    required=False,
                    unique=False,
                ),
                Attribute(
                    name="food_type",
                    attr_type=AttributeType.STRING,
                    required=False,
                    unique=False,
                ),
            ],
        )
    )
    trip_data_ontology.add_entity(
        Entity(
            label="Attraction",
            attributes=[
                Attribute(
                    name="name",
                    attr_type=AttributeType.STRING,
                    required=True,
                    unique=True,
                ),
                Attribute(
                    name="description",
                    attr_type=AttributeType.STRING,
                    required=False,
                    unique=False,
                ),
                Attribute(
                    name="type",
                    attr_type=AttributeType.STRING,
                    required=False,
                    unique=False,
                ),
            ],
        )
    )
    trip_data_ontology.add_relation(
        Relation(
            label="IN_COUNTRY",
            source="City",
            target="Country",
        )
    )
    trip_data_ontology.add_relation(
        Relation(
            label="IN_CITY",
            source="Restaurant",
            target="City",
        )
    )
    trip_data_ontology.add_relation(
        Relation(
            label="IN_CITY",
            source="Attraction",
            target="City",
        )
    )
    return trip_data_ontology
