from dataclasses import dataclass, field
from typing import Any, Iterator, Protocol

type Entity = int
type Archetype = frozenset[type]
type ComponentMap = dict[type, Any]
type EntityComponentStorage = dict[Entity, ComponentMap]


class System(Protocol):
    @classmethod
    def process(cls, world: "World") -> None: ...


@dataclass
class World:
    registry: dict[Archetype, EntityComponentStorage] = field(default_factory=dict)
    next_entity_index: int = 0
    systems: list[type[System]] = field(default_factory=list)

    def generate_id(self) -> int:
        self.next_entity_index += 1
        return self.next_entity_index

    def spawn(self, *components) -> int:
        archetype = frozenset([type(c) for c in components])

        if archetype not in self.registry:
            self.registry[archetype] = {}

        entity = self.generate_id()
        self.registry[archetype][entity] = {type(c): c for c in components}
        return entity

    def query(
        self, *component_types, exact=False
    ) -> Iterator[tuple[Entity, ComponentMap]]:
        component_subset_signature = frozenset(component_types)
        for archetype, entity_component_storage in self.registry.items():
            if (exact and component_subset_signature == archetype) or (
                not exact and component_subset_signature.issubset(archetype)
            ):
                for entity, component_storage in entity_component_storage.items():
                    yield entity, component_storage

    def get_components(self, entity: int, pop=False):
        for componentstorage in self.registry.values():
            if entity in componentstorage:
                if pop:
                    return componentstorage.pop(entity, None)
                else:
                    return componentstorage[entity]
        return None

    def tick(self, **args):
        for system in self.systems:
            system.process(self, **args)