Core

class BaseDiagramElement:
    """
    Base class for any object that belongs to a diagram.

    Provides access to the current diagram context and allows
    attaching structured properties (e.g. key-value tables).
    """

    allowed_diagram_types: tuple[DiagramType, ...] | None = None

    _diagram: Diagram

    def __init__(self, **kwargs: Any) -> None:
        """
        Initializes the element and adds it to the current diagram context.
        """
        self._diagram = current_diagram()
        self._contribute_to_diagram()
        self.properties = DiagramElementProperties()

    def _check_diagram_type(self) -> None:
        if not self.allowed_diagram_types:
            return None

        if self._diagram.type not in self.allowed_diagram_types:
            element_name = self.__class__.__name__
            diagram_type = self._diagram.type.value
            allowed = ", ".join([dt.value for dt in self.allowed_diagram_types])

            raise ValueError(
                f"{element_name} is not allowed in {diagram_type}. "
                f"Allowed diagram types: {allowed}."
            )

        return None

    def _contribute_to_diagram(self) -> None:
        self._check_diagram_type()
        self._diagram.add_base_element(self)

    def set_property_header(self, *args: str) -> Self:
        """
        Sets the column headers for the element's property table.

        This must be called either before adding any property rows, or
        the header length must match the number of values.

        Args:
            *args: Column names to use as the property header.

        Returns:
            The updated diagram element.

        Raises:
            ValueError: If header length does not match the number of values.
        """
        if not args:
            raise ValueError("The header cannot be empty")

        if self.properties.properties:
            expected_header_length = self.properties.properties[0]

            if len(args) != len(expected_header_length):
                raise ValueError(
                    "The header length does not match the number of values"
                )

        self.properties.header = list(args)

        return self

    def without_property_header(self) -> Self:
        """
        Disables the rendering of the header row in the property table.

        Returns:
            The updated diagram element.
        """
        self.properties.show_header = False

        return self

    def add_property(self, *args: str) -> Self:
        """
        Adds a row to the property table.

        The number of arguments must match the number of header columns.

        Args:
            *args: Values for each column in the property row.

        Returns:
            The updated diagram element.

        Raises:
            ValueError: If the number of values does not match the
                header length.
        """
        if len(args) != len(self.properties.header):
            raise ValueError(
                "The number of values does not match the header length"
            )

        self.properties.properties.append(list(args))

        return self

    @property
    def diagram(self) -> Diagram:
        """Returns the current diagram context."""
        return self._diagram

class Element(BaseDiagramElement, abc.ABC):
    """
    Base class for all C4 elements (e.g. Person, System, Container, Component).

    Elements are automatically registered in the current diagram context.
    """

    allowed_diagram_types: tuple[DiagramType, ...] | None = None

    _diagram: Diagram

    alias: str
    label: str
    tags: list[str] | None
    base_shape: str | None
    technology: str | None

    def __init__(
        self,
        label: Required[str] = REQUIRED,
        description: str | None = None,
        sprite: str | None = None,
        tags: list[str] | None = None,
        link: str | None = None,
        type_: str | None = None,
        alias: Maybe[str] = MISSING,
    ) -> None:
        """
        Initialize a new diagram element. Automatically adds the element to the
        current diagram.

        Args:
            label: Display name for the element. Required.
            description: Optional description text.
            sprite: Optional sprite/icon reference for rendering.
            tags: Optional tags for styling or grouping.
            link: Optional URL associated with the element.
            type_: Optional custom type or stereotype label.
            alias: Unique identifier for the element. If not provided, it is
                autogenerated from the label.

        Raises:
            ValueError: If `label` is not provided.
        """
        self.label = self._check_label(label)
        self.alias = self._check_alias(alias, self.label)
        self.sprite = sprite
        self.type = type_
        self.tags = tags
        self.link = link

        self.description = description
        self.base_shape = None
        self.technology = None

        super().__init__()

    @overload
    def __rshift__(self, other: str) -> Relationship: ...

    @overload
    def __rshift__(self, other: Element) -> _EdgeDraft: ...

    def __rshift__(self, other: str | Element) -> Relationship | _EdgeDraft:
        """
        Enables:
          - Self >> "label" >> Element2   (pending relationship)
          - Self >> Element2 | "label"    (draft for later '| "label"')
        """
        if isinstance(other, str):
            # self >> "label" >> element2
            return Relationship(label=other, from_element=self)

        if isinstance(other, Element):
            # Draft for: self >> element2 | "label"
            return _EdgeDraft(source=self, destination=other)

        return NotImplemented

    @overload
    def __lshift__(self, other: str) -> Relationship: ...

    @overload
    def __lshift__(self, other: Element) -> _EdgeDraft: ...

    def __lshift__(self, other: str | Element) -> Relationship | _EdgeDraft:
        """
        Enables:
          - Element2 << "label" << Self   (pending relationship)
          - Element2 << self | "label"      (draft for later '| "label"')
        """
        if isinstance(other, str):
            # element1 << "label" << element2
            return Relationship(label=other, to_element=self)

        if isinstance(other, Element):
            # Draft for: element1 >> element2 | "label"
            return _EdgeDraft(source=other, destination=self)

        return NotImplemented

    def __rrshift__(self, other: list[Relationship]) -> list[Relationship]:
        """
        Enables: [Relationship(...), ...] >> Element.
        """
        if isinstance(other, list) and all(
            isinstance(r, Relationship) for r in other
        ):
            return [r._connect(r.from_element, destination=self) for r in other]

        return NotImplemented  # pragma: no cover

    def __rlshift__(self, other: list[Relationship]) -> list[Relationship]:
        """
        Enables: Element << [Relationship(...), ...].
        """
        if isinstance(other, list) and all(
            isinstance(r, Relationship) for r in other
        ):
            return [
                r._connect(source=self, destination=r.to_element) for r in other
            ]

        return NotImplemented  # pragma: no cover

    def _check_label(self, label: str | Required) -> str:
        if label is REQUIRED:
            raise ValueError("The 'label' argument is required")

        return cast(str, label)

    def _check_alias(self, alias: Maybe[str], label: str) -> str:
        if alias is MISSING:
            alias = self._generate_alias(label)

        return cast(str, alias)

    @override
    def _contribute_to_diagram(self) -> None:
        self._check_diagram_type()
        self._diagram.add(self)

    def uses(
        self,
        other: _TElement,
        label: str,
        relationship_type: RelationshipType = RelationshipType.REL,
        **kwargs: Any,
    ) -> Relationship:
        """
        Declare that this element uses another.

        Args:
            other: The element being used.
            label: Description of the interaction.
            relationship_type: Type of arrow to use.
            kwargs: Optional relationship kwargs.

        Returns:
            The created relationship.
        """
        relationship_class = Relationship.get_relationship_by_type(
            relationship_type
        )
        return relationship_class(
            from_element=self,  # type: ignore[arg-type]
            to_element=other,  # type: ignore[arg-type]
            label=label,
            **kwargs,
        )

    def used_by(
        self,
        other: _TElement,
        label: str,
        relationship_type: RelationshipType = RelationshipType.REL,
        **kwargs: Any,
    ) -> Relationship:
        """
        Declare that another element uses this element.

        Args:
            other: The element that uses this element.
            label: Description of the interaction.
            relationship_type: Type of arrow to use.
            kwargs: Optional relationship kwargs.

        Returns:
            The created relationship.
        """
        relationship_class = Relationship.get_relationship_by_type(
            relationship_type
        )
        return relationship_class(
            from_element=other,  # type: ignore[arg-type]
            to_element=self,  # type: ignore[arg-type]
            label=label,
            **kwargs,
        )

    def _generate_alias(self, label: str) -> str:
        return current_diagram().generate_alias(label=label)

    @override
    def __str__(self) -> str:
        """Returns the string representation of the element."""
        cls_name = self.__class__.__name__
        return f"{cls_name}(alias={self.alias!r}, label={self.label!r})"

    def __repr__(self) -> str:
        cls_name = self.__class__.__name__
        attrs = [
            f"{self.label!r}",
        ]

        if self.description:
            attrs.append(f"{self.description!r}")

        if self.sprite:
            attrs.append(f"sprite={self.sprite!r}")

        if self.type:
            attrs.append(f"type_={self.type!r}")

        if self.tags:
            attrs.append(f"tags={self.tags!r}")

        if self.link:
            attrs.append(f"link={self.link!r}")

        if self.technology:
            attrs.append(f"technology={self.technology!r}")

        if self.base_shape:
            attrs.append(f"base_shape={self.base_shape!r}")

        attrs.append(f"alias={self.alias!r}")

        args = ", ".join(attrs)
        return f"{cls_name}({args})"

class ElementWithTechnology(Element):
    """
    Base class for elements that define a `technology` field.
    """

    def __init__(
        self,
        label: Required[str] = REQUIRED,
        description: str | None = None,
        technology: str | None = None,
        sprite: str | None = None,
        tags: list[str] | None = None,
        link: str | None = None,
        alias: Maybe[str] = MISSING,
    ) -> None:
        """
        Initialize a new diagram element.

        Args:
            label: Display name for the element. Required.
            description: Optional description text.
            technology: Optional technology.
            sprite: Optional sprite/icon reference for rendering.
            tags: Optional tags for styling or grouping.
            link: Optional URL associated with the element.
            alias: Unique identifier for the element. If not provided, it is
                autogenerated from the label.
        """
        super().__init__(
            alias=alias,
            label=label,
            description=description,
            sprite=sprite,
            tags=tags,
            link=link,
        )

        self.technology = technology

class Boundary(Element):
    """
    Represents a boundary element that groups other elements.

    Boundaries can be nested, and manage their own child elements.
    """

    def __init__(
        self,
        label: Required[str] = REQUIRED,
        description: str | None = None,
        type_: str | None = None,
        tags: list[str] | None = None,
        link: str | None = None,
        alias: Maybe[str] = MISSING,
    ) -> None:
        """
        Initialize a new boundary element.

        Args:
            label: Human-readable name for the boundary. Required.
            description: Optional description.
            type_: Optional stereotype or visual marker.
            tags: Optional tags for styling or grouping.
            link: Optional hyperlink associated with the boundary.
            alias: Unique identifier for the boundary.
                If not provided, one is autogenerated.

        Notes:
            - If the boundary is created within another boundary context, it is
              added as a nested boundary.
            - Otherwise, it is added directly to the current diagram.
        """
        self._parent = get_boundary()

        super().__init__(
            label=label,
            alias=alias,
            description=description,
            type_=type_,
            tags=tags,
            link=link,
        )

        self._elements: list[Element] = []
        self._relationships: list[Relationship] = []
        self._boundaries: list[Boundary] = []

        self.__ordered_elements: list[BaseDiagramElement] = []

    @override
    def _contribute_to_diagram(self) -> None:
        self._check_diagram_type()
        self._diagram.add_boundary(self)

    @property
    def elements(self) -> list[Element]:
        """
        Returns the list of diagram elements added to this boundary.

        Returns:
            Child elements grouped under this boundary.
        """
        return self._elements

    @property
    def boundaries(self) -> list[Boundary]:
        """
        Returns the list of nested boundaries inside this boundary.

        Returns:
            Child boundaries nested within this boundary.
        """
        return self._boundaries

    @property
    def ordered_elements(self) -> list[BaseDiagramElement]:
        """
        Return the boundary elements in their order of definition,
        preserving the original declaration order for rendering.
        """
        return self.__ordered_elements

    @property
    def relationships(self) -> list[Relationship]:
        """
        Returns all relationships defined in the boundary.
        """
        return self._relationships

    def __enter__(self) -> Self:
        """
        Enter the boundary context.

        Returns:
            The boundary instance now active as context.
        """
        set_boundary(self)
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        """
        Exit the boundary context and restore the previous boundary.
        """
        set_boundary(self._parent)

    def add(self, element: _TElement) -> _TElement:
        """
        Add a diagram element to this boundary.

        Args:
            element: The element to add.

        Returns:
            The added element.
        """
        self._elements.append(element)
        self.__ordered_elements.append(element)

        return element

    def add_boundary(self, boundary: _TBoundary) -> _TBoundary:
        """
        Add a nested boundary to this boundary.

        Args:
            boundary: The boundary to add.

        Returns:
            The added boundary.
        """
        self._boundaries.append(boundary)
        self.__ordered_elements.append(boundary)

        return boundary

    def add_relationship(self, relationship: _TRelationship) -> _TRelationship:
        """
        Add a relationship between elements.

        Args:
            relationship: The relationship to add.

        Returns:
            The added relationship.
        """
        self._relationships.append(relationship)
        self.__ordered_elements.append(relationship)

        return relationship

@unique
class RelationshipType(EnumDescriptionsMixin, StrEnum):
    """
    Enum representing different types of relationships between
    diagram elements.
    """

    REL = "REL"
    BI_REL = "BI_REL"
    REL_BACK = "REL_BACK"
    REL_NEIGHBOR = "REL_NEIGHBOR"
    BI_REL_NEIGHBOR = "BI_REL_NEIGHBOR"
    REL_BACK_NEIGHBOR = "REL_BACK_NEIGHBOR"
    REL_D = "REL_D"
    REL_DOWN = "REL_DOWN"
    BI_REL_D = "BI_REL_D"
    BI_REL_DOWN = "BI_REL_DOWN"
    REL_U = "REL_U"
    REL_UP = "REL_UP"
    BI_REL_U = "BI_REL_U"
    BI_REL_UP = "BI_REL_UP"
    REL_L = "REL_L"
    REL_LEFT = "REL_LEFT"
    BI_REL_L = "BI_REL_L"
    BI_REL_LEFT = "BI_REL_LEFT"
    REL_R = "REL_R"
    REL_RIGHT = "REL_RIGHT"
    BI_REL_R = "BI_REL_R"
    BI_REL_RIGHT = "BI_REL_RIGHT"

    @classmethod
    def get_descriptions(cls) -> dict[RelationshipType, str]:
        """Return the Enum items description used in documentation."""
        return {
            cls.BI_REL: "A bidirectional relationship between two elements.",
            cls.BI_REL_DOWN: "A bidirectional downward relationship.",
            cls.BI_REL_D: (
                "A bidirectional downward relationship. "
                "Shorthand for `BI_REL_DOWN`."
            ),
            cls.BI_REL_LEFT: "A bidirectional leftward relationship.",
            cls.BI_REL_L: (
                "A bidirectional leftward relationship. "
                "Shorthand for `BI_REL_LEFT`."
            ),
            cls.BI_REL_NEIGHBOR: (
                "A bidirectional neighboring relationship between two elements."
            ),
            cls.BI_REL_RIGHT: "A bidirectional rightward relationship.",
            cls.BI_REL_R: (
                "A bidirectional rightward relationship. "
                "Shorthand for `BI_REL_RIGHT`."
            ),
            cls.BI_REL_UP: "A bidirectional upward relationship.",
            cls.BI_REL_U: (
                "A bidirectional upward relationship. "
                "Shorthand for `BI_REL_UP`."
            ),
            cls.REL: "A unidirectional relationship between two elements.",
            cls.REL_BACK: "A unidirectional relationship pointing backward.",
            cls.REL_BACK_NEIGHBOR: (
                "A unidirectional relationship combining backward "
                "and neighboring semantics."
            ),
            cls.REL_DOWN: "A unidirectional downward relationship.",
            cls.REL_D: (
                "A unidirectional downward relationship. "
                "Shorthand for `REL_DOWN`."
            ),
            cls.REL_LEFT: "A unidirectional leftward relationship.",
            cls.REL_L: (
                "A unidirectional leftward relationship. "
                "Shorthand for `REL_LEFT`."
            ),
            cls.REL_NEIGHBOR: (
                "A unidirectional relationship representing a lateral "
                "or neighboring interaction."
            ),
            cls.REL_RIGHT: "A unidirectional rightward relationship.",
            cls.REL_R: (
                "A unidirectional rightward relationship. "
                "Shorthand for `REL_RIGHT`."
            ),
            cls.REL_UP: "A unidirectional upward relationship.",
            cls.REL_U: (
                "A unidirectional upward relationship. Shorthand for `REL_UP`."
            ),
        }

class Relationship(BaseDiagramElement):
    """
    Represents a connection between two elements.

    Supports fluent chaining using `>>` and `<<` operators, and
    subclass registration for each `RelationshipType`.
    """

    __relationship_by_type: ClassVar[
        dict[RelationshipType, type[Relationship]]
    ] = {}

    relationship_type: RelationshipType = RelationshipType.REL

    def __init__(
        self,
        label: str | None = None,
        description: str | None = None,
        technology: str | None = None,
        sprite: str | None = None,
        tags: list[str] | None = None,
        link: str | None = None,
        index: str | BaseIndex | None = None,
        from_element: _TElement | None = None,
        to_element: _TElement | None = None,
        relationship_type: RelationshipType | None = None,
    ) -> None:
        """
        Initialize a relationship between two elements.

        Args:
            label: The label shown on the relationship edge.
            description: Additional details about the relationship.
            technology: The technology used in the communication.
            sprite: Optional sprite to represent the relationship.
            tags: Optional tags for styling or grouping.
            link: URL link associated with the relationship.
            index: Index associated with the relationship.
            from_element: The source element. Optional.
            to_element: The destination element. Optional.
            relationship_type: Type of the relationship.
                Defaults to the class-level `relationship_type`.

        Notes:
            If both `from_element` and `to_element` are provided,
            the relationship will be registered in the current
            diagram immediately.
        """
        self.from_element = from_element
        self.to_element = to_element
        self.label = label
        self.technology = technology
        self.description = description
        self.sprite = sprite
        self.tags = tags
        self.link = link
        self.index = index

        self.relationship_type = relationship_type or self.relationship_type

        super().__init__()

    @override
    def __init_subclass__(cls, *args: Any, **kwargs: Any) -> None:
        """
        Registers the relationship subclass under its unique
        `relationship_type`.
        """
        super().__init_subclass__(*args, **kwargs)

        relationship_type = getattr(cls, "relationship_type", None)
        if (
            relationship_type is None
            or relationship_type in cls.__relationship_by_type
        ):
            raise TypeError(
                f"Please provide an unique `relationship_type` for this"
                f" class {cls.__name__}"
            )

        cls.__relationship_by_type[relationship_type] = cls

    def get_participants(self) -> tuple[_TElement, _TElement]:
        if not self.from_element:
            raise ValueError("from_element not provided")

        if not self.to_element:
            raise ValueError("to_element not provided")

        return self.from_element, self.to_element  # type: ignore[return-value]

    @overload
    def __rshift__(
        self, other: _TElement
    ) -> Relationship: ...  # pragma: no cover

    @overload
    def __rshift__(
        self, other: list[_TElement]
    ) -> list[Relationship]: ...  # pragma: no cover

    def __rshift__(
        self, other: _TElement | list[_TElement]
    ) -> Relationship | list[Relationship]:
        """Implements Self >> _TElement and Self >> [_TElement]."""
        self._ensure_not_completed()

        return self._connect(source=self.from_element, destination=other)  # type: ignore[arg-type,type-var]

    @overload
    def __lshift__(
        self, other: _TElement
    ) -> Relationship: ...  # pragma: no cover

    @overload
    def __lshift__(
        self, other: list[_TElement]
    ) -> list[Relationship]: ...  # pragma: no cover

    def __lshift__(
        self, other: _TElement | list[_TElement]
    ) -> Relationship | list[Relationship]:
        """Implements Self << _TElement and Self << [_TElement]."""
        self._ensure_not_completed()

        return self._connect(source=other, destination=self.to_element)  # type: ignore[arg-type,type-var]

    @overload
    def __rrshift__(
        self, other: _TElement
    ) -> Relationship: ...  # pragma: no cover

    @overload
    def __rrshift__(
        self, other: list[_TElement]
    ) -> list[Relationship]: ...  # pragma: no cover

    def __rrshift__(
        self, other: _TElement | list[_TElement]
    ) -> Relationship | list[Relationship]:
        """Called for [_TElement] >> Self or _TElement >> Self."""
        self._ensure_not_completed()

        return self._connect(source=other, destination=self.to_element)  # type: ignore[arg-type,type-var]

    @overload
    def __rlshift__(
        self, other: _TElement
    ) -> Relationship: ...  # pragma: no cover

    @overload
    def __rlshift__(
        self, other: list[_TElement]
    ) -> list[Relationship]: ...  # pragma: no cover

    def __rlshift__(
        self, other: _TElement | list[_TElement]
    ) -> Relationship | list[Relationship]:
        """Called for [_TElement] << Self or _TElement << Self."""
        self._ensure_not_completed()

        return self._connect(source=self.from_element, destination=other)  # type: ignore[arg-type,type-var]

    def __repr__(self) -> str:
        cls_name = self.__class__.__name__
        attrs = [
            f"{self.label!r}",
        ]

        if self.description:
            attrs.append(f"{self.description!r}")

        repr_attrs = [
            "technology",
            "sprite",
            "tags",
            "link",
            "index",
        ]

        for attr in repr_attrs:
            value = getattr(self, attr)
            if value:
                attrs.append(f"{attr}={value!r}")

        args = ", ".join(attrs)
        return f"{cls_name}({args})"

    @overload
    def _connect(
        self, source: None, destination: None
    ) -> NoReturn: ...  # pragma: no cover

    @overload
    def _connect(
        self, source: _TElement, destination: None
    ) -> NoReturn: ...  # pragma: no cover

    @overload
    def _connect(
        self, source: None, destination: _TElement
    ) -> NoReturn: ...  # pragma: no cover

    @overload
    def _connect(
        self, source: _TElement, destination: _TElement
    ) -> Relationship: ...  # pragma: no cover

    @overload
    def _connect(
        self, source: list[_TElement], destination: None
    ) -> NoReturn: ...  # pragma: no cover

    @overload
    def _connect(
        self, source: None, destination: list[_TElement]
    ) -> NoReturn: ...  # pragma: no cover

    @overload
    def _connect(
        self, source: list[_TElement], destination: list[_TElement]
    ) -> NoReturn: ...  # pragma: no cover

    @overload
    def _connect(
        self, source: list[_TElement], destination: _TElement
    ) -> list[Relationship]: ...  # pragma: no cover

    @overload
    def _connect(
        self, source: _TElement, destination: list[_TElement]
    ) -> list[Relationship]: ...  # pragma: no cover

    def _connect(
        self,
        source: _TElement | list[_TElement] | None,
        destination: _TElement | list[_TElement] | None,
    ) -> Relationship | list[Relationship]:
        self._ensure_not_completed()

        if not source and not destination:
            raise ValueError("Either source or destination must be provided")

        if isinstance(source, list) and isinstance(destination, list):
            raise ValueError(  # noqa: TRY004
                "Either source or destination must be a single element"
            )

        if isinstance(source, list):
            from_iter = source
            to_iter: Iterable[_TElement] = repeat(destination)  # type: ignore[arg-type]
        elif isinstance(destination, list):
            from_iter: Iterable[_TElement] = repeat(source)  # type: ignore[no-redef]
            to_iter = destination
        else:
            # Both are single elements
            return self.copy(from_element=source, to_element=destination)

        return [
            self.copy(from_element=src, to_element=dst)
            for src, dst in zip(from_iter, to_iter, strict=False)
        ]

    def _ensure_not_completed(self) -> None:
        if self.from_element and self.to_element:
            raise ValueError(
                "Cannot modify relationship with both specified elements"
            )

    @override
    def _contribute_to_diagram(self) -> None:
        if self.from_element and self.to_element:
            self._diagram.add_relationship(self)

    def get_attrs(self) -> dict[str, Any]:
        """
        Returns a dictionary of all relationship attributes.
        """
        return {
            "from_element": self.from_element,
            "to_element": self.to_element,
            "label": self.label,
            "technology": self.technology,
            "description": self.description,
            "sprite": self.sprite,
            "tags": self.tags,
            "link": self.link,
            "index": self.index,
            "relationship_type": self.relationship_type,
        }

    def copy(self, **overrides: Any) -> Relationship:
        """
        Clone the relationship, optionally overriding fields.
        """
        attrs = {**self.get_attrs(), **overrides}

        cls = self.get_relationship_by_type(self.relationship_type)

        relationship_copy = cls(**attrs)

        if self.properties.properties:
            relationship_copy.properties = copy.deepcopy(self.properties)

        return relationship_copy

    @classmethod
    def get_relationship_by_type(
        cls, relationship_type: RelationshipType
    ) -> type[Relationship]:
        """
        Retrieve the relationship class associated with the
        given RelationshipType.

        Args:
            relationship_type: The enum value representing the
                type of relationship.

        Returns:
            The corresponding Relationship subclass.

        Raises:
            KeyError: If no class is registered for the provided
                relationship type.
        """
        return cls.__relationship_by_type[relationship_type]