class SpatialGrid {
    constructor(cellSize = 10) {
        this.cellSize = cellSize;
        this.buckets = new Map(); // Key: "gridX,gridY" -> Set of entities
    }

    _getKey(x, y) {
        const cx = Math.floor(x / this.cellSize);
        const cy = Math.floor(y / this.cellSize);
        return `${cx},${cy}`;
    }

    add(entity) {
        // Entity must have gridX and gridY
        const key = this._getKey(entity.gridX, entity.gridY);
        if (!this.buckets.has(key)) {
            this.buckets.set(key, new Set());
        }
        this.buckets.get(key).add(entity);
        entity._spatialKey = key;
    }

    updateEntity(entity) {
        const oldKey = entity._spatialKey;
        const newKey = this._getKey(entity.gridX, entity.gridY);

        if (oldKey !== newKey) {
            if (oldKey && this.buckets.has(oldKey)) {
                this.buckets.get(oldKey).delete(entity);
                if (this.buckets.get(oldKey).size === 0) {
                    this.buckets.delete(oldKey);
                }
            }
            this.add(entity);
        }
    }

    remove(entity) {
        const key = entity._spatialKey;
        if (key && this.buckets.has(key)) {
            this.buckets.get(key).delete(entity);
            if (this.buckets.get(key).size === 0) {
                this.buckets.delete(key);
            }
        }
        entity._spatialKey = null;
    }

    // Najdi entitete v bližini (v sosednjih bucketih)
    query(x, y) {
        const results = [];
        const cx = Math.floor(x / this.cellSize);
        const cy = Math.floor(y / this.cellSize);

        // Preveri 3x3 sosednjih bucketov (center + 8 sosedov)
        for (let dy = -1; dy <= 1; dy++) {
            for (let dx = -1; dx <= 1; dx++) {
                const key = `${cx + dx},${cy + dy}`;
                if (this.buckets.has(key)) {
                    for (const ent of this.buckets.get(key)) {
                        results.push(ent);
                    }
                }
            }
        }
        return results;
    }
}