from collections import deque
from enum import Enum
from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any


MOVES = {
    "^": (0, -1),
    ">": (1, 0),
    "v": (0, 1),
    "<": (-1, 0),
}

LARGER = {
    "#": "##",
    "O": "[]",
    ".": "..",
    "@": "@.",
}


class Tile(str, Enum):
    EMPTY = "."
    WALL = "#"
    BOX = "O"
    BOX_LEFT = "["
    BOX_RIGHT = "]"


def sokoban(grid: Grid, start: Coordinate, moves: str) -> Grid:
    for move in moves:
        dir_ = MOVES[move]
        next_pos = start + dir_
        if grid.get(next_pos) == Tile.WALL:
            continue
        elif grid.get(next_pos) == Tile.EMPTY:
            start = next_pos
            continue

        # neither a wall, nor empty space? There must be a box in the way
        to_move = set()
        wall_in_the_way = False
        queue = deque([next_pos])
        while queue:
            pos = queue.popleft()
            pos_tile = grid.get(pos)
            if pos_tile == Tile.EMPTY:
                continue

            if pos_tile == Tile.WALL:
                wall_in_the_way = True
                break

            if (pos, pos_tile) in to_move:
                continue

            to_move.add((pos, pos_tile))
            queue.append(pos + dir_)
            if pos_tile in [Tile.BOX_LEFT, Tile.BOX_RIGHT]:
                if pos_tile == Tile.BOX_LEFT:
                    queue.append(pos + (1, 0))
                else:
                    queue.append(pos + (-1, 0))

        if wall_in_the_way:
            continue

        for box, _ in to_move:
            grid.set(box, Tile.EMPTY)

        for box, box_tile in to_move:
            grid.set(box + dir_, box_tile)

        start = next_pos

    return grid


class Day(AOCDay):
    inputs = [
        [
            (2028, "input15_test2"),
            (10092, "input15_test"),
            (1515788, "input15"),
        ],
        [
            (9021, "input15_test"),
            (1516544, "input15"),
        ],
    ]

    def parse_input(self, part2: bool = False) -> tuple[Grid, Coordinate, str]:
        map_, moves = self.getMultiLineInputAsArray()
        if part2:
            map_ = ["".join([LARGER[x] for x in y]) for y in map_]
        grid = Grid.from_data(map_, default=Tile.EMPTY, translate={r"\.#O\[\]": Tile})
        start = list(grid.find("@"))[0]
        grid.set(start, ".")

        return grid, start, "".join(moves)

    def part1(self) -> Any:
        grid = sokoban(*self.parse_input())
        return sum(100 * c.y + c.x for c in grid.find("O"))

    def part2(self) -> Any:
        grid = sokoban(*self.parse_input(True))
        return sum(100 * c.y + c.x for c in grid.find("["))


if __name__ == "__main__":
    day = Day(2024, 15)
    day.run(verbose=True)