package coffee.liz.ecs;

import coffee.liz.ecs.events.ComponentAdded;
import coffee.liz.ecs.events.ComponentRemoved;
import coffee.liz.ecs.events.EntityEvent;
import coffee.liz.ecs.model.Component;
import coffee.liz.ecs.model.Entity;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/** Incremental cache for component-to-entity lookups. */
public final class ComponentCache {
    private final Map<Class<? extends Component>, Set<Entity>> entitiesByComponent = new HashMap<>();

    /**
     * Index all current components on a newly tracked entity.
     *
     * @param entity
     *            entity to index
     */
    public void addEntity(final Entity entity) {
        entity.componentTypes().forEach(componentType -> addComponentEntry(componentType, entity));
    }

    /**
     * Remove all current components for a no-longer-tracked entity.
     *
     * @param entity
     *            entity to remove from the cache
     */
    public void removeEntity(final Entity entity) {
        entity.componentTypes().forEach(componentType -> removeComponentEntry(componentType, entity));
    }

    /**
     * Apply a mutation event to the cache.
     *
     * @param event
     *            mutation event emitted by an entity
     */
    public void onEntityEvent(final EntityEvent event) {
        switch (event) {
            case ComponentAdded a -> addComponentEntry(a.getComponentClazz(), a.getEntity());
            case ComponentRemoved r -> removeComponentEntry(r.getComponentType(), r.getEntity());
            default -> throw new IllegalStateException("Unexpected value: " + event);
        }
    }

    /**
     * Get all entities that contain a component type.
     *
     * @param componentType
     *            component class to query
     * @return entities containing that component
     */
    public Set<Entity> entitiesWith(final Class<? extends Component> componentType) {
        return entitiesByComponent.getOrDefault(componentType, Collections.emptySet());
    }

    /**
     * Resolve entities that contain every queried component type.
     *
     * @param componentTypes
     *            required component classes
     * @return entities that have all queried component types
     */
    public Set<Entity> entitiesMatchingAllOf(final Collection<Class<? extends Component>> componentTypes) {
        if (componentTypes.isEmpty()) {
            return Collections.emptySet();
        }

        Set<Entity> smallestBucket = null;
        for (final Class<? extends Component> componentType : componentTypes) {
            final Set<Entity> entities = entitiesWith(componentType);
            if (entities.isEmpty()) {
                return Collections.emptySet();
            }
            if (smallestBucket == null || entities.size() < smallestBucket.size()) {
                smallestBucket = entities;
            }
        }

        final Set<Entity> matches = new HashSet<>();
        if (smallestBucket == null) {
            return matches;
        }

        for (final Entity entity : smallestBucket) {
            if (entity.hasAll(componentTypes)) {
                matches.add(entity);
            }
        }
        return matches;
    }

    /**
     * Remove all cached component indexes.
     */
    public void clear() {
        entitiesByComponent.clear();
    }

    private void addComponentEntry(final Class<? extends Component> componentType, final Entity entity) {
        entitiesByComponent.computeIfAbsent(componentType, _unused -> new HashSet<>()).add(entity);
    }

    private void removeComponentEntry(final Class<? extends Component> componentType, final Entity entity) {
        final Set<Entity> entities = entitiesByComponent.get(componentType);
        if (entities == null) {
            return;
        }

        entities.remove(entity);
        if (entities.isEmpty()) {
            entitiesByComponent.remove(componentType);
        }
    }
}