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

import org.apache.logging.log4j.LogManager;
import net.minecraft.util.thread.ProcessorHandle;
import java.util.function.BiConsumer;
import net.minecraft.util.Unit;
import java.util.Iterator;
import java.util.function.Supplier;
import javax.annotation.Nullable;
import java.util.concurrent.CompletionException;
import java.io.IOException;
import java.util.concurrent.CompletionStage;
import java.util.function.Function;
import com.mojang.datafixers.util.Either;
import java.util.concurrent.CompletableFuture;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.Util;
import com.google.common.collect.Maps;
import java.io.File;
import net.minecraft.world.level.ChunkPos;
import java.util.Map;
import net.minecraft.util.thread.StrictQueue;
import net.minecraft.util.thread.ProcessorMailbox;
import java.util.concurrent.atomic.AtomicBoolean;
import org.apache.logging.log4j.Logger;

public class IOWorker implements AutoCloseable {
    private static final Logger LOGGER;
    private final AtomicBoolean shutdownRequested;
    private final ProcessorMailbox<StrictQueue.IntRunnable> mailbox;
    private final RegionFileStorage storage;
    private final Map<ChunkPos, PendingStore> pendingWrites;
    
    protected IOWorker(final File file, final boolean boolean2, final String string) {
        this.shutdownRequested = new AtomicBoolean();
        this.pendingWrites = Maps.newLinkedHashMap();
        this.storage = new RegionFileStorage(file, boolean2);
        this.mailbox = new ProcessorMailbox<StrictQueue.IntRunnable>(new StrictQueue.FixedPriorityQueue(Priority.values().length), Util.ioPool(), "IOWorker-" + string);
    }
    
    public CompletableFuture<Void> store(final ChunkPos bph, final CompoundTag le) {
        final PendingStore a4;
        return this.submitTask(() -> {
            a4 = this.pendingWrites.computeIfAbsent(bph, bph -> new PendingStore(le));
            a4.data = le;
            return (com.mojang.datafixers.util.Either<Object, Exception>)Either.left(a4.result);
        }).<Void>thenCompose(Function.identity());
    }
    
    @Nullable
    public CompoundTag load(final ChunkPos bph) throws IOException {
        final PendingStore a3;
        CompoundTag le4;
        final CompletableFuture<CompoundTag> completableFuture3 = this.<CompoundTag>submitTask(() -> {
            a3 = this.pendingWrites.get(bph);
            if (a3 != null) {
                return Either.left(a3.data);
            }
            else {
                try {
                    le4 = this.storage.read(bph);
                    return Either.left(le4);
                }
                catch (Exception exception4) {
                    IOWorker.LOGGER.warn("Failed to read chunk {}", bph, exception4);
                    return Either.right(exception4);
                }
            }
        });
        try {
            return completableFuture3.join();
        }
        catch (CompletionException completionException4) {
            if (completionException4.getCause() instanceof IOException) {
                throw (IOException)completionException4.getCause();
            }
            throw completionException4;
        }
    }
    
    public CompletableFuture<Void> synchronize() {
        final CompletableFuture<Void> completableFuture2 = this.submitTask(() -> Either.left(CompletableFuture.allOf(this.pendingWrites.values().stream().map(a -> a.result).<CompletableFuture<?>>toArray(CompletableFuture[]::new)))).<Void>thenCompose(Function.identity());
        return completableFuture2.<Void>thenCompose(void1 -> this.submitTask(() -> {
            try {
                this.storage.flush();
                return Either.left(null);
            }
            catch (Exception exception2) {
                IOWorker.LOGGER.warn("Failed to synchronized chunks", (Throwable)exception2);
                return Either.right(exception2);
            }
        }));
    }
    
    private <T> CompletableFuture<T> submitTask(final Supplier<Either<T, Exception>> supplier) {
        return this.mailbox.<T>askEither(amp -> new StrictQueue.IntRunnable(Priority.HIGH.ordinal(), () -> {
            if (!this.shutdownRequested.get()) {
                amp.tell(supplier.get());
            }
            this.tellStorePending();
        }));
    }
    
    private void storePendingChunk() {
        final Iterator<Map.Entry<ChunkPos, PendingStore>> iterator2 = this.pendingWrites.entrySet().iterator();
        if (!iterator2.hasNext()) {
            return;
        }
        final Map.Entry<ChunkPos, PendingStore> entry3 = iterator2.next();
        iterator2.remove();
        this.runStore(entry3.getKey(), entry3.getValue());
        this.tellStorePending();
    }
    
    private void tellStorePending() {
        this.mailbox.tell(new StrictQueue.IntRunnable(Priority.LOW.ordinal(), this::storePendingChunk));
    }
    
    private void runStore(final ChunkPos bph, final PendingStore a) {
        try {
            this.storage.write(bph, a.data);
            a.result.complete(null);
        }
        catch (Exception exception4) {
            IOWorker.LOGGER.error("Failed to store chunk {}", bph, exception4);
            a.result.completeExceptionally(exception4);
        }
    }
    
    @Override
    public void close() throws IOException {
        if (!this.shutdownRequested.compareAndSet(false, true)) {
            return;
        }
        final CompletableFuture<Unit> completableFuture2 = this.mailbox.<Unit>ask(amp -> new StrictQueue.IntRunnable(Priority.HIGH.ordinal(), () -> amp.tell(Unit.INSTANCE)));
        try {
            completableFuture2.join();
        }
        catch (CompletionException completionException3) {
            if (completionException3.getCause() instanceof IOException) {
                throw (IOException)completionException3.getCause();
            }
            throw completionException3;
        }
        this.mailbox.close();
        this.pendingWrites.forEach(this::runStore);
        this.pendingWrites.clear();
        try {
            this.storage.close();
        }
        catch (Exception exception3) {
            IOWorker.LOGGER.error("Failed to close storage", (Throwable)exception3);
        }
    }
    
    static {
        LOGGER = LogManager.getLogger();
    }
    
    enum Priority {
        HIGH, 
        LOW;
    }
    
    static class PendingStore {
        private CompoundTag data;
        private final CompletableFuture<Void> result;
        
        public PendingStore(final CompoundTag le) {
            this.result = new CompletableFuture<Void>();
            this.data = le;
        }
    }
}
