package coffee.liz.ecs;

import coffee.liz.ecs.model.Component;
import coffee.liz.ecs.model.Entity;
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.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
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;

	public DAGWorld(final Map<Class<? extends System<T>>, System<T>> systems) {
		this.systems = systems;
		this.systemExecutionOrder = buildExecutionOrder(systems.values().stream().toList());
		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> query(final Collection<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());
	}

	@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 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));
		final Map<Class<?>, Integer> inDegree = new HashMap<>();
		final Map<Class<?>, Set<Class<?>>> adjacencyList = new HashMap<>();

		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 close() throws Exception {
		for (final System<T> system : systemExecutionOrder) {
			system.close();
		}
		componentCache.clear();
		entities.clear();
	}
}