package com.thebrokenrail.twine.component;

import com.thebrokenrail.twine.Twine;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.nbt.Tag;
import net.minecraft.server.world.ServerWorld;
import net.minecraft.util.math.ChunkPos;
import net.minecraft.world.PersistentState;

import java.util.*;

public class StageDataComponent extends PersistentState {
    public StageDataComponent(String key) {
        super(key);
    }

    public static class StageData {
        public int x;
        public int z;
        public long time = 0;
    }

    private final Map<UUID, StageData[]> data = new HashMap<>();

    public StageData[] getData(UUID uuid) {
        return data.computeIfAbsent(uuid, k -> new StageData[Twine.STAGE_COUNT]);
    }

    private static final int CHUNK_RADIUS = 16;
    private static final String STAGE_DATA_ID = "twine_stage_data";

    public static int findStageOfChunk(ChunkPos pos, StageDataComponent.StageData[] data) {
        for (int i = Twine.STAGE_COUNT - 1; i >= 0; i--) {
            StageDataComponent.StageData stage = data[i];
            if (stage != null && Math.abs(stage.x - pos.x) <= CHUNK_RADIUS && Math.abs(stage.z - pos.z) <= CHUNK_RADIUS) {
                return i;
            }
        }
        data[0] = new StageDataComponent.StageData();
        data[0].x = pos.x;
        data[0].z = pos.z;
        return 0;
    }

    public int findStageOfChunk(ChunkPos pos) {
        Collection<StageData[]> values = data.values();
        int stage = 0;
        for (StageData[] value : values) {
            stage = Math.max(stage, findStageOfChunk(pos, value));
        }
        return stage;
    }

    public static StageDataComponent getFromWorld(ServerWorld world) {
        return world.getPersistentStateManager().getOrCreate(() -> new StageDataComponent(STAGE_DATA_ID), STAGE_DATA_ID);
    }

    @Override
    public void fromTag(CompoundTag compoundTag) {
        data.clear();
        Tag uuidTag = compoundTag.get("Data");
        if (uuidTag instanceof ListTag) {
            for (int i = 0 ; i < ((ListTag) uuidTag).size(); i++) {
                CompoundTag obj = ((ListTag) uuidTag).getCompound(i);
                UUID uuid = obj.getUuid("UUID");

                Tag tag = obj.get("Stages");
                StageData[] newData = new StageData[Twine.STAGE_COUNT];
                if (tag instanceof ListTag) {
                    for (int k = 0; k < ((ListTag) tag).size(); k++) {
                        Tag newTag = ((ListTag) tag).get(k);
                        if (newTag instanceof CompoundTag) {
                            StageData stage = new StageData();
                            stage.x = ((CompoundTag) newTag).getInt("X");
                            stage.z = ((CompoundTag) newTag).getInt("Z");
                            stage.time = ((CompoundTag) newTag).getLong("Time");
                            newData[((CompoundTag) newTag).getInt("Index")] = stage;
                        }
                    }
                    data.put(uuid, newData);
                }
            }
        }
    }

    @Override
    public CompoundTag toTag(CompoundTag compoundTag) {
        ListTag list = new ListTag();
        for (Map.Entry<UUID, StageData[]> entry : data.entrySet()) {
            CompoundTag obj = new CompoundTag();
            obj.putUuid("UUID", entry.getKey());
            ListTag newList = new ListTag();
            for (int i = 0; i < Twine.STAGE_COUNT; i++) {
                if (entry.getValue()[i] != null) {
                    CompoundTag tag = new CompoundTag();
                    tag.putInt("X", entry.getValue()[i].x);
                    tag.putInt("Z", entry.getValue()[i].z);
                    tag.putLong("Time", entry.getValue()[i].time);
                    tag.putInt("Index", i);
                    newList.add(tag);
                }
            }
            obj.put("Stages", newList);
            list.add(obj);
        }
        compoundTag.put("Data", list);
        return compoundTag;
    }
}