from __future__ import annotations
from minigrid.core.grid import Grid
from minigrid.core.world_object import Wall

class MultiRoomGridColor(Grid):
    def __init__(self, width: int, height: int):
        super().__init__(width, height)

    def wall_rect(self, x: int, y: int, w: int, h: int, c: str | None = None):
        """
        Create a rectangle of walls. If `c` is None, use the default wall behavior.
        Otherwise, use the specified color `c` for the walls.
        """
        if c is None:
            # Default behavior (use default wall color)
            self.horz_wall(x, y, w)
            self.horz_wall(x, y + h - 1, w)
            self.vert_wall(x, y, h)
            self.vert_wall(x + w - 1, y, h)
        else:
            # Use the specified color `c` for the walls
            self.horz_wall(x, y, w, obj_type=lambda: Wall(c))
            self.horz_wall(x, y + h - 1, w, obj_type=lambda: Wall(c))
            self.vert_wall(x, y, h, obj_type=lambda: Wall(c))
            self.vert_wall(x + w - 1, y, h, obj_type=lambda: Wall(c))

