package coffee.liz.ecs;

import coffee.liz.ecs.model.Component;
import coffee.liz.ecs.model.Entity;
import coffee.liz.ecs.model.Query;
import coffee.liz.ecs.model.System;
import coffee.liz.ecs.model.World;

import lombok.RequiredArgsConstructor;
import lombok.extern.log4j.Log4j2;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/** World that updates in {@link System#getDependencies()} topological order. */
@Log4j2
@RequiredArgsConstructor
public class DAGWorld<T> implements World<T> {
    /** All entities in the world. */
    protected final Set<Entity> entities = Collections.synchronizedSet(new HashSet<>());

    /** Cache mapping component types to entities having that component. */
    private final Map<Class<? extends Component>, Set<Entity>> componentCache = Collections
            .synchronizedMap(new HashMap<>());

    /** Deterministic ID's for spawned entities. */
    private final AtomicInteger nextEntityId = new AtomicInteger(0);

    /** All registered systems. */
    protected final Map<Class<? extends System<T>>, System<T>> systems;

    /** Ordered list of systems for execution. */
    private final List<System<T>> systemExecutionOrder;

    @SafeVarargs
    public DAGWorld(final System<T>... systems) {
        this.systems = singletonClazzMap(systems);
        this.systemExecutionOrder = buildExecutionOrder(Arrays.asList(systems));
        log.debug("Executing in order: {}", systemExecutionOrder);
    }

    @Override
    public Entity createEntity() {
        final Entity entity = Entity.builder().id(nextEntityId.incrementAndGet()).build();
        entities.add(entity);
        return entity;
    }

    @Override
    public void removeEntity(final Entity entity) {
        entity.getComponentMap().keySet().forEach(componentType -> {
            final Set<Entity> cachedEntities = componentCache.get(componentType);
            if (cachedEntities != null) {
                cachedEntities.remove(entity);
            }
        });
        entities.remove(entity);
    }

    @Override
    public Set<Entity> resolve(final Query query) {
        final Set<Class<? extends Component>> components = query.queryingComponents();
        return switch (query.filter()) {
            case ALL_OF -> resolveAllOf(components);
            case ANY_OF -> resolveAnyOf(components);
            case NONE_OF -> resolveNoneOf(components);
        };
    }

    @Override
    public void update(final T state, final Duration duration) {
        systemExecutionOrder.forEach(system -> {
            refreshComponentCache();
            system.update(this, state, duration);
        });
        refreshComponentCache();
    }

    @SuppressWarnings("unchecked")
    @Override
    public <S extends System<T>> S getSystem(final Class<S> system) {
        return (S) systems.get(system);
    }

    private void refreshComponentCache() {
        componentCache.clear();
        entities.forEach(entity -> entity.getComponentMap().keySet().forEach(
                componentType -> componentCache.computeIfAbsent(componentType, _ -> new HashSet<>()).add(entity)));
    }

    private Set<Entity> resolveAllOf(final Set<Class<? extends Component>> components) {
        if (components.isEmpty()) {
            return Set.copyOf(entities);
        }

        final Class<? extends Component> firstType = components.iterator().next();
        final Set<Entity> candidates = componentCache.get(firstType);
        if (candidates == null) {
            return Collections.emptySet();
        }

        return candidates.stream().filter(entity -> components.stream().allMatch(entity::has))
                .collect(Collectors.toSet());
    }

    private Set<Entity> resolveAnyOf(final Set<Class<? extends Component>> components) {
        if (components.isEmpty()) {
            return Collections.emptySet();
        }

        return entities.stream().filter(entity -> components.stream().anyMatch(entity::has))
                .collect(Collectors.toSet());
    }

    private Set<Entity> resolveNoneOf(final Set<Class<? extends Component>> components) {
        if (components.isEmpty()) {
            return Set.copyOf(entities);
        }

        return entities.stream().filter(entity -> components.stream().noneMatch(entity::has))
                .collect(Collectors.toSet());
    }

    private List<System<T>> buildExecutionOrder(final Collection<System<T>> systems) {
        if (systems.isEmpty()) {
            return Collections.emptyList();
        }

        final Map<Class<?>, System<T>> systemMap = systems.stream()
                .collect(Collectors.toMap(System::getClass, system -> system, (_, b) -> b, LinkedHashMap::new));
        final Map<Class<?>, Integer> inDegree = new LinkedHashMap<>();
        final Map<Class<?>, Set<Class<?>>> adjacencyList = new LinkedHashMap<>();

        systems.forEach(system -> {
            final Class<?> systemClass = system.getClass();
            inDegree.put(systemClass, 0);
            adjacencyList.put(systemClass, new HashSet<>());
        });

        systems.forEach(system -> {
            system.getDependencies().forEach(dependency -> {
                if (systemMap.containsKey(dependency)) {
                    adjacencyList.get(dependency).add(system.getClass());
                    inDegree.merge(system.getClass(), 1, Integer::sum);
                }
            });
        });

        // Kahn's algorithm
        final List<System<T>> result = new ArrayList<>();

        final Queue<Class<?>> queue = new LinkedList<>(
                inDegree.entrySet().stream().filter(entry -> entry.getValue() == 0).map(Map.Entry::getKey).toList());

        while (!queue.isEmpty()) {
            final Class<?> currentClass = queue.poll();
            result.add(systemMap.get(currentClass));

            adjacencyList.getOrDefault(currentClass, Collections.emptySet()).forEach(dependent -> {
                final int newInDegree = inDegree.get(dependent) - 1;
                inDegree.put(dependent, newInDegree);
                if (newInDegree == 0) {
                    queue.add(dependent);
                }
            });
        }

        if (result.size() != systems.size()) {
            throw new IllegalStateException("Circular dependency detected in systems");
        }

        return Collections.unmodifiableList(result);
    }

    @Override
    public void dispose() {
        for (final System<T> system : systemExecutionOrder) {
            system.dispose();
        }
        componentCache.clear();
        entities.clear();
    }

    @SuppressWarnings("unchecked")
    private static <T> Map<Class<? extends T>, T> singletonClazzMap(final T... singletons) {
        final boolean areSingletons = Arrays.stream(singletons).map(t -> (Class<? extends System<T>>) t.getClass())
                .distinct().count() == singletons.length;
        if (!areSingletons) {
            throw new IllegalArgumentException("Only one instance may be used per clazz");
        }

        return Arrays.stream(singletons).map(t -> Map.entry((Class<? extends T>) t.getClass(), t))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    }
}