"""
File consists of several classes for the different elements of a device.
"""
from __future__ import annotations
from enum import Enum, auto
import uuid
from typing import List, Dict, Optional
import json


class AggregationLayer(Enum):
    SYSTEM = auto()
    ASSEMBLY = auto()
    SUBASSEMBLY = auto()
    COMPONENT = auto()


class LegoComponent:
    def __init__(
        self,
        label: Optional[str] = None,
        datasheet: Optional[dict] = None,
        *more_properties: dict,
        **kwargs,
    ) -> None:
        self.uuid: uuid.UUID = uuid.uuid4()
        self.parent: None | LegoAssembly = None
        self.layer: AggregationLayer = AggregationLayer.COMPONENT
        self.properties: dict = {}
        if label is not None:
            self.properties["label"] = label
        if datasheet is not None:
            self.properties.update(datasheet)
        for prop in more_properties:
            if isinstance(prop, dict):
                self.properties.update(prop)
            else:
                raise ValueError(f"Unexpected argument type: {type(more_properties)}")
        for key, value in kwargs.items():
            self.properties[key] = value

    def clone(self, new_label: Optional[str] = None) -> LegoComponent:
        if new_label is None:
            new_label = self.properties["label"]
        clone = LegoComponent(None, None, self.properties)
        clone.properties["label"] = new_label
        return clone

    def get_root_assembly(self):
        if self.parent is None:
            return None
        current_assembly = self.parent
        while current_assembly.parent is not None:
            current_assembly = current_assembly.parent
        return current_assembly

    def to_dict(self) -> Dict:
        dict_ = {
            "uuid": self.uuid,
            "properties": self.properties,
            "layer": self.layer,
        }
        return {"component": dict_}

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return f"LegoComponent {self.properties['label']} [{self.uuid}]"


class LegoAssembly:
    def __init__(
        self,
        layer: AggregationLayer,
        label: Optional[str] = None,
        *properties: dict,
        **kwargs,
    ) -> None:
        self.uuid: uuid.UUID = uuid.uuid4()
        self.parent: None | LegoAssembly = None
        self.properties: dict = {}
        self.layer: AggregationLayer = layer
        if label is not None:
            self.properties["label"] = label
        for prop in properties:
            if isinstance(prop, dict):
                self.properties.update(prop)
            else:
                raise ValueError(f"Unexpected argument type: {type(properties)}")
        for key, value in kwargs.items():
            self.properties[key] = value
        self.components: List[LegoComponent] = []
        self.assemblies: List[LegoAssembly] = []

    def add_component(self, component: LegoComponent | List[LegoComponent]) -> None:
        if isinstance(component, list):
            for c in component:
                self.add_component(c)
            return

        if not isinstance(component, LegoComponent):
            raise TypeError(
                f"Argument should be of type {LegoComponent.__name__}, "
                f"got {type(component).__name__} instead."
            )

        if self.get_root_assembly().contains_uuid(component.uuid):
            raise AssertionError(
                f"This assembly or a subassembly already contains the component with ID "
                f"{component.uuid}."
            )
        component.parent = self
        self.components.append(component)

    def add_assembly(self, assembly: LegoAssembly | List[LegoAssembly]) -> None:
        if isinstance(assembly, list):
            for a in assembly:
                self.add_assembly(a)
            return

        if not isinstance(assembly, LegoAssembly):
            raise TypeError(
                f"Argument should be of type {LegoAssembly.__name__}, "
                f"got {type(assembly).__name__} instead."
            )

        if self.get_root_assembly().contains_uuid(assembly.uuid):
            raise AssertionError(
                f"This assembly or a subassembly already contains the assembly with ID "
                f"{assembly.uuid}."
            )
        assembly.parent = self
        self.assemblies.append(assembly)

    def add(
        self, part: LegoAssembly | LegoComponent | List[LegoAssembly | LegoComponent]
    ) -> None:
        if isinstance(part, LegoComponent):
            self.add_component(part)
        elif isinstance(part, LegoAssembly):
            self.add_assembly(part)
        elif isinstance(part, list):
            for p in part:
                self.add(p)
        else:
            raise TypeError(
                f"Argument should be of types {LegoAssembly.__name__}, {LegoComponent.__name__} or a list of them, "
                f"got {type(part).__name__} instead."
            )

    def children(self) -> Dict[str, List[LegoComponent] | List[LegoAssembly]]:
        return {"components": self.components, "assemblies": self.assemblies}

    def get_component_list(self, max_depth: int = -1) -> List[LegoComponent]:
        component_list = []
        component_list.extend(self.components)
        if max_depth > 0:
            for assembly in self.assemblies:
                component_list.extend(assembly.get_component_list(max_depth - 1))
        return component_list

    def get_root_assembly(self) -> LegoAssembly:
        current_assembly = self
        while current_assembly.parent is not None:
            current_assembly = current_assembly.parent
        return current_assembly

    def contains_uuid(self, uuid_: uuid.UUID):
        # check component ids
        component_ids = list(map(lambda c: c.uuid, self.components))
        if uuid_ in component_ids:
            return True
        # check assembly ids
        assembly_ids = list(map(lambda a: a.uuid, self.assemblies))
        if uuid_ in assembly_ids:
            return True
        # recursively check assemblies
        for assembly in self.assemblies:
            if assembly.contains_uuid(uuid_):
                return True
        return False

    def to_dict(self) -> Dict:
        dict_ = {
            "uuid": self.uuid,
            "properties": self.properties,
            "layer": self.layer,
        }
        # store components
        dict_["components"] = [component.to_dict() for component in self.components]
        dict_["assemblies"] = [assembly.to_dict() for assembly in self.assemblies]
        return {"assembly": dict_}

    def __repr__(self):
        return f"LegoAssembly {self.properties['label']} [{self.uuid}]"

    def clone(self, label: str = None) -> LegoAssembly:
        if label is None:
            label = self.label
        clone = LegoAssembly(label, copy.deepcopy(self.properties), self.layer)
        for component in self.components:
            clone.add_component(component.clone())
        for assembly in self.assemblies:
            clone.add_assembly(assembly.clone())
        return clone


def print_assembly_tree(root, level=0, is_last=False):
    if not isinstance(root, LegoAssembly):
        raise TypeError(
            f"Argument should be of type {LegoAssembly.__name__}, "
            f"got {type(root).__name__} instead."
        )
    # print component
    assembly_padding = ""
    if level > 0:
        assembly_padding += "│   " * (level - 1)
        if is_last:
            assembly_padding += "└── "
        else:
            assembly_padding += "├── "
    print(f"{assembly_padding}{root}")
    # recursively print child components
    for i, assembly in enumerate(root.assemblies):
        is_last_ = i == len(root.assemblies) - 1 and len(root.components) == 0
        print_assembly_tree(assembly, level + 1, is_last_)
    # print items
    for i, item in enumerate(root.components):
        component_padding = "│   " * level if not is_last else "    "
        component_padding += "├── " if i < len(root.components) - 1 else "└── "
        print(f"{component_padding}{item}")


def correct_aggregation_hierarchy(root: LegoAssembly, strict: bool = False):
    if not isinstance(root, LegoAssembly):
        raise TypeError(
            f"Argument should be of type {LegoAssembly.__name__}, "
            f"got {type(root).__name__} instead."
        )
    higher_level = operator.le
    if strict:
        higher_level = operator.lt
    for component in root.components:
        if not higher_level(root.layer.value, component.layer.value):
            return False
    for assembly in root.assemblies:
        if not higher_level(root.layer.value, assembly.layer.value):
            return False
        if not correct_aggregation_hierarchy(assembly, strict):
            return False
    return True


class KPIEncoder(json.JSONEncoder):
    def default(self, o):
        if isinstance(o, uuid.UUID):
            return "kpi-" + str(o)
        if isinstance(o, (AggregationLayer)):
            return "kpi-" + o.name
        return super().default(o)