package net.minecraft.gametest.framework;

import com.google.common.collect.Lists;
import com.google.common.collect.Streams;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import net.minecraft.core.Holder;
import net.minecraft.core.Holder.Reference;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.world.level.block.Rotation;

public class GameTestBatchFactory {
	private static final int MAX_TESTS_PER_BATCH = 50;
	public static final GameTestBatchFactory.TestDecorator DIRECT = (reference, serverLevel) -> Stream.of(
		new GameTestInfo(reference, Rotation.NONE, serverLevel, RetryOptions.noRetries())
	);

	public static List<GameTestBatch> divideIntoBatches(
		Collection<Reference<GameTestInstance>> instances, GameTestBatchFactory.TestDecorator decorator, ServerLevel level
	) {
		Map<Holder<TestEnvironmentDefinition>, List<GameTestInfo>> map = (Map<Holder<TestEnvironmentDefinition>, List<GameTestInfo>>)instances.stream()
			.flatMap(reference -> decorator.decorate(reference, level))
			.collect(Collectors.groupingBy(gameTestInfo -> gameTestInfo.getTest().batch()));
		return map.entrySet().stream().flatMap(entry -> {
			Holder<TestEnvironmentDefinition> holder = (Holder<TestEnvironmentDefinition>)entry.getKey();
			List<GameTestInfo> list = (List<GameTestInfo>)entry.getValue();
			return Streams.mapWithIndex(Lists.partition(list, 50).stream(), (listx, l) -> toGameTestBatch(listx, holder, (int)l));
		}).toList();
	}

	public static GameTestRunner.GameTestBatcher fromGameTestInfo() {
		return fromGameTestInfo(50);
	}

	public static GameTestRunner.GameTestBatcher fromGameTestInfo(int maxTests) {
		return collection -> {
			Map<Holder<TestEnvironmentDefinition>, List<GameTestInfo>> map = (Map<Holder<TestEnvironmentDefinition>, List<GameTestInfo>>)collection.stream()
				.filter(Objects::nonNull)
				.collect(Collectors.groupingBy(gameTestInfo -> gameTestInfo.getTest().batch()));
			return map.entrySet().stream().flatMap(entry -> {
				Holder<TestEnvironmentDefinition> holder = (Holder<TestEnvironmentDefinition>)entry.getKey();
				List<GameTestInfo> list = (List<GameTestInfo>)entry.getValue();
				return Streams.mapWithIndex(Lists.partition(list, maxTests).stream(), (listx, l) -> toGameTestBatch(List.copyOf(listx), holder, (int)l));
			}).toList();
		};
	}

	public static GameTestBatch toGameTestBatch(Collection<GameTestInfo> gameTestInfos, Holder<TestEnvironmentDefinition> environment, int index) {
		return new GameTestBatch(index, gameTestInfos, environment);
	}

	@FunctionalInterface
	public interface TestDecorator {
		Stream<GameTestInfo> decorate(Reference<GameTestInstance> reference, ServerLevel serverLevel);
	}
}