package net.minecraft.server.packs.resources;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicInteger;
import net.minecraft.Util;
import net.minecraft.util.Unit;
import org.jetbrains.annotations.Nullable;

public class SimpleReloadInstance<S> implements ReloadInstance {
	private static final int PREPARATION_PROGRESS_WEIGHT = 2;
	private static final int EXTRA_RELOAD_PROGRESS_WEIGHT = 2;
	private static final int LISTENER_PROGRESS_WEIGHT = 1;
	final CompletableFuture<Unit> allPreparations = new CompletableFuture();
	@Nullable
	private CompletableFuture<List<S>> allDone;
	final Set<PreparableReloadListener> preparingListeners;
	private final int listenerCount;
	private final AtomicInteger startedTasks = new AtomicInteger();
	private final AtomicInteger finishedTasks = new AtomicInteger();
	private final AtomicInteger startedReloads = new AtomicInteger();
	private final AtomicInteger finishedReloads = new AtomicInteger();

	public static ReloadInstance of(
		ResourceManager resourceManager,
		List<PreparableReloadListener> listeners,
		Executor backgroundExecutor,
		Executor gameExecutor,
		CompletableFuture<Unit> alsoWaitedFor
	) {
		SimpleReloadInstance<Void> simpleReloadInstance = new SimpleReloadInstance<>(listeners);
		simpleReloadInstance.startTasks(backgroundExecutor, gameExecutor, resourceManager, listeners, SimpleReloadInstance.StateFactory.SIMPLE, alsoWaitedFor);
		return simpleReloadInstance;
	}

	protected SimpleReloadInstance(List<PreparableReloadListener> preparingListeners) {
		this.listenerCount = preparingListeners.size();
		this.preparingListeners = new HashSet(preparingListeners);
	}

	protected void startTasks(
		Executor backgroundExecutor,
		Executor gameExecutor,
		ResourceManager resourceManager,
		List<PreparableReloadListener> listeners,
		SimpleReloadInstance.StateFactory<S> stateFactory,
		CompletableFuture<?> alsoWaitedFor
	) {
		this.allDone = this.prepareTasks(backgroundExecutor, gameExecutor, resourceManager, listeners, stateFactory, alsoWaitedFor);
	}

	protected CompletableFuture<List<S>> prepareTasks(
		Executor backgroundExecutor,
		Executor gameExectutor,
		ResourceManager resourceManager,
		List<PreparableReloadListener> listeners,
		SimpleReloadInstance.StateFactory<S> stateFactory,
		CompletableFuture<?> alsoWaitedFor
	) {
		Executor executor = runnable -> {
			this.startedTasks.incrementAndGet();
			backgroundExecutor.execute(() -> {
				runnable.run();
				this.finishedTasks.incrementAndGet();
			});
		};
		Executor executor2 = runnable -> {
			this.startedReloads.incrementAndGet();
			gameExectutor.execute(() -> {
				runnable.run();
				this.finishedReloads.incrementAndGet();
			});
		};
		this.startedTasks.incrementAndGet();
		alsoWaitedFor.thenRun(this.finishedTasks::incrementAndGet);
		CompletableFuture<?> completableFuture = alsoWaitedFor;
		List<CompletableFuture<S>> list = new ArrayList();

		for (PreparableReloadListener preparableReloadListener : listeners) {
			PreparableReloadListener.PreparationBarrier preparationBarrier = this.createBarrierForListener(preparableReloadListener, completableFuture, gameExectutor);
			CompletableFuture<S> completableFuture2 = stateFactory.create(preparationBarrier, resourceManager, preparableReloadListener, executor, executor2);
			list.add(completableFuture2);
			completableFuture = completableFuture2;
		}

		return Util.sequenceFailFast(list);
	}

	private PreparableReloadListener.PreparationBarrier createBarrierForListener(
		PreparableReloadListener listener, CompletableFuture<?> alsoWaitedFor, Executor executor
	) {
		return new PreparableReloadListener.PreparationBarrier() {
			@Override
			public <T> CompletableFuture<T> wait(T object) {
				executor.execute(() -> {
					SimpleReloadInstance.this.preparingListeners.remove(listener);
					if (SimpleReloadInstance.this.preparingListeners.isEmpty()) {
						SimpleReloadInstance.this.allPreparations.complete(Unit.INSTANCE);
					}
				});
				return SimpleReloadInstance.this.allPreparations.thenCombine(alsoWaitedFor, (unit, object2) -> object);
			}
		};
	}

	@Override
	public CompletableFuture<?> done() {
		return (CompletableFuture<?>)Objects.requireNonNull(this.allDone, "not started");
	}

	@Override
	public float getActualProgress() {
		int i = this.listenerCount - this.preparingListeners.size();
		float f = weightProgress(this.finishedTasks.get(), this.finishedReloads.get(), i);
		float g = weightProgress(this.startedTasks.get(), this.startedReloads.get(), this.listenerCount);
		return f / g;
	}

	private static int weightProgress(int tasks, int reloads, int listeners) {
		return tasks * 2 + reloads * 2 + listeners * 1;
	}

	public static ReloadInstance create(
		ResourceManager resourceManager,
		List<PreparableReloadListener> listeners,
		Executor backgroundExecutor,
		Executor gameExecutor,
		CompletableFuture<Unit> alsoWaitedFor,
		boolean profiled
	) {
		return profiled
			? ProfiledReloadInstance.of(resourceManager, listeners, backgroundExecutor, gameExecutor, alsoWaitedFor)
			: of(resourceManager, listeners, backgroundExecutor, gameExecutor, alsoWaitedFor);
	}

	@FunctionalInterface
	protected interface StateFactory<S> {
		SimpleReloadInstance.StateFactory<Void> SIMPLE = (preparationBarrier, resourceManager, preparableReloadListener, executor, executor2) -> preparableReloadListener.reload(
			preparationBarrier, resourceManager, executor, executor2
		);

		CompletableFuture<S> create(
			PreparableReloadListener.PreparationBarrier preparationBarrier,
			ResourceManager resourceManager,
			PreparableReloadListener preparableReloadListener,
			Executor executor,
			Executor executor2
		);
	}
}