from typing import Iterator


class World:
    def __init__(self):
        self.entities: set[Entity] = set()
        self.mapping: dict[type, set[Entity]] = {}

    def query(self, *needed: type, without: tuple[type] = []) -> Iterator['Entity']:
        if not needed:
            for entity in self.entities:
                for without_type in without:
                    if without_type in entity:
                        break
                else:
                    yield entity
        else:
            first_component = needed[0]
            for entity in self.mapping.get(first_component, set()):
                for component_type in needed[1:]:
                    if entity not in self.mapping.get(component_type, set()):
                        break
                else:
                    for without_type in without:
                        if without_type in entity:
                            break
                    else:
                        yield entity

class Entity:
    def __init__(self, world: World):
        self.world = world
        self.world.entities.add(self)
        self.components: dict[type, object] = {}

    def __getitem__(self, component_type):
        return self.components[component_type]
    
    def __setitem__(self, component_type, component):
        self.components[component_type] = component
        self.world.mapping.setdefault(component_type, set()).add(self)
    
    def __delitem__(self, component_type):
        del self.components[component_type]
        self.world.mapping[component_type].remove(self)
        if len(self.world.mapping[component_type]) == 0:
            del self.world.mapping[component_type]
          
    def __contains__(self, component_type):
        return component_type in self.components
    
    def __del__(self): # TODO: Réparer ça marche pas car il n'est pas détruit car il est toujours dans World
        self.world.entities.remove(self)
        for component_type in self.components:
            self.world.mapping[component_type].remove(self)
            if len(self.world.mapping[component_type]) == 0:
                del self.world.mapping[component_type]