package net.minecraft.server.level;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableSet;
import com.mojang.logging.LogUtils;
import it.unimi.dsi.fastutil.longs.Long2ByteMap;
import it.unimi.dsi.fastutil.longs.Long2ByteOpenHashMap;
import it.unimi.dsi.fastutil.longs.Long2IntMap;
import it.unimi.dsi.fastutil.longs.Long2IntMaps;
import it.unimi.dsi.fastutil.longs.Long2IntOpenHashMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectFunction;
import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectOpenHashMap;
import it.unimi.dsi.fastutil.longs.LongIterator;
import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
import it.unimi.dsi.fastutil.longs.LongSet;
import it.unimi.dsi.fastutil.longs.Long2ObjectMap.Entry;
import it.unimi.dsi.fastutil.objects.ObjectIterator;
import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
import it.unimi.dsi.fastutil.objects.ObjectSet;
import it.unimi.dsi.fastutil.objects.ReferenceOpenHashSet;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import net.minecraft.core.SectionPos;
import net.minecraft.util.SortedArraySet;
import net.minecraft.util.thread.TaskScheduler;
import net.minecraft.world.level.ChunkPos;
import net.minecraft.world.level.chunk.LevelChunk;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;

public abstract class DistanceManager {
	static final Logger LOGGER = LogUtils.getLogger();
	static final int PLAYER_TICKET_LEVEL = ChunkLevel.byStatus(FullChunkStatus.ENTITY_TICKING);
	private static final int INITIAL_TICKET_LIST_CAPACITY = 4;
	final Long2ObjectMap<ObjectSet<ServerPlayer>> playersPerChunk = new Long2ObjectOpenHashMap<>();
	final Long2ObjectOpenHashMap<SortedArraySet<Ticket<?>>> tickets = new Long2ObjectOpenHashMap<>();
	private final DistanceManager.ChunkTicketTracker ticketTracker = new DistanceManager.ChunkTicketTracker();
	private final DistanceManager.FixedPlayerDistanceChunkTracker naturalSpawnChunkCounter = new DistanceManager.FixedPlayerDistanceChunkTracker(8);
	private final TickingTracker tickingTicketsTracker = new TickingTracker();
	private final DistanceManager.PlayerTicketTracker playerTicketManager = new DistanceManager.PlayerTicketTracker(32);
	final Set<ChunkHolder> chunksToUpdateFutures = new ReferenceOpenHashSet<>();
	final ThrottlingChunkTaskDispatcher ticketDispatcher;
	final LongSet ticketsToRelease = new LongOpenHashSet();
	final Executor mainThreadExecutor;
	private long ticketTickCounter;
	private int simulationDistance = 10;

	protected DistanceManager(Executor dispatcher, Executor mainThreadExecutor) {
		TaskScheduler<Runnable> taskScheduler = TaskScheduler.wrapExecutor("player ticket throttler", mainThreadExecutor);
		this.ticketDispatcher = new ThrottlingChunkTaskDispatcher(taskScheduler, dispatcher, 4);
		this.mainThreadExecutor = mainThreadExecutor;
	}

	protected void purgeStaleTickets() {
		this.ticketTickCounter++;
		ObjectIterator<Entry<SortedArraySet<Ticket<?>>>> objectIterator = this.tickets.long2ObjectEntrySet().fastIterator();

		while (objectIterator.hasNext()) {
			Entry<SortedArraySet<Ticket<?>>> entry = (Entry<SortedArraySet<Ticket<?>>>)objectIterator.next();
			Iterator<Ticket<?>> iterator = ((SortedArraySet)entry.getValue()).iterator();
			boolean bl = false;

			while (iterator.hasNext()) {
				Ticket<?> ticket = (Ticket<?>)iterator.next();
				if (ticket.timedOut(this.ticketTickCounter)) {
					iterator.remove();
					bl = true;
					this.tickingTicketsTracker.removeTicket(entry.getLongKey(), ticket);
				}
			}

			if (bl) {
				this.ticketTracker.update(entry.getLongKey(), getTicketLevelAt((SortedArraySet<Ticket<?>>)entry.getValue()), false);
			}

			if (((SortedArraySet)entry.getValue()).isEmpty()) {
				objectIterator.remove();
			}
		}
	}

	/**
	 * Gets the {@linkplain net.minecraft.server.level.Ticket#getTicketLevel level} of the ticket.
	 */
	private static int getTicketLevelAt(SortedArraySet<Ticket<?>> tickets) {
		return !tickets.isEmpty() ? tickets.first().getTicketLevel() : ChunkLevel.MAX_LEVEL + 1;
	}

	protected abstract boolean isChunkToRemove(long chunkPos);

	@Nullable
	protected abstract ChunkHolder getChunk(long chunkPos);

	@Nullable
	protected abstract ChunkHolder updateChunkScheduling(long chunkPos, int newLevel, @Nullable ChunkHolder holder, int oldLevel);

	public boolean runAllUpdates(ChunkMap chunkMap) {
		this.naturalSpawnChunkCounter.runAllUpdates();
		this.tickingTicketsTracker.runAllUpdates();
		this.playerTicketManager.runAllUpdates();
		int i = Integer.MAX_VALUE - this.ticketTracker.runDistanceUpdates(Integer.MAX_VALUE);
		boolean bl = i != 0;
		if (bl) {
		}

		if (!this.chunksToUpdateFutures.isEmpty()) {
			for (ChunkHolder chunkHolder : this.chunksToUpdateFutures) {
				chunkHolder.updateHighestAllowedStatus(chunkMap);
			}

			for (ChunkHolder chunkHolder : this.chunksToUpdateFutures) {
				chunkHolder.updateFutures(chunkMap, this.mainThreadExecutor);
			}

			this.chunksToUpdateFutures.clear();
			return true;
		} else {
			if (!this.ticketsToRelease.isEmpty()) {
				LongIterator longIterator = this.ticketsToRelease.iterator();

				while (longIterator.hasNext()) {
					long l = longIterator.nextLong();
					if (this.getTickets(l).stream().anyMatch(ticket -> ticket.getType() == TicketType.PLAYER)) {
						ChunkHolder chunkHolder2 = chunkMap.getUpdatingChunkIfPresent(l);
						if (chunkHolder2 == null) {
							throw new IllegalStateException();
						}

						CompletableFuture<ChunkResult<LevelChunk>> completableFuture = chunkHolder2.getEntityTickingChunkFuture();
						completableFuture.thenAccept(chunkResult -> this.mainThreadExecutor.execute(() -> this.ticketDispatcher.release(l, () -> {}, false)));
					}
				}

				this.ticketsToRelease.clear();
			}

			return bl;
		}
	}

	void addTicket(long chunkPos, Ticket<?> ticket) {
		SortedArraySet<Ticket<?>> sortedArraySet = this.getTickets(chunkPos);
		int i = getTicketLevelAt(sortedArraySet);
		Ticket<?> ticket2 = sortedArraySet.addOrGet(ticket);
		ticket2.setCreatedTick(this.ticketTickCounter);
		if (ticket.getTicketLevel() < i) {
			this.ticketTracker.update(chunkPos, ticket.getTicketLevel(), true);
		}
	}

	void removeTicket(long chunkPos, Ticket<?> ticket) {
		SortedArraySet<Ticket<?>> sortedArraySet = this.getTickets(chunkPos);
		if (sortedArraySet.remove(ticket)) {
		}

		if (sortedArraySet.isEmpty()) {
			this.tickets.remove(chunkPos);
		}

		this.ticketTracker.update(chunkPos, getTicketLevelAt(sortedArraySet), false);
	}

	public <T> void addTicket(TicketType<T> type, ChunkPos pos, int level, T value) {
		this.addTicket(pos.toLong(), new Ticket<>(type, level, value));
	}

	public <T> void removeTicket(TicketType<T> type, ChunkPos pos, int level, T value) {
		Ticket<T> ticket = new Ticket<>(type, level, value);
		this.removeTicket(pos.toLong(), ticket);
	}

	public <T> void addRegionTicket(TicketType<T> type, ChunkPos pos, int distance, T value) {
		Ticket<T> ticket = new Ticket<>(type, ChunkLevel.byStatus(FullChunkStatus.FULL) - distance, value);
		long l = pos.toLong();
		this.addTicket(l, ticket);
		this.tickingTicketsTracker.addTicket(l, ticket);
	}

	public <T> void removeRegionTicket(TicketType<T> type, ChunkPos pos, int distance, T value) {
		Ticket<T> ticket = new Ticket<>(type, ChunkLevel.byStatus(FullChunkStatus.FULL) - distance, value);
		long l = pos.toLong();
		this.removeTicket(l, ticket);
		this.tickingTicketsTracker.removeTicket(l, ticket);
	}

	private SortedArraySet<Ticket<?>> getTickets(long chunkPos) {
		return this.tickets
			.computeIfAbsent(chunkPos, (Long2ObjectFunction<? extends SortedArraySet<Ticket<?>>>)(l -> (SortedArraySet<Ticket<?>>)SortedArraySet.create(4)));
	}

	protected void updateChunkForced(ChunkPos pos, boolean add) {
		Ticket<ChunkPos> ticket = new Ticket<>(TicketType.FORCED, ChunkMap.FORCED_TICKET_LEVEL, pos);
		long l = pos.toLong();
		if (add) {
			this.addTicket(l, ticket);
			this.tickingTicketsTracker.addTicket(l, ticket);
		} else {
			this.removeTicket(l, ticket);
			this.tickingTicketsTracker.removeTicket(l, ticket);
		}
	}

	public void addPlayer(SectionPos sectionPos, ServerPlayer player) {
		ChunkPos chunkPos = sectionPos.chunk();
		long l = chunkPos.toLong();
		this.playersPerChunk.computeIfAbsent(l, (Long2ObjectFunction<? extends ObjectSet<ServerPlayer>>)(lx -> new ObjectOpenHashSet<>())).add(player);
		this.naturalSpawnChunkCounter.update(l, 0, true);
		this.playerTicketManager.update(l, 0, true);
		this.tickingTicketsTracker.addTicket(TicketType.PLAYER, chunkPos, this.getPlayerTicketLevel(), chunkPos);
	}

	public void removePlayer(SectionPos sectionPos, ServerPlayer player) {
		ChunkPos chunkPos = sectionPos.chunk();
		long l = chunkPos.toLong();
		ObjectSet<ServerPlayer> objectSet = this.playersPerChunk.get(l);
		objectSet.remove(player);
		if (objectSet.isEmpty()) {
			this.playersPerChunk.remove(l);
			this.naturalSpawnChunkCounter.update(l, Integer.MAX_VALUE, false);
			this.playerTicketManager.update(l, Integer.MAX_VALUE, false);
			this.tickingTicketsTracker.removeTicket(TicketType.PLAYER, chunkPos, this.getPlayerTicketLevel(), chunkPos);
		}
	}

	private int getPlayerTicketLevel() {
		return Math.max(0, ChunkLevel.byStatus(FullChunkStatus.ENTITY_TICKING) - this.simulationDistance);
	}

	public boolean inEntityTickingRange(long chunkPos) {
		return ChunkLevel.isEntityTicking(this.tickingTicketsTracker.getLevel(chunkPos));
	}

	public boolean inBlockTickingRange(long chunkPos) {
		return ChunkLevel.isBlockTicking(this.tickingTicketsTracker.getLevel(chunkPos));
	}

	protected String getTicketDebugString(long chunkPos) {
		SortedArraySet<Ticket<?>> sortedArraySet = this.tickets.get(chunkPos);
		return sortedArraySet != null && !sortedArraySet.isEmpty() ? sortedArraySet.first().toString() : "no_ticket";
	}

	protected void updatePlayerTickets(int viewDistance) {
		this.playerTicketManager.updateViewDistance(viewDistance);
	}

	public void updateSimulationDistance(int simulationDistance) {
		if (simulationDistance != this.simulationDistance) {
			this.simulationDistance = simulationDistance;
			this.tickingTicketsTracker.replacePlayerTicketsLevel(this.getPlayerTicketLevel());
		}
	}

	/**
	 * Returns the number of chunks taken into account when calculating the mob cap
	 */
	public int getNaturalSpawnChunkCount() {
		this.naturalSpawnChunkCounter.runAllUpdates();
		return this.naturalSpawnChunkCounter.chunks.size();
	}

	public boolean hasPlayersNearby(long chunkPos) {
		this.naturalSpawnChunkCounter.runAllUpdates();
		return this.naturalSpawnChunkCounter.chunks.containsKey(chunkPos);
	}

	public LongIterator getSpawnCandidateChunks() {
		this.naturalSpawnChunkCounter.runAllUpdates();
		return this.naturalSpawnChunkCounter.chunks.keySet().iterator();
	}

	public String getDebugStatus() {
		return this.ticketDispatcher.getDebugStatus();
	}

	private void dumpTickets(String filename) {
		try {
			FileOutputStream fileOutputStream = new FileOutputStream(new File(filename));

			try {
				for (Entry<SortedArraySet<Ticket<?>>> entry : this.tickets.long2ObjectEntrySet()) {
					ChunkPos chunkPos = new ChunkPos(entry.getLongKey());

					for (Ticket<?> ticket : (SortedArraySet)entry.getValue()) {
						fileOutputStream.write(
							(chunkPos.x + "\t" + chunkPos.z + "\t" + ticket.getType() + "\t" + ticket.getTicketLevel() + "\t\n").getBytes(StandardCharsets.UTF_8)
						);
					}
				}
			} catch (Throwable var9) {
				try {
					fileOutputStream.close();
				} catch (Throwable var8) {
					var9.addSuppressed(var8);
				}

				throw var9;
			}

			fileOutputStream.close();
		} catch (IOException var10) {
			LOGGER.error("Failed to dump tickets to {}", filename, var10);
		}
	}

	@VisibleForTesting
	TickingTracker tickingTracker() {
		return this.tickingTicketsTracker;
	}

	public LongSet getTickingChunks() {
		return this.tickingTicketsTracker.getTickingChunks();
	}

	public void removeTicketsOnClosing() {
		ImmutableSet<TicketType<?>> immutableSet = ImmutableSet.of(TicketType.UNKNOWN, TicketType.POST_TELEPORT);
		ObjectIterator<Entry<SortedArraySet<Ticket<?>>>> objectIterator = this.tickets.long2ObjectEntrySet().fastIterator();

		while (objectIterator.hasNext()) {
			Entry<SortedArraySet<Ticket<?>>> entry = (Entry<SortedArraySet<Ticket<?>>>)objectIterator.next();
			Iterator<Ticket<?>> iterator = ((SortedArraySet)entry.getValue()).iterator();
			boolean bl = false;

			while (iterator.hasNext()) {
				Ticket<?> ticket = (Ticket<?>)iterator.next();
				if (!immutableSet.contains(ticket.getType())) {
					iterator.remove();
					bl = true;
					this.tickingTicketsTracker.removeTicket(entry.getLongKey(), ticket);
				}
			}

			if (bl) {
				this.ticketTracker.update(entry.getLongKey(), getTicketLevelAt((SortedArraySet<Ticket<?>>)entry.getValue()), false);
			}

			if (((SortedArraySet)entry.getValue()).isEmpty()) {
				objectIterator.remove();
			}
		}
	}

	public boolean hasTickets() {
		return !this.tickets.isEmpty();
	}

	class ChunkTicketTracker extends ChunkTracker {
		private static final int MAX_LEVEL = ChunkLevel.MAX_LEVEL + 1;

		public ChunkTicketTracker() {
			super(MAX_LEVEL + 1, 16, 256);
		}

		@Override
		protected int getLevelFromSource(long pos) {
			SortedArraySet<Ticket<?>> sortedArraySet = DistanceManager.this.tickets.get(pos);
			if (sortedArraySet == null) {
				return Integer.MAX_VALUE;
			} else {
				return sortedArraySet.isEmpty() ? Integer.MAX_VALUE : sortedArraySet.first().getTicketLevel();
			}
		}

		@Override
		protected int getLevel(long chunkPos) {
			if (!DistanceManager.this.isChunkToRemove(chunkPos)) {
				ChunkHolder chunkHolder = DistanceManager.this.getChunk(chunkPos);
				if (chunkHolder != null) {
					return chunkHolder.getTicketLevel();
				}
			}

			return MAX_LEVEL;
		}

		@Override
		protected void setLevel(long chunkPos, int level) {
			ChunkHolder chunkHolder = DistanceManager.this.getChunk(chunkPos);
			int i = chunkHolder == null ? MAX_LEVEL : chunkHolder.getTicketLevel();
			if (i != level) {
				chunkHolder = DistanceManager.this.updateChunkScheduling(chunkPos, level, chunkHolder, i);
				if (chunkHolder != null) {
					DistanceManager.this.chunksToUpdateFutures.add(chunkHolder);
				}
			}
		}

		public int runDistanceUpdates(int toUpdateCount) {
			return this.runUpdates(toUpdateCount);
		}
	}

	class FixedPlayerDistanceChunkTracker extends ChunkTracker {
		/**
		 * Chunks that are at most {@link #range} chunks away from the closest player.
		 */
		protected final Long2ByteMap chunks = new Long2ByteOpenHashMap();
		protected final int maxDistance;

		protected FixedPlayerDistanceChunkTracker(final int maxDistance) {
			super(maxDistance + 2, 16, 256);
			this.maxDistance = maxDistance;
			this.chunks.defaultReturnValue((byte)(maxDistance + 2));
		}

		@Override
		protected int getLevel(long chunkPos) {
			return this.chunks.get(chunkPos);
		}

		@Override
		protected void setLevel(long chunkPos, int level) {
			byte b;
			if (level > this.maxDistance) {
				b = this.chunks.remove(chunkPos);
			} else {
				b = this.chunks.put(chunkPos, (byte)level);
			}

			this.onLevelChange(chunkPos, b, level);
		}

		/**
		 * Called after {@link PlayerChunkTracker#setLevel(long, int)} puts/removes chunk into/from {@link #chunksInRange}.
		 * 
		 * @param oldLevel Previous level of the chunk if it was smaller than {@link #range}, {@code range + 2} otherwise.
		 */
		protected void onLevelChange(long chunkPos, int oldLevel, int newLevel) {
		}

		@Override
		protected int getLevelFromSource(long pos) {
			return this.havePlayer(pos) ? 0 : Integer.MAX_VALUE;
		}

		private boolean havePlayer(long chunkPos) {
			ObjectSet<ServerPlayer> objectSet = DistanceManager.this.playersPerChunk.get(chunkPos);
			return objectSet != null && !objectSet.isEmpty();
		}

		public void runAllUpdates() {
			this.runUpdates(Integer.MAX_VALUE);
		}

		private void dumpChunks(String filename) {
			try {
				FileOutputStream fileOutputStream = new FileOutputStream(new File(filename));

				try {
					for (it.unimi.dsi.fastutil.longs.Long2ByteMap.Entry entry : this.chunks.long2ByteEntrySet()) {
						ChunkPos chunkPos = new ChunkPos(entry.getLongKey());
						String string = Byte.toString(entry.getByteValue());
						fileOutputStream.write((chunkPos.x + "\t" + chunkPos.z + "\t" + string + "\n").getBytes(StandardCharsets.UTF_8));
					}
				} catch (Throwable var8) {
					try {
						fileOutputStream.close();
					} catch (Throwable var7) {
						var8.addSuppressed(var7);
					}

					throw var8;
				}

				fileOutputStream.close();
			} catch (IOException var9) {
				DistanceManager.LOGGER.error("Failed to dump chunks to {}", filename, var9);
			}
		}
	}

	class PlayerTicketTracker extends DistanceManager.FixedPlayerDistanceChunkTracker {
		private int viewDistance;
		private final Long2IntMap queueLevels = Long2IntMaps.synchronize(new Long2IntOpenHashMap());
		private final LongSet toUpdate = new LongOpenHashSet();

		protected PlayerTicketTracker(final int maxDistance) {
			super(maxDistance);
			this.viewDistance = 0;
			this.queueLevels.defaultReturnValue(maxDistance + 2);
		}

		@Override
		protected void onLevelChange(long chunkPos, int oldLevel, int newLevel) {
			this.toUpdate.add(chunkPos);
		}

		public void updateViewDistance(int viewDistance) {
			for (it.unimi.dsi.fastutil.longs.Long2ByteMap.Entry entry : this.chunks.long2ByteEntrySet()) {
				byte b = entry.getByteValue();
				long l = entry.getLongKey();
				this.onLevelChange(l, b, this.haveTicketFor(b), b <= viewDistance);
			}

			this.viewDistance = viewDistance;
		}

		private void onLevelChange(long chunkPos, int level, boolean hadTicket, boolean hasTicket) {
			if (hadTicket != hasTicket) {
				Ticket<?> ticket = new Ticket<>(TicketType.PLAYER, DistanceManager.PLAYER_TICKET_LEVEL, new ChunkPos(chunkPos));
				if (hasTicket) {
					DistanceManager.this.ticketDispatcher.submit(() -> DistanceManager.this.mainThreadExecutor.execute(() -> {
						if (this.haveTicketFor(this.getLevel(chunkPos))) {
							DistanceManager.this.addTicket(chunkPos, ticket);
							DistanceManager.this.ticketsToRelease.add(chunkPos);
						} else {
							DistanceManager.this.ticketDispatcher.release(chunkPos, () -> {}, false);
						}
					}), chunkPos, () -> level);
				} else {
					DistanceManager.this.ticketDispatcher
						.release(chunkPos, () -> DistanceManager.this.mainThreadExecutor.execute(() -> DistanceManager.this.removeTicket(chunkPos, ticket)), true);
				}
			}
		}

		@Override
		public void runAllUpdates() {
			super.runAllUpdates();
			if (!this.toUpdate.isEmpty()) {
				LongIterator longIterator = this.toUpdate.iterator();

				while (longIterator.hasNext()) {
					long l = longIterator.nextLong();
					int i = this.queueLevels.get(l);
					int j = this.getLevel(l);
					if (i != j) {
						DistanceManager.this.ticketDispatcher.onLevelChange(new ChunkPos(l), () -> this.queueLevels.get(l), j, ix -> {
							if (ix >= this.queueLevels.defaultReturnValue()) {
								this.queueLevels.remove(l);
							} else {
								this.queueLevels.put(l, ix);
							}
						});
						this.onLevelChange(l, j, this.haveTicketFor(i), this.haveTicketFor(j));
					}
				}

				this.toUpdate.clear();
			}
		}

		private boolean haveTicketFor(int level) {
			return level <= this.viewDistance;
		}
	}
}