package eu.mhsl.minenet.minigames.util;

import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
import net.minestom.server.instance.InstanceContainer;
import net.minestom.server.instance.batch.Batch;
import net.minestom.server.instance.batch.ChunkBatch;
import net.minestom.server.utils.chunk.ChunkUtils;

import java.lang.reflect.Field;
import java.util.concurrent.CompletableFuture;

public class BatchUtil {
    public static long[] getAffectedChunks(Batch batch) {
        try {
            Field field = batch.getClass().getDeclaredField("chunkBatchesMap");
            field.setAccessible(true);

            @SuppressWarnings("unchecked")
            Long2ObjectMap<ChunkBatch> chunkBatchesMap = (Long2ObjectMap<ChunkBatch>) field.get(batch);

            return chunkBatchesMap.keySet().toLongArray();
        } catch (NoSuchFieldException | IllegalAccessException e) {
            throw new RuntimeException(e);
        }
    }

    public static void loadAndApplyBatch(Batch batch, InstanceContainer instance, Runnable onFinish) {
        batch.awaitReady();
        long[] affectedChunks = BatchUtil.getAffectedChunks(batch);
        ChunkUtils.optionalLoadAll(instance, affectedChunks, null).thenRun(() -> batch.apply(instance, onFinish));
    }

    public static void loadAndApplyBatchBlocking(Batch batch, InstanceContainer instance) {
        CompletableFuture<Void> future = new CompletableFuture<>();
        batch.awaitReady();
        long[] affectedChunks = BatchUtil.getAffectedChunks(batch);
        CompletableFuture<Void> loadChunksTask = ChunkUtils.optionalLoadAll(instance, affectedChunks, null);

        Runnable completerTask = () -> {
            System.out.println("COMPLETE");
            future.complete(null);
        };

        loadChunksTask.thenRun(() -> {
            System.out.println("BEGIN");
            batch.apply(instance, completerTask);
        });

        try {
            future.get();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
