class Entity(BaseModel):
    """The base class of all ontology classes.

    Args:
        id (str): IRI
    """

    id: str = Field(alias="@id", description="IRI (possibly relative)", title="@id")

    @computed_field(alias="@type", title="@type", description="JSON-LD @type")
    @property
    def type(self) -> str:
        return type(self).__name__

    model_config = ConfigDict(populate_by_name=True, serialize_by_alias=True)

class JSONLDGraph(BaseModel):
    """Encapsulates a JSON-LD document/graph.

    This is the return type of the make_model() function, and
    ontology_graph(), shacl_graph() class methods.

    Args:
        context (dict): JSON-LD context (mapped to @context)
        graph (List): Default graph containing entities (mapped to @graph)
    """

    context: SkipJsonSchema[dict] = Field(
        default={
            "@vocab": "http://example.com/vocab/",
            "@base": "http://example.com/",
        },
        alias="@context",
        title="@context",
        description="JSON-LD context",
    )

    graph: List = Field(
        default=[], alias="@graph", title="@graph", description="Default graph"
    )

    model_config = ConfigDict(populate_by_name=True, serialize_by_alias=True)

    @classmethod
    def ontology_graph(cls) -> "JSONLDGraph":
        """Generate an ontology graph from the classes in the ontology.

        Returns:
            JSONLDGraph: With _OntologyClass and _OntologyProperty (internal classes) instances.
        """

        # Collect unique entity types from the model
        entity_classes = cls._collect_entity_classes_from_model()

        # Generate class and property definitions
        ontology_entities: List[_OntologyClass | _OntologyProperty] = []
        properties_seen = set()

        for entity_class in entity_classes:
            # Add class definition
            class_def = cls._create_class_definition(entity_class, entity_classes)
            ontology_entities.append(class_def)

            # Add property definitions
            property_defs = cls._create_property_definitions(
                entity_class, properties_seen
            )
            ontology_entities.extend(property_defs)

        # Extract vocab from context
        vocab = cls.model_fields["context"].default.get(
            "@vocab", "http://example.com/vocab/"
        )

        # Build ontology context
        ontology_context = {
            "@vocab": vocab,
            "@base": vocab,
            "rdfs": "http://www.w3.org/2000/01/rdf-schema#",
            "owl": "http://www.w3.org/2002/07/owl#",
        }

        # Return as JSONLDGraph
        return JSONLDGraph(context=ontology_context, graph=ontology_entities)

    @classmethod
    def shacl_graph(cls) -> "JSONLDGraph":
        """Generate SHACL shapes graph from the classes in the ontology.

        Returns:
            JSONLDGraph: With _NodeShape and _PropertyShape (internal classes) instances.
        """

        # Collect unique entity types from the model
        entity_classes = cls._collect_entity_classes_from_model()

        # Generate node shapes
        shacl_shapes: List[_NodeShape] = []

        for entity_class in entity_classes:
            # Create node shape for this class
            node_shape = cls._create_node_shape(entity_class)
            shacl_shapes.append(node_shape)

        # Extract vocab from context
        vocab = cls.model_fields["context"].default.get(
            "@vocab", "http://example.com/vocab/"
        )

        # Build SHACL context
        shacl_context = {
            "@vocab": vocab,
            "@base": vocab,
            "sh": "http://www.w3.org/ns/shacl#",
            "xsd": "http://www.w3.org/2001/XMLSchema#",
        }

        # Return as JSONLDGraph
        return JSONLDGraph(context=shacl_context, graph=shacl_shapes)

    @classmethod
    def _collect_entity_classes_from_model(cls) -> set:
        """Collect all Entity classes from the model's graph field type annotation.

        Returns:
            set: All Entity classes found in the model's graph field annotation.
        """
        entity_classes = set()

        # Get the graph field's annotation
        graph_field = cls.model_fields.get("graph")
        if graph_field and graph_field.annotation:
            annotation = graph_field.annotation

            # Extract the entity type from List[EntityType] or List[A|B|C]
            origin = get_origin(annotation)
            if origin is list or origin is List:
                args = get_args(annotation)
                if args:
                    inner_type = args[0]

                    # Check if inner type is a Union (A|B|C or Union[A, B, C])
                    inner_origin = get_origin(inner_type)
                    # Handle both Union[A, B, C] and A|B|C (UnionType in Python 3.10+)
                    if inner_origin is Union or isinstance(inner_type, UnionType):
                        # Handle List[A|B|C] - extract all union members
                        union_args = get_args(inner_type)
                        for entity_type in union_args:
                            if isinstance(entity_type, type) and issubclass(
                                entity_type, Entity
                            ):
                                entity_classes.add(entity_type)
                                entity_classes.update(
                                    cls._get_all_subclasses(entity_type)
                                )
                    else:
                        # Handle List[EntityType] - single entity type
                        if isinstance(inner_type, type) and issubclass(
                            inner_type, Entity
                        ):
                            entity_classes.add(inner_type)
                            entity_classes.update(cls._get_all_subclasses(inner_type))

        return entity_classes

    @classmethod
    def _get_all_subclasses(cls, entity_class) -> set:
        """Recursively get all subclasses of an entity class.

        Args:
            entity_class: The entity class to find subclasses for.

        Returns:
            set: All subclasses of the given entity class.
        """
        subclasses = set()
        for subclass in entity_class.__subclasses__():
            subclasses.add(subclass)
            subclasses.update(cls._get_all_subclasses(subclass))
        return subclasses

    @classmethod
    def _create_class_definition(
        cls, entity_class, entity_classes: set
    ) -> _OntologyClass:
        """Create an RDFS class definition for an Entity class.

        Args:
            entity_class: The Entity class to create a definition for.
            entity_classes: Set of all entity classes in the model.

        Returns:
            _OntologyClass: RDFS class definition.
        """
        class_def = _OntologyClass(
            id=entity_class.__name__,
            label=entity_class.__name__,
            comment=entity_class.__doc__.strip() if entity_class.__doc__ else None,
        )

        # Add inheritance relationships
        for base in entity_class.__bases__:
            if base != Entity and issubclass(base, Entity):
                class_def.subClassOf = Relation(id=base.__name__)
                entity_classes.add(base)

        return class_def

    @classmethod
    def _create_property_definitions(
        cls, entity_class, properties_seen: set
    ) -> List[_OntologyProperty]:
        """Create property definitions for fields defined on a class.

        Args:
            entity_class: The Entity class to create property definitions for.
            properties_seen: Set of property names already processed.

        Returns:
            List[_OntologyProperty]: Property definitions for the class fields.
        """
        property_defs = []
        parent_fields = cls._get_parent_fields(entity_class)

        for field_name, field_info in entity_class.model_fields.items():
            # Skip special fields and inherited fields
            if field_name in ["id", "type"] or field_name in parent_fields:
                continue

            # Skip duplicates
            if field_name in properties_seen:
                continue
            properties_seen.add(field_name)

            # Create property definition
            prop_def = cls._create_single_property_definition(
                entity_class, field_name, field_info, field_name
            )
            property_defs.append(prop_def)

        return property_defs

    @classmethod
    def _get_parent_fields(cls, entity_class) -> set:
        """Get all field names from parent classes.

        Args:
            entity_class: The Entity class to get parent fields for.

        Returns:
            set: Field names from all parent classes.
        """
        parent_fields = set()
        for base in entity_class.__bases__:
            if base != Entity and issubclass(base, Entity):
                parent_fields.update(base.model_fields.keys())
        return parent_fields

    @classmethod
    def _create_single_property_definition(
        cls, entity_class, field_name: str, field_info, prop_name: str
    ) -> _OntologyProperty:
        """Create a single property definition.

        Args:
            entity_class: The Entity class containing the field.
            field_name: Name of the field.
            field_info: Pydantic field information.
            prop_name: Name for the property.

        Returns:
            _OntologyProperty: Property definition for the field.
        """
        # Extract type and metadata
        field_type, metadata = cls._extract_field_type_and_metadata(
            entity_class, field_name, field_info
        )

        # Determine if ObjectProperty or DatatypeProperty
        is_relation = cls._is_relation_type(field_type)

        prop_def = _OntologyProperty(
            id=prop_name,
            type="owl:ObjectProperty" if is_relation else "owl:DatatypeProperty",
            label=prop_name,
            domain=Relation(id=entity_class.__name__),
            comment=field_info.description,
        )

        # Add properties to _OntologyProperty according to annotation
        for meta in metadata:
            if isinstance(meta, RDFSAnnotation.DOMAIN):
                prop_def.domain = Relation(id=meta.value)
            if isinstance(meta, RDFSAnnotation.RANGE):
                prop_def.range = Relation(id=meta.value)

        return prop_def

    @classmethod
    def _create_node_shape(cls, entity_class) -> _NodeShape:
        """Create a SHACL node shape for an Entity class.

        Args:
            entity_class: The Entity class to create a node shape for.

        Returns:
            _NodeShape: SHACL node shape for the entity class.
        """
        # Create property shapes for all fields
        property_shapes = []

        for field_name, field_info in entity_class.model_fields.items():
            # Skip special fields
            if field_name in ["id", "type"]:
                continue

            # Create property shape
            prop_shape = cls._create_property_shape(
                entity_class, field_name, field_info
            )
            property_shapes.append(prop_shape)

        # Create node shape
        node_shape = _NodeShape(
            id=f"{entity_class.__name__}Shape",
            targetClass=Relation(id=entity_class.__name__),
            property=property_shapes,
        )

        return node_shape

    @classmethod
    def _create_property_shape(
        cls, entity_class, field_name: str, field_info
    ) -> _PropertyShape:
        """Create a SHACL property shape for a field.

        Args:
            entity_class: The Entity class containing the field.
            field_name: Name of the field.
            field_info: Pydantic field information.

        Returns:
            _PropertyShape: SHACL property shape for the field.
        """
        # Extract type and metadata
        field_type, metadata = cls._extract_field_type_and_metadata(
            entity_class, field_name, field_info
        )

        # Determine if this is a relation
        is_relation = cls._is_relation_type(field_type)

        # Create base property shape
        prop_shape = _PropertyShape(
            id=f"{entity_class.__name__}Shape_{field_name}",
            path=Relation(id=field_name),
            name=field_name,
            description=field_info.description,
        )

        # Map Python types to XSD datatypes
        datatype_map = {
            str: "xsd:string",
            int: "xsd:integer",
            float: "xsd:decimal",
            bool: "xsd:boolean",
        }

        # Set datatype or class based on field type
        if is_relation:
            prop_shape.nodeKind = Relation(id="sh:IRI")
        else:
            if field_type in datatype_map:
                prop_shape.datatype = Relation(id=datatype_map[field_type])

        # Check if field is required (not Optional)
        if not field_info.is_required():
            # Optional field - no minCount constraint
            pass
        else:
            # Required field has minCount constraint
            prop_shape.minCount = 1

        # Apply SHACL annotations from metadata
        for meta in metadata:
            if isinstance(meta, SHACLAnnotation.DATATYPE):
                prop_shape.datatype = Relation(id=meta.value)
            elif isinstance(meta, SHACLAnnotation.MAX_COUNT):
                prop_shape.maxCount = meta.value
            elif isinstance(meta, SHACLAnnotation.MIN_COUNT):
                prop_shape.minCount = meta.value
            elif isinstance(meta, SHACLAnnotation.PATTERN):
                prop_shape.pattern = meta.value
            elif isinstance(meta, SHACLAnnotation.MIN_LENGTH):
                prop_shape.minLength = meta.value
            elif isinstance(meta, SHACLAnnotation.MAX_LENGTH):
                prop_shape.maxLength = meta.value
            elif isinstance(meta, SHACLAnnotation.MIN_INCLUSIVE):
                prop_shape.minInclusive = meta.value
            elif isinstance(meta, SHACLAnnotation.MAX_INCLUSIVE):
                prop_shape.maxInclusive = meta.value
            elif isinstance(meta, SHACLAnnotation.MIN_EXCLUSIVE):
                prop_shape.minExclusive = meta.value
            elif isinstance(meta, SHACLAnnotation.MAX_EXCLUSIVE):
                prop_shape.maxExclusive = meta.value
            elif isinstance(meta, SHACLAnnotation.NODE_KIND):
                prop_shape.nodeKind = Relation(id=meta.value)
            elif isinstance(meta, SHACLAnnotation.CLASS):
                prop_shape.shclass = Relation(id=meta.value)

        return prop_shape

    @classmethod
    def _extract_field_type_and_metadata(
        cls, entity_class, field_name: str, field_info
    ) -> tuple:
        """Extract the field type and metadata from annotations.

        Args:
            entity_class: The Entity class containing the field.
            field_name: Name of the field.
            field_info: Pydantic field information.

        Returns:
            tuple: Field type and metadata tuple.
        """
        field_type = field_info.annotation
        metadata = ()

        # Get original annotation
        original_annotation = None
        if (
            hasattr(entity_class, "__annotations__")
            and field_name in entity_class.__annotations__
        ):
            original_annotation = entity_class.__annotations__[field_name]

        # Extract metadata from Annotated
        if original_annotation and get_origin(original_annotation) is Annotated:
            args = get_args(original_annotation)
            field_type = args[0]
            metadata = args[1:]

        # Handle Optional types
        origin = getattr(field_type, "__origin__", None)
        if origin is not None:
            args = getattr(field_type, "__args__", ())
            if args:
                field_type = args[0]

        return field_type, metadata

    @classmethod
    def _is_relation_type(cls, field_type) -> bool:
        """Check if a field type is a Relation (ObjectProperty).

        Args:
            field_type: The field type to check.

        Returns:
            bool: True if the field type is a Relation, False otherwise.
        """
        try:
            if isinstance(field_type, type) and issubclass(field_type, Relation):
                return True
        except TypeError:
            pass
        return False

class Relation(BaseModel):
    """This class should be the type of Entity attributes for them to be considered as IRIs.

    Args:
        id (str): IRI of relation
    """

    id: str = Field(alias="@id", title="@id", description="IRI (possibly relative)")
    model_config = ConfigDict(populate_by_name=True, serialize_by_alias=True)

class RDFSAnnotation:
    """Provides methods for setting RDFS annotations.

    Encapsulates dataclasses that validate and provide type information for RDFS annotations.
    These annotations are used in the construction of the ontology graph.
    """

    @dataclass
    class DOMAIN:
        """Dataclass that holds rdfs:domain annotation for a property.

        rdfs:domain is an instance of rdf:Property that is used to state that
        any resource that has a given property is an instance of one or more classes.

        Args:
            value (str): Name of RDFS class
        """

        value: Annotated[str, AfterValidator(val_no_whitespace)]

    @dataclass
    class RANGE:
        """Dataclass that holds rdfs:range annotation for a property.

        rdfs:range is an instance of rdf:Property that is used to state
        that the values of a property are instances of one or more classes.

        Args:
            value (str): Name of RDFS class
        """

        value: Annotated[str, AfterValidator(val_no_whitespace)]

    @staticmethod
    def domain(value: str) -> DOMAIN:
        """RDFSAnnotation.DOMAIN factory

        Args:
            value (str): Name of RDFS class

        Returns:
            RDFSAnnotation.DOMAIN
        """
        return RDFSAnnotation.DOMAIN(value=value)

    @staticmethod
    def range(value: str) -> RANGE:
        """RDFSAnnotation.RANGE factory

        Args:
            value (str): Name of RDFS class

        Returns:
            RDFSAnnotation.RANGE
        """
        return RDFSAnnotation.RANGE(value=value)

class SHACLAnnotation:
    """Provides methods for setting SHACL annotations.

    Encapsulates dataclasses that validate and provide type information for SHACL annotations.
    These annotations are used in the construction of the SHACL graph.
    """

    @dataclass
    class DATATYPE:
        """Dataclass that holds sh:datatype annotation for a property.

        The values of sh:datatype in a shape are IRIs (e.g. xsd:integer).
        A shape has at most one value for sh:datatype.

        Args:
            value (str): IRI of datatype
        """

        value: Annotated[str, AfterValidator(val_datatype)]

    @dataclass
    class MAX_COUNT:
        """Dataclass that holds sh:maxCount annotation for a property.

        The maximum cardinality.
        The values of sh:maxCount in a property shape are literals with datatype xsd:integer.

        Args:
            value (int): Max cardinality of property
        """

        value: Annotated[int, AfterValidator(val_non_negative_int)]

    @dataclass
    class MIN_COUNT:
        """Dataclass that holds sh:minCount annotation for a property.

        The minimum cardinality.
        The values of sh:minCount in a property shape are literals with datatype xsd:integer.

        Args:
            value (int): Min cardinality of property
        """

        value: Annotated[int, AfterValidator(val_non_negative_int)]

    @dataclass
    class PATTERN:
        """Dataclass that holds sh:pattern annotation for a property.

        String-based constraint.
        A regular expression that all value nodes need to match.
        The values of sh:pattern in a shape are valid pattern arguments for the SPARQL REGEX function.

        Args:
            value (str): SPARQL regex (validation not implemented!)
        """

        value: Annotated[str, AfterValidator(val_regex_pattern)]

    @dataclass
    class MIN_LENGTH:
        """Dataclass that holds sh:minLength annotation for a property.

        String-based constraint.
        The minimum length. The values of sh:minLength in a shape are literals with datatype xsd:integer.
        A shape has at most one value for sh:minLength.

        Args:
            value (int): Min length of RDF literal
        """

        value: Annotated[int, AfterValidator(val_non_negative_int)]

    @dataclass
    class MAX_LENGTH:
        """Dataclass that holds sh:maxLength annotation for a property.

        String-based constraint.
        The maximum length. The values of sh:maxLength in a shape are literals with datatype xsd:integer.
        A shape has at most one value for sh:maxLength.

        Args:
            value (int): Max length of RDF literal
        """

        value: Annotated[int, AfterValidator(val_non_negative_int)]

    @dataclass
    class MIN_INCLUSIVE:
        """Dataclass that holds sh:minInclusive annotation for a property.

        Value range constraint.
        The minimum inclusive value. The values of sh:minInclusive in a shape are literals.
        A shape has at most one value for sh:minInclusive.

        Args:
            value (int | float): Min inclusive value of RDF literal
        """

        value: int | float

    @dataclass
    class MAX_INCLUSIVE:
        """Dataclass that holds sh:maxInclusive annotation for a property.

        Value range constraint.
        The maximum inclusive value. The values of sh:maxInclusive in a shape are literals.
        A shape has at most one value for sh:maxInclusive.

        Args:
            value (int | float): Max inclusive value of RDF literal
        """

        value: int | float

    @dataclass
    class MIN_EXCLUSIVE:
        """Dataclass that holds sh:minExclusive annotation for a property.

        Value range constraint.
        The minimum exclusive value. The values of sh:minExclusive in a shape are literals.
        A shape has at most one value for sh:minExclusive.

        Args:
            value (int | float): Min exclusive value of RDF literal
        """

        value: int | float

    @dataclass
    class MAX_EXCLUSIVE:
        """Dataclass that holds sh:maxExclusive annotation for a property.

        Value range constraint.
        The maximum exclusive value. The values of sh:maxExclusive in a shape are literals.
        A shape has at most one value for sh:maxExclusive.

        Args:
            value (int | float): Max exclusive value of RDF literal
        """

        value: int | float

    @dataclass
    class NODE_KIND:
        """Dataclass that holds sh:nodeKind annotation for a property.

        Value type constraint.
        The values of sh:nodeKind in a shape are one of the following six instances of the class sh:NodeKind:
        sh:BlankNode, sh:IRI, sh:Literal sh:BlankNodeOrIRI, sh:BlankNodeOrLiteral and sh:IRIOrLiteral.
        A shape has at most one value for sh:nodeKind.

        Args:
            value (str): Instance of node kind
        """

        value: Annotated[str, AfterValidator(val_node_kind)]

    @dataclass
    class CLASS:
        """Dataclass that holds sh:class annotation for a property.

        Value type constraint.
        The type of all value nodes. The values of sh:class in a shape are IRIs.

        Args:
            value (str): IRI of class
        """

        value: Annotated[str, AfterValidator(val_no_whitespace)]

    @staticmethod
    def datatype(value: str) -> DATATYPE:
        """SHACLAnnotation.DATATYPE factory

        Args:
            value (str): IRI of datatype

        Returns:
            SHACLAnnotation.DATATYPE
        """
        return SHACLAnnotation.DATATYPE(value)

    @staticmethod
    def maxCount(value: int) -> MAX_COUNT:
        """SHACLAnnotation.MAX_COUNT factory

        Args:
            value (int): Max cardinality of property

        Returns:
            SHACLAnnotation.MAX_COUNT
        """
        return SHACLAnnotation.MAX_COUNT(value)

    @staticmethod
    def minCount(value: int) -> MIN_COUNT:
        """SHACLAnnotation.MIN_COUNT factory

        Args:
            value (int): Min cardinality of property

        Returns:
            SHACLAnnotation.MIN_COUNT
        """
        return SHACLAnnotation.MIN_COUNT(value)

    @staticmethod
    def pattern(value: str) -> PATTERN:
        """SHACLAnnotation.PATTERN factory

        Args:
            value (str): SPARQL regex (validation not implemented!)

        Returns:
            SHACLAnnotation.PATTERN
        """
        return SHACLAnnotation.PATTERN(value)

    @staticmethod
    def minLength(value: int) -> MIN_LENGTH:
        """SHACLAnnotation.MIN_LENGTH factory

        Args:
            value (int): Min length of RDF literal

        Returns:
            SHACLAnnotation.MIN_LENGTH
        """
        return SHACLAnnotation.MIN_LENGTH(value)

    @staticmethod
    def maxLength(value: int) -> MAX_LENGTH:
        """SHACLAnnotation.MAX_LENGTH factory

        Args:
            value (int): Max length of RDF literal

        Returns:
            SHACLAnnotation.MAX_LENGTH
        """
        return SHACLAnnotation.MAX_LENGTH(value)

    @staticmethod
    def minInclusive(value: float) -> MIN_INCLUSIVE:
        """SHACLAnnotation.MIN_INCLUSIVE factory

        Args:
            value (int | float): Min inclusive value of RDF literal

        Returns:
            SHACLAnnotation.MIN_INCLUSIVE
        """
        return SHACLAnnotation.MIN_INCLUSIVE(value)

    @staticmethod
    def maxInclusive(value: float) -> MAX_INCLUSIVE:
        """SHACLAnnotation.MAX_INCLUSIVE factory

        Args:
            value (int | float): Max inclusive value of RDF literal

        Returns:
            SHACLAnnotation.MAX_INCLUSIVE
        """
        return SHACLAnnotation.MAX_INCLUSIVE(value)

    @staticmethod
    def minExclusive(value: float) -> MIN_EXCLUSIVE:
        """SHACLAnnotation.MIN_EXCLUSIVE factory

        Args:
            value (int | float): Min exclusive value of RDF literal

        Returns:
            SHACLAnnotation.MIN_EXCLUSIVE
        """
        return SHACLAnnotation.MIN_EXCLUSIVE(value)

    @staticmethod
    def maxExclusive(value: float) -> MAX_EXCLUSIVE:
        """SHACLAnnotation.MAX_EXCLUSIVE factory

        Args:
            value (int | float): Max exclusive value of RDF literal

        Returns:
            SHACLAnnotation.MAX_EXCLUSIVE
        """
        return SHACLAnnotation.MAX_EXCLUSIVE(value)

    @staticmethod
    def nodeKind(value: str) -> NODE_KIND:
        """SHACLAnnotation.NODE_KIND factory

        Args:
            value (str): Instance of node kind

        Returns:
            SHACLAnnotation.NODE_KIND
        """
        return SHACLAnnotation.NODE_KIND(value)

    @staticmethod
    def shclass(value: str) -> CLASS:
        """SHACLAnnotation.CLASS factory

        Args:
            value (str): IRI of class

        Returns:
            SHACLAnnotation.CLASS
        """
        return SHACLAnnotation.CLASS(value)