from typing import Dict, Iterable, Optional, Union

from pydantic import BaseModel

from structurizr2csv.enums import Terminology
from structurizr2csv.utils import camel_case
from structurizr2csv.workspace import (
    Component,
    ComponentView,
    Container,
    ContainerView,
    DeploymentView,
    DynamicView,
    ElementView,
    Location,
    Model,
    Person,
    RelationshipView,
    SoftwareSystem,
    SystemContextView,
    SystemLandscapeView,
    Views,
    Workspace,
)

View = Union[
    SystemLandscapeView,
    SystemContextView,
    ContainerView,
    ComponentView,
    DynamicView,
    DeploymentView,
]
Element = Union[Person, SoftwareSystem, Container, Component]


class WorkspaceWrapper:
    """Extends the structurizr's workspace parser to add convenient features.

    The file `workspace.py` is automatically generated by `datamodel-codegen`.
    To ease further updates, we don't want to modify that file by hand, so we
    use this class to dynamically add some attributes and methods.
    """

    def __init__(self, workspace: Workspace):

        self._workspace = workspace

        if self.model:
            # this is how you dynamically add attributes to a pydantic model
            object.__setattr__(
                self.model, "object_ids", self.retrieve_model_object_ids()
            )

            for view in self.iterate_views():
                for element in view.elements or []:
                    self.extend_element(element)
                for relationship in view.relationships or []:
                    self.extend_element(relationship)

    @property
    def model(self) -> Optional[Model]:
        return self._workspace.model

    @property
    def views(self) -> Optional[Views]:
        return self._workspace.views

    def extend_element(self, elt: Union[ElementView, RelationshipView]):
        """Do a lot of dirty stuff to add attributes to pydantic model.

        AFAIK there's no cleaner way to extend a pydantic model that has already been
        defined. Subclassing etc. doesn't work as with other "normal" classes.
        """
        assert self.model, "Workspace's model is empty - no elements to extend"  # nosec

        definition = self.model.object_ids[elt.id]  # type: ignore

        # this is how you dynamically add attributes to a pydantic model...
        object.__setattr__(
            definition,
            "c4_type",
            Terminology[camel_case(definition.__class__.__name__)],
        )
        object.__setattr__(
            definition,
            "safe_description",
            (
                '"{}"'.format(definition.description.replace('"', '\\"'))
                if definition.description
                else ""
            ),
        )

        if isinstance(definition, (Person, SoftwareSystem)):
            object.__setattr__(
                definition,
                "is_internal",
                # compare with values, because the auto-generated workspace might
                # define other enums called "Location1", "Location2"...
                definition.location.value == Location.Internal.value
                if definition.location
                else False,
            )

        object.__setattr__(elt, "definition", definition)
        object.__setattr__(
            elt, "c4_type", definition.c4_type  # pyright: reportGeneralTypeIssues=false
        )

    def retrieve_model_object_ids(self) -> Dict:
        """Recursively retrieve all model objects' IDs."""
        ids = {}

        def map_objects(container: BaseModel, attrname: str):
            for obj in getattr(container, attrname, []) or []:

                ids.setdefault(obj.id, obj)

                map_objects(obj, "relationships")
                map_objects(obj, "containers")
                map_objects(obj, "components")

        map_objects(self.model, "people")
        map_objects(self.model, "softwareSystems")

        return ids

    def iterate_views(self) -> Iterable[View]:
        """Yield all views for this workspace."""
        for view_type in (
            attr
            for attr in self.views.__dict__.keys()
            if (attr.endswith("Views") or attr == "Persons")
        ):
            views = getattr(self.views, view_type)
            if views:
                for view in getattr(self.views, view_type):
                    yield view
