from enum import Enum
from kennel.engine.components.component import Component, ComponentType
from typing import List, Optional


class EntityType(str, Enum):
    LASER = "LASER"


class Entity:
    def __init__(
        self, entity_type: EntityType, id: str, component_list: List[Component]
    ):
        self.entity_type = entity_type
        self.id = id
        self.components = {}
        for component in component_list:
            self.add_component(component)

    def get_component(self, component_type: ComponentType) -> Optional[Component]:
        return self.components[component_type]

    def add_component(self, component: Component) -> None:
        self.components[component.component_type] = component

    def to_dict(self) -> dict:
        return {
            "entity_type": self.entity_type,
            "id": self.id,
            "components": {k: v.dict() for k, v in self.components.items()},
        }


class EntityManager:
    def __init__(self):
        self.entities = {}
        self.component_entities = {}

    def update(self) -> None:
        self.component_entities = {}
        for entity in self.entities.values():
            for component in entity.components.values():
                if component.component_type not in self.component_entities:
                    self.component_entities[component.component_type] = set()
                self.component_entities[component.component_type].add(entity)

    def get_entities_with_component(
        self, component_type: ComponentType
    ) -> List[Entity]:
        return self.component_entities.get(component_type, [])

    def add_entity(self, entity: Entity) -> None:
        self.entities[entity.id] = entity
        for component in entity.components.values():
            if component.component_type not in self.component_entities:
                self.component_entities[component.component_type] = set()
            self.component_entities[component.component_type].add(entity)

    def remove_entity(self, entity_id: str) -> None:
        if entity_id not in self.entities:
            return
        entity = self.entities[entity_id]
        for component in entity.components.values():
            if component.component_type in self.component_entities:
                self.component_entities[component.component_type].remove(entity)
        del self.entities[entity_id]

    def get_entity(self, entity_id: str) -> Optional[Entity]:
        return self.entities.get(entity_id)

    def to_dict(self) -> dict:
        return {k: v.to_dict() for k, v in self.entities.items()}