package coffee.liz.ecs.grid

import coffee.liz.ecs.Entity
import coffee.liz.ecs.Rect
import coffee.liz.ecs.Vec2

data class Grid(
    val dimensions: Vec2,
    val grid: MutableList<MutableList<MutableSet<Entity>>> =
        MutableList(dimensions.y.toInt()) {
            MutableList(dimensions.x.toInt()) { mutableSetOf() }
        },
) {
    fun at(position: Vec2): MutableSet<Entity> {
        if (!Rect(Vec2(0f, 0f), dimensions).contains(position)) {
            throw IllegalArgumentException()
        }
        return grid[position.y.toInt()][position.x.toInt()]
    }

    fun add(entity: Entity) {
        if (!entity.has(GridPosition::class)) {
            throw IllegalArgumentException("Entity must have a GridPosition component to be added to the grid")
        }
        val position = entity.get(GridPosition::class).position
        at(position).add(entity)
    }

    fun reassignCellEntities(entities: Collection<Entity>) {
        val remainingEntities = entities.filter { it.has(GridPosition::class) }.toMutableSet()
        for (y in grid.indices) {
            for (x in grid[y].indices) {
                val vec = Vec2(x.toFloat(), y.toFloat())
                val entities = at(vec)
                entities.retainAll(
                    entities
                        .filter {
                            if (!it.has(GridPosition::class)) {
                                return@filter false
                            }
                            val currentPosition = it.get(GridPosition::class).position
                            (vec == currentPosition).also { stillInCell ->
                                if (!stillInCell) {
                                    at(currentPosition).add(it)
                                }
                            }
                        }.toSet(),
                )
                remainingEntities.removeAll(entities)
            }
        }
        remainingEntities.forEach { add(it) }
    }
}