from trauma.api.data.dto import AgeGroup, ContactDetails
from trauma.core.database import MongoBaseModel


class EntityModel(MongoBaseModel):
    id: str
    name: str
    ageGroups: list[AgeGroup]
    treatmentAreas: list[str]
    treatmentMethods: list[str]
    description: str = ''
    contactDetails: ContactDetails


class EntityModelExtended(EntityModel):
    highlightedAgeGroup: AgeGroup | None = None
    highlightedTreatmentArea: str | None = None
    highlightedTreatmentMethod: str | None = None
    topMatch: bool = False
    score: float | None = None