package net.minecraft.world.level.chunk.storage;

import com.google.common.collect.ImmutableList;
import com.mojang.datafixers.DataFixer;
import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
import it.unimi.dsi.fastutil.longs.LongSet;
import java.io.File;
import java.io.IOException;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import net.minecraft.SharedConstants;
import net.minecraft.nbt.GameProfileSerializer;
import net.minecraft.nbt.NBTTagCompound;
import net.minecraft.nbt.NBTTagIntArray;
import net.minecraft.nbt.NBTTagList;
import net.minecraft.server.level.WorldServer;
import net.minecraft.util.datafix.DataFixTypes;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.EntityTypes;
import net.minecraft.world.level.ChunkCoordIntPair;
import net.minecraft.world.level.World;
import net.minecraft.world.level.entity.ChunkEntities;
import net.minecraft.world.level.entity.EntityPersistentStorage;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

public class EntityStorage implements EntityPersistentStorage<Entity> {

    private static final Logger LOGGER = LogManager.getLogger();
    private static final String ENTITIES_TAG = "Entities";
    private static final String POSITION_TAG = "Position";
    private final WorldServer level;
    private final IOWorker worker;
    private final LongSet emptyChunks = new LongOpenHashSet();
    private final Executor mainThreadExecutor;
    protected final DataFixer fixerUpper;

    public EntityStorage(WorldServer worldserver, File file, DataFixer datafixer, boolean flag, Executor executor) {
        this.level = worldserver;
        this.fixerUpper = datafixer;
        this.mainThreadExecutor = executor;
        this.worker = new IOWorker(file, flag, "entities");
    }

    @Override
    public CompletableFuture<ChunkEntities<Entity>> a(ChunkCoordIntPair chunkcoordintpair) {
        return this.emptyChunks.contains(chunkcoordintpair.pair()) ? CompletableFuture.completedFuture(b(chunkcoordintpair)) : this.worker.b(chunkcoordintpair).thenApplyAsync((nbttagcompound) -> {
            if (nbttagcompound == null) {
                this.emptyChunks.add(chunkcoordintpair.pair());
                return b(chunkcoordintpair);
            } else {
                try {
                    ChunkCoordIntPair chunkcoordintpair1 = b(nbttagcompound);

                    if (!Objects.equals(chunkcoordintpair, chunkcoordintpair1)) {
                        EntityStorage.LOGGER.error("Chunk file at {} is in the wrong location. (Expected {}, got {})", chunkcoordintpair, chunkcoordintpair, chunkcoordintpair1);
                    }
                } catch (Exception exception) {
                    EntityStorage.LOGGER.warn("Failed to parse chunk {} position info", chunkcoordintpair, exception);
                }

                NBTTagCompound nbttagcompound1 = this.c(nbttagcompound);
                NBTTagList nbttaglist = nbttagcompound1.getList("Entities", 10);
                List<Entity> list = (List) EntityTypes.a((List) nbttaglist, (World) this.level).collect(ImmutableList.toImmutableList());

                return new ChunkEntities<>(chunkcoordintpair, list);
            }
        }, this.mainThreadExecutor);
    }

    private static ChunkCoordIntPair b(NBTTagCompound nbttagcompound) {
        int[] aint = nbttagcompound.getIntArray("Position");

        return new ChunkCoordIntPair(aint[0], aint[1]);
    }

    private static void a(NBTTagCompound nbttagcompound, ChunkCoordIntPair chunkcoordintpair) {
        nbttagcompound.set("Position", new NBTTagIntArray(new int[]{chunkcoordintpair.x, chunkcoordintpair.z}));
    }

    private static ChunkEntities<Entity> b(ChunkCoordIntPair chunkcoordintpair) {
        return new ChunkEntities<>(chunkcoordintpair, ImmutableList.of());
    }

    @Override
    public void a(ChunkEntities<Entity> chunkentities) {
        ChunkCoordIntPair chunkcoordintpair = chunkentities.a();

        if (chunkentities.c()) {
            if (this.emptyChunks.add(chunkcoordintpair.pair())) {
                this.worker.a(chunkcoordintpair, (NBTTagCompound) null);
            }

        } else {
            NBTTagList nbttaglist = new NBTTagList();

            chunkentities.b().forEach((entity) -> {
                NBTTagCompound nbttagcompound = new NBTTagCompound();

                if (entity.e(nbttagcompound)) {
                    nbttaglist.add(nbttagcompound);
                }

            });
            NBTTagCompound nbttagcompound = new NBTTagCompound();

            nbttagcompound.setInt("DataVersion", SharedConstants.getGameVersion().getWorldVersion());
            nbttagcompound.set("Entities", nbttaglist);
            a(nbttagcompound, chunkcoordintpair);
            this.worker.a(chunkcoordintpair, nbttagcompound).exceptionally((throwable) -> {
                EntityStorage.LOGGER.error("Failed to store chunk {}", chunkcoordintpair, throwable);
                return null;
            });
            this.emptyChunks.remove(chunkcoordintpair.pair());
        }
    }

    @Override
    public void a() {
        this.worker.a().join();
    }

    private NBTTagCompound c(NBTTagCompound nbttagcompound) {
        int i = a(nbttagcompound);

        return GameProfileSerializer.a(this.fixerUpper, DataFixTypes.ENTITY_CHUNK, nbttagcompound, i);
    }

    public static int a(NBTTagCompound nbttagcompound) {
        return nbttagcompound.hasKeyOfType("DataVersion", 99) ? nbttagcompound.getInt("DataVersion") : -1;
    }

    @Override
    public void close() throws IOException {
        this.worker.close();
    }
}