package com.tiedup.remake.cells;

import com.tiedup.remake.core.SystemMessageManager;
import com.tiedup.remake.core.TiedUpMod;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import net.minecraft.core.BlockPos;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.nbt.Tag;
import net.minecraft.server.MinecraftServer;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.world.level.ChunkPos;
import net.minecraft.world.level.saveddata.SavedData;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

/**
 * Global registry for Cell System V2 data.
 *
 * Named CellRegistryV2 to coexist with v1 CellRegistry during migration.
 * Uses "tiedup_cell_registry_v2" as the SavedData name.
 *
 * Provides spatial indices for fast lookups by wall position, interior position,
 * core position, chunk, and camp.
 */
public class CellRegistryV2 extends SavedData {

    private static final String DATA_NAME = "tiedup_cell_registry_v2";

    /** Reservation timeout in ticks (30 seconds = 600 ticks) */
    private static final long RESERVATION_TIMEOUT_TICKS = 600L;

    // ==================== RESERVATION ====================

    private static class CellReservation {

        private final UUID kidnapperUUID;
        private final long expiryTime;

        public CellReservation(UUID kidnapperUUID, long expiryTime) {
            this.kidnapperUUID = kidnapperUUID;
            this.expiryTime = expiryTime;
        }

        public UUID getKidnapperUUID() {
            return kidnapperUUID;
        }

        public boolean isExpired(long currentTime) {
            return currentTime >= expiryTime;
        }
    }

    // ==================== STORAGE ====================

    // Primary storage
    private final Map<UUID, CellDataV2> cells = new ConcurrentHashMap<>();

    // Indices (rebuilt on load)
    private final Map<BlockPos, UUID> wallToCell = new ConcurrentHashMap<>();
    private final Map<BlockPos, UUID> interiorToCell =
        new ConcurrentHashMap<>();
    private final Map<BlockPos, UUID> coreToCell = new ConcurrentHashMap<>();

    // Spatial + camp indices
    private final Map<ChunkPos, Set<UUID>> cellsByChunk =
        new ConcurrentHashMap<>();
    private final Map<UUID, Set<UUID>> cellsByCamp = new ConcurrentHashMap<>();

    // Breach tracking index (breached wall position → cell ID)
    private final Map<BlockPos, UUID> breachedToCell =
        new ConcurrentHashMap<>();

    // Reservations (not persisted)
    private final Map<UUID, CellReservation> reservations =
        new ConcurrentHashMap<>();

    // ==================== STATIC ACCESS ====================

    public static CellRegistryV2 get(ServerLevel level) {
        return level
            .getDataStorage()
            .computeIfAbsent(
                CellRegistryV2::load,
                CellRegistryV2::new,
                DATA_NAME
            );
    }

    public static CellRegistryV2 get(MinecraftServer server) {
        return get(server.overworld());
    }

    // ==================== CELL LIFECYCLE ====================

    /**
     * Create a new cell from a flood-fill result.
     */
    public CellDataV2 createCell(
        BlockPos corePos,
        FloodFillResult result,
        @Nullable UUID ownerId
    ) {
        CellDataV2 cell = new CellDataV2(corePos, result);
        if (ownerId != null) {
            cell.setOwnerId(ownerId);
        }

        cells.put(cell.getId(), cell);

        // Register in indices
        coreToCell.put(corePos.immutable(), cell.getId());
        for (BlockPos pos : cell.getWallBlocks()) {
            wallToCell.put(pos.immutable(), cell.getId());
        }
        for (BlockPos pos : cell.getInteriorBlocks()) {
            interiorToCell.put(pos.immutable(), cell.getId());
        }

        addToSpatialIndex(cell);
        setDirty();
        return cell;
    }

    /**
     * Register a pre-constructed CellDataV2 (used by migration and structure loading).
     * The cell must already have its ID and corePos set.
     */
    public void registerExistingCell(CellDataV2 cell) {
        cells.put(cell.getId(), cell);
        coreToCell.put(cell.getCorePos().immutable(), cell.getId());

        for (BlockPos pos : cell.getWallBlocks()) {
            wallToCell.put(pos.immutable(), cell.getId());
        }
        for (BlockPos pos : cell.getInteriorBlocks()) {
            interiorToCell.put(pos.immutable(), cell.getId());
        }

        addToSpatialIndex(cell);
        setDirty();
    }

    /**
     * Remove a cell from the registry and all indices.
     */
    public void removeCell(UUID cellId) {
        CellDataV2 cell = cells.remove(cellId);
        if (cell == null) return;

        coreToCell.remove(cell.getCorePos());
        for (BlockPos pos : cell.getWallBlocks()) {
            wallToCell.remove(pos);
        }
        for (BlockPos pos : cell.getBreachedPositions()) {
            breachedToCell.remove(pos);
        }
        for (BlockPos pos : cell.getInteriorBlocks()) {
            interiorToCell.remove(pos);
        }

        removeFromSpatialIndex(cell);
        reservations.remove(cellId);
        setDirty();
    }

    /**
     * Rescan a cell with a new flood-fill result.
     * Clears old indices and repopulates with new geometry.
     */
    public void rescanCell(UUID cellId, FloodFillResult newResult) {
        CellDataV2 cell = cells.get(cellId);
        if (cell == null) return;

        // Clear old indices for this cell
        for (BlockPos pos : cell.getWallBlocks()) {
            wallToCell.remove(pos);
        }
        for (BlockPos pos : cell.getBreachedPositions()) {
            wallToCell.remove(pos);
            breachedToCell.remove(pos);
        }
        for (BlockPos pos : cell.getInteriorBlocks()) {
            interiorToCell.remove(pos);
        }
        removeFromSpatialIndex(cell);

        // Update geometry
        cell.updateGeometry(newResult);

        // Rebuild indices
        for (BlockPos pos : cell.getWallBlocks()) {
            wallToCell.put(pos.immutable(), cellId);
        }
        for (BlockPos pos : cell.getInteriorBlocks()) {
            interiorToCell.put(pos.immutable(), cellId);
        }
        addToSpatialIndex(cell);

        setDirty();
    }

    // ==================== QUERIES ====================

    @Nullable
    public CellDataV2 getCell(UUID cellId) {
        return cells.get(cellId);
    }

    @Nullable
    public CellDataV2 getCellAtCore(BlockPos corePos) {
        UUID cellId = coreToCell.get(corePos);
        return cellId != null ? cells.get(cellId) : null;
    }

    @Nullable
    public CellDataV2 getCellContaining(BlockPos pos) {
        UUID cellId = interiorToCell.get(pos);
        return cellId != null ? cells.get(cellId) : null;
    }

    @Nullable
    public CellDataV2 getCellByWall(BlockPos pos) {
        UUID cellId = wallToCell.get(pos);
        return cellId != null ? cells.get(cellId) : null;
    }

    @Nullable
    public UUID getCellIdAtWall(BlockPos pos) {
        return wallToCell.get(pos);
    }

    public boolean isInsideAnyCell(BlockPos pos) {
        return interiorToCell.containsKey(pos);
    }

    public Collection<CellDataV2> getAllCells() {
        return Collections.unmodifiableCollection(cells.values());
    }

    public int getCellCount() {
        return cells.size();
    }

    public List<CellDataV2> getCellsByCamp(UUID campId) {
        Set<UUID> cellIds = cellsByCamp.get(campId);
        if (cellIds == null) return Collections.emptyList();

        List<CellDataV2> result = new ArrayList<>();
        for (UUID cellId : cellIds) {
            CellDataV2 cell = cells.get(cellId);
            if (cell != null) {
                result.add(cell);
            }
        }
        return result;
    }

    public List<CellDataV2> findCellsNear(BlockPos center, double radius) {
        List<CellDataV2> nearby = new ArrayList<>();
        double radiusSq = radius * radius;

        int chunkRadius = (int) Math.ceil(radius / 16.0) + 1;
        ChunkPos centerChunk = new ChunkPos(center);

        for (int dx = -chunkRadius; dx <= chunkRadius; dx++) {
            for (int dz = -chunkRadius; dz <= chunkRadius; dz++) {
                ChunkPos checkChunk = new ChunkPos(
                    centerChunk.x + dx,
                    centerChunk.z + dz
                );
                Set<UUID> cellsInChunk = cellsByChunk.get(checkChunk);

                if (cellsInChunk != null) {
                    for (UUID cellId : cellsInChunk) {
                        CellDataV2 cell = cells.get(cellId);
                        if (
                            cell != null &&
                            cell.getCorePos().distSqr(center) <= radiusSq
                        ) {
                            nearby.add(cell);
                        }
                    }
                }
            }
        }
        return nearby;
    }

    @Nullable
    public CellDataV2 findCellByPrisoner(UUID prisonerId) {
        for (CellDataV2 cell : cells.values()) {
            if (cell.hasPrisoner(prisonerId)) {
                return cell;
            }
        }
        return null;
    }

    @Nullable
    public CellDataV2 getCellByName(String name) {
        if (name == null || name.isEmpty()) return null;
        for (CellDataV2 cell : cells.values()) {
            if (name.equals(cell.getName())) {
                return cell;
            }
        }
        return null;
    }

    public List<CellDataV2> getCellsByOwner(UUID ownerId) {
        if (ownerId == null) return Collections.emptyList();
        return cells
            .values()
            .stream()
            .filter(c -> ownerId.equals(c.getOwnerId()))
            .collect(Collectors.toList());
    }

    public int getCellCountOwnedBy(UUID ownerId) {
        if (ownerId == null) return 0;
        return (int) cells
            .values()
            .stream()
            .filter(c -> ownerId.equals(c.getOwnerId()))
            .count();
    }

    @Nullable
    public UUID getNextCellId(@Nullable UUID currentId) {
        if (cells.isEmpty()) return null;
        List<UUID> ids = new ArrayList<>(cells.keySet());
        if (currentId == null) return ids.get(0);
        int index = ids.indexOf(currentId);
        if (index < 0 || index >= ids.size() - 1) return ids.get(0);
        return ids.get(index + 1);
    }

    // ==================== CAMP QUERIES ====================

    public List<UUID> getPrisonersInCamp(UUID campId) {
        if (campId == null) return Collections.emptyList();
        return getCellsByCamp(campId)
            .stream()
            .flatMap(cell -> cell.getPrisonerIds().stream())
            .collect(Collectors.toList());
    }

    public int getPrisonerCountInCamp(UUID campId) {
        if (campId == null) return 0;
        return getCellsByCamp(campId)
            .stream()
            .mapToInt(CellDataV2::getPrisonerCount)
            .sum();
    }

    @Nullable
    public CellDataV2 findAvailableCellInCamp(UUID campId) {
        if (campId == null) return null;
        for (CellDataV2 cell : getCellsByCamp(campId)) {
            if (!cell.isFull()) {
                return cell;
            }
        }
        return null;
    }

    public boolean hasCampCells(UUID campId) {
        if (campId == null) return false;
        Set<UUID> cellIds = cellsByCamp.get(campId);
        return cellIds != null && !cellIds.isEmpty();
    }

    /**
     * Update the camp index for a cell after ownership change.
     * Removes from old camp index, adds to new if camp-owned.
     */
    public void updateCampIndex(CellDataV2 cell, @Nullable UUID oldOwnerId) {
        if (oldOwnerId != null) {
            Set<UUID> oldCampCells = cellsByCamp.get(oldOwnerId);
            if (oldCampCells != null) {
                oldCampCells.remove(cell.getId());
                if (oldCampCells.isEmpty()) {
                    cellsByCamp.remove(oldOwnerId);
                }
            }
        }

        if (cell.isCampOwned() && cell.getOwnerId() != null) {
            cellsByCamp
                .computeIfAbsent(cell.getOwnerId(), k ->
                    ConcurrentHashMap.newKeySet()
                )
                .add(cell.getId());
        }

        setDirty();
    }

    // ==================== PRISONER MANAGEMENT ====================

    public synchronized boolean assignPrisoner(UUID cellId, UUID prisonerId) {
        CellDataV2 cell = cells.get(cellId);
        if (cell == null || cell.isFull()) return false;

        // Ensure prisoner uniqueness across cells
        CellDataV2 existingCell = findCellByPrisoner(prisonerId);
        if (existingCell != null) {
            if (existingCell.getId().equals(cellId)) {
                return true; // Already in this cell
            }
            return false; // Already in another cell
        }

        if (cell.addPrisoner(prisonerId)) {
            setDirty();
            return true;
        }
        return false;
    }

    public boolean releasePrisoner(
        UUID cellId,
        UUID prisonerId,
        MinecraftServer server
    ) {
        CellDataV2 cell = cells.get(cellId);
        if (cell == null) return false;

        if (cell.removePrisoner(prisonerId)) {
            // Synchronize with PrisonerManager
            if (cell.isCampOwned() && cell.getCampId() != null) {
                com.tiedup.remake.prison.PrisonerManager manager =
                    com.tiedup.remake.prison.PrisonerManager.get(
                        server.overworld()
                    );
                com.tiedup.remake.prison.PrisonerState currentState =
                    manager.getState(prisonerId);

                boolean isBeingExtracted =
                    currentState ==
                    com.tiedup.remake.prison.PrisonerState.WORKING;

                if (
                    !isBeingExtracted &&
                    currentState ==
                    com.tiedup.remake.prison.PrisonerState.IMPRISONED
                ) {
                    manager.release(
                        prisonerId,
                        server.overworld().getGameTime()
                    );
                }
            }

            setDirty();
            return true;
        }
        return false;
    }

    public boolean assignPrisonerWithNotification(
        UUID cellId,
        UUID prisonerId,
        MinecraftServer server,
        String prisonerName
    ) {
        CellDataV2 cell = cells.get(cellId);
        if (cell == null || cell.isFull()) return false;

        if (cell.addPrisoner(prisonerId)) {
            setDirty();
            if (cell.hasOwner()) {
                notifyOwner(
                    server,
                    cell.getOwnerId(),
                    SystemMessageManager.MessageCategory.PRISONER_ARRIVED,
                    prisonerName
                );
            }
            return true;
        }
        return false;
    }

    public boolean releasePrisonerWithNotification(
        UUID cellId,
        UUID prisonerId,
        MinecraftServer server,
        String prisonerName,
        boolean escaped
    ) {
        CellDataV2 cell = cells.get(cellId);
        if (cell == null) return false;

        if (cell.removePrisoner(prisonerId)) {
            // Synchronize with PrisonerManager
            if (cell.isCampOwned() && cell.getCampId() != null) {
                com.tiedup.remake.prison.PrisonerManager manager =
                    com.tiedup.remake.prison.PrisonerManager.get(
                        server.overworld()
                    );
                com.tiedup.remake.prison.PrisonerState currentState =
                    manager.getState(prisonerId);

                if (
                    currentState ==
                    com.tiedup.remake.prison.PrisonerState.IMPRISONED
                ) {
                    manager.release(
                        prisonerId,
                        server.overworld().getGameTime()
                    );
                }
            }

            setDirty();

            if (cell.hasOwner()) {
                SystemMessageManager.MessageCategory category = escaped
                    ? SystemMessageManager.MessageCategory.PRISONER_ESCAPED
                    : SystemMessageManager.MessageCategory.PRISONER_RELEASED;
                notifyOwner(server, cell.getOwnerId(), category, prisonerName);
            }
            return true;
        }
        return false;
    }

    public int releasePrisonerFromAllCells(UUID prisonerId) {
        int count = 0;
        for (CellDataV2 cell : cells.values()) {
            if (cell.removePrisoner(prisonerId)) {
                count++;
            }
        }
        if (count > 0) {
            setDirty();
        }
        return count;
    }

    /** Offline timeout for cleanup: 30 minutes */
    private static final long OFFLINE_TIMEOUT_MS = 30 * 60 * 1000L;

    public int cleanupEscapedPrisoners(
        ServerLevel level,
        com.tiedup.remake.state.CollarRegistry collarRegistry,
        double maxDistance
    ) {
        int removed = 0;

        for (CellDataV2 cell : cells.values()) {
            List<UUID> toRemove = new ArrayList<>();

            for (UUID prisonerId : cell.getPrisonerIds()) {
                boolean shouldRemove = false;
                String reason = null;

                ServerPlayer prisoner = level
                    .getServer()
                    .getPlayerList()
                    .getPlayer(prisonerId);

                if (prisoner == null) {
                    Long timestamp = cell.getPrisonerTimestamp(prisonerId);
                    long ts =
                        timestamp != null
                            ? timestamp
                            : System.currentTimeMillis();
                    long offlineDuration = System.currentTimeMillis() - ts;

                    if (offlineDuration > OFFLINE_TIMEOUT_MS) {
                        shouldRemove = true;
                        reason =
                            "offline for too long (" +
                            (offlineDuration / 60000) +
                            " minutes)";
                    } else {
                        continue;
                    }
                } else {
                    // Use corePos for distance check (V2 uses core position, not spawnPoint)
                    double distSq = prisoner
                        .blockPosition()
                        .distSqr(cell.getCorePos());
                    if (distSq > maxDistance * maxDistance) {
                        shouldRemove = true;
                        reason =
                            "too far from cell (" +
                            (int) Math.sqrt(distSq) +
                            " blocks)";
                    }

                    if (
                        !shouldRemove && !collarRegistry.hasOwners(prisonerId)
                    ) {
                        shouldRemove = true;
                        reason = "no collar registered";
                    }

                    if (!shouldRemove) {
                        com.tiedup.remake.state.IBondageState state =
                            com.tiedup.remake.util.KidnappedHelper.getKidnappedState(
                                prisoner
                            );
                        if (state == null || !state.isCaptive()) {
                            shouldRemove = true;
                            reason = "no longer captive";
                        }
                    }
                }

                if (shouldRemove) {
                    toRemove.add(prisonerId);
                    TiedUpMod.LOGGER.info(
                        "[CellRegistryV2] Removing escaped prisoner {} from cell {} - reason: {}",
                        prisonerId.toString().substring(0, 8),
                        cell.getId().toString().substring(0, 8),
                        reason
                    );
                }
            }

            for (UUID id : toRemove) {
                cell.removePrisoner(id);

                if (cell.isCampOwned() && cell.getCampId() != null) {
                    com.tiedup.remake.prison.PrisonerManager manager =
                        com.tiedup.remake.prison.PrisonerManager.get(
                            level.getServer().overworld()
                        );
                    com.tiedup.remake.prison.PrisonerState currentState =
                        manager.getState(id);

                    if (
                        currentState ==
                        com.tiedup.remake.prison.PrisonerState.IMPRISONED
                    ) {
                        com.tiedup.remake.prison.service.PrisonerService.get().escape(
                            level,
                            id,
                            "offline_cleanup"
                        );
                    }
                }

                removed++;
            }
        }

        if (removed > 0) {
            setDirty();
        }

        return removed;
    }

    // ==================== NOTIFICATIONS ====================

    private void notifyOwner(
        MinecraftServer server,
        UUID ownerId,
        SystemMessageManager.MessageCategory category,
        String prisonerName
    ) {
        if (server == null || ownerId == null) return;
        ServerPlayer owner = server.getPlayerList().getPlayer(ownerId);
        if (owner != null) {
            String template = SystemMessageManager.getTemplate(category);
            String formattedMessage = String.format(template, prisonerName);
            SystemMessageManager.sendToPlayer(
                owner,
                category,
                formattedMessage
            );
        }
    }

    // ==================== BREACH MANAGEMENT ====================

    /**
     * Record a wall breach: updates CellDataV2 and indices atomically.
     */
    public void addBreach(UUID cellId, BlockPos pos) {
        CellDataV2 cell = cells.get(cellId);
        if (cell == null) return;

        cell.addBreach(pos);
        wallToCell.remove(pos);
        breachedToCell.put(pos.immutable(), cellId);
        setDirty();
    }

    /**
     * Repair a wall breach: updates CellDataV2 and indices atomically.
     */
    public void repairBreach(UUID cellId, BlockPos pos) {
        CellDataV2 cell = cells.get(cellId);
        if (cell == null) return;

        cell.repairBreach(pos);
        breachedToCell.remove(pos);
        wallToCell.put(pos.immutable(), cellId);
        setDirty();
    }

    /**
     * Get the cell ID for a breached wall position.
     */
    @Nullable
    public UUID getCellIdAtBreach(BlockPos pos) {
        return breachedToCell.get(pos);
    }

    // ==================== RESERVATIONS ====================

    public boolean reserveCell(UUID cellId, UUID kidnapperUUID, long gameTime) {
        cleanupExpiredReservations(gameTime);

        long expiryTime = gameTime + RESERVATION_TIMEOUT_TICKS;
        CellReservation existing = reservations.get(cellId);

        if (existing != null) {
            if (existing.getKidnapperUUID().equals(kidnapperUUID)) {
                reservations.put(
                    cellId,
                    new CellReservation(kidnapperUUID, expiryTime)
                );
                return true;
            }
            if (!existing.isExpired(gameTime)) {
                return false;
            }
        }

        reservations.put(
            cellId,
            new CellReservation(kidnapperUUID, expiryTime)
        );
        return true;
    }

    public boolean consumeReservation(UUID cellId, UUID kidnapperUUID) {
        CellReservation reservation = reservations.remove(cellId);
        if (reservation == null) return false;
        return reservation.getKidnapperUUID().equals(kidnapperUUID);
    }

    public boolean isReservedByOther(
        UUID cellId,
        @Nullable UUID kidnapperUUID,
        long gameTime
    ) {
        CellReservation reservation = reservations.get(cellId);
        if (reservation == null) return false;
        if (reservation.isExpired(gameTime)) {
            reservations.remove(cellId);
            return false;
        }
        return (
            kidnapperUUID == null ||
            !reservation.getKidnapperUUID().equals(kidnapperUUID)
        );
    }

    public void cancelReservation(UUID cellId, UUID kidnapperUUID) {
        CellReservation reservation = reservations.get(cellId);
        if (
            reservation != null &&
            reservation.getKidnapperUUID().equals(kidnapperUUID)
        ) {
            reservations.remove(cellId);
        }
    }

    private void cleanupExpiredReservations(long gameTime) {
        reservations
            .entrySet()
            .removeIf(entry -> entry.getValue().isExpired(gameTime));
    }

    // ==================== SPATIAL INDEX ====================

    private void addToSpatialIndex(CellDataV2 cell) {
        ChunkPos chunkPos = new ChunkPos(cell.getCorePos());
        cellsByChunk
            .computeIfAbsent(chunkPos, k -> ConcurrentHashMap.newKeySet())
            .add(cell.getId());

        // Add to camp index if camp-owned
        if (
            cell.getOwnerType() == CellOwnerType.CAMP &&
            cell.getOwnerId() != null
        ) {
            cellsByCamp
                .computeIfAbsent(cell.getOwnerId(), k ->
                    ConcurrentHashMap.newKeySet()
                )
                .add(cell.getId());
        }
    }

    private void removeFromSpatialIndex(CellDataV2 cell) {
        ChunkPos chunkPos = new ChunkPos(cell.getCorePos());
        Set<UUID> cellsInChunk = cellsByChunk.get(chunkPos);
        if (cellsInChunk != null) {
            cellsInChunk.remove(cell.getId());
            if (cellsInChunk.isEmpty()) {
                cellsByChunk.remove(chunkPos);
            }
        }

        if (
            cell.getOwnerType() == CellOwnerType.CAMP &&
            cell.getOwnerId() != null
        ) {
            Set<UUID> cellsInCamp = cellsByCamp.get(cell.getOwnerId());
            if (cellsInCamp != null) {
                cellsInCamp.remove(cell.getId());
                if (cellsInCamp.isEmpty()) {
                    cellsByCamp.remove(cell.getOwnerId());
                }
            }
        }
    }

    // ==================== INDEX REBUILD ====================

    private void rebuildIndices() {
        wallToCell.clear();
        interiorToCell.clear();
        coreToCell.clear();
        breachedToCell.clear();
        cellsByChunk.clear();
        cellsByCamp.clear();

        for (CellDataV2 cell : cells.values()) {
            coreToCell.put(cell.getCorePos(), cell.getId());
            for (BlockPos pos : cell.getWallBlocks()) {
                wallToCell.put(pos, cell.getId());
            }
            for (BlockPos pos : cell.getBreachedPositions()) {
                breachedToCell.put(pos, cell.getId());
            }
            for (BlockPos pos : cell.getInteriorBlocks()) {
                interiorToCell.put(pos, cell.getId());
            }
            addToSpatialIndex(cell);
        }
    }

    // ==================== PERSISTENCE ====================

    @Override
    public @NotNull CompoundTag save(@NotNull CompoundTag tag) {
        ListTag cellList = new ListTag();
        for (CellDataV2 cell : cells.values()) {
            cellList.add(cell.save());
        }
        tag.put("cells", cellList);
        return tag;
    }

    public static CellRegistryV2 load(CompoundTag tag) {
        CellRegistryV2 registry = new CellRegistryV2();

        if (tag.contains("cells")) {
            ListTag cellList = tag.getList("cells", Tag.TAG_COMPOUND);
            for (int i = 0; i < cellList.size(); i++) {
                CellDataV2 cell = CellDataV2.load(cellList.getCompound(i));
                if (cell != null) {
                    registry.cells.put(cell.getId(), cell);
                }
            }
        }

        registry.rebuildIndices();
        return registry;
    }

    // ==================== DEBUG ====================

    public String toDebugString() {
        StringBuilder sb = new StringBuilder();
        sb.append("CellRegistryV2:\n");
        sb.append("  Total cells: ").append(cells.size()).append("\n");
        sb.append("  Wall index: ").append(wallToCell.size()).append("\n");
        sb
            .append("  Interior index: ")
            .append(interiorToCell.size())
            .append("\n");
        sb.append("  Core index: ").append(coreToCell.size()).append("\n");

        for (CellDataV2 cell : cells.values()) {
            sb.append("  ").append(cell.toString()).append("\n");
        }
        return sb.toString();
    }
}