from collections import deque
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": "[]",
    ".": "..",
    "@": "@.",
}


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) == ".":
            start += dir_
            continue
        elif grid.get(next_pos) == "#":
            continue
        elif grid.get(next_pos) == "O":
            pos_behind_box = next_pos + dir_
            while grid.get(pos_behind_box) == "O":
                pos_behind_box += dir_

            if grid.get(pos_behind_box) == "#":
                continue

            grid.set(pos_behind_box, "O")
            start = next_pos
            grid.set(start, ".")
            continue

        if move in ["<", ">"]:
            # easy mode
            pos_behind_box = next_pos + dir_
            while grid.get(pos_behind_box) in ["[", "]"]:
                pos_behind_box += dir_ * 2

            if grid.get(pos_behind_box) == "#":
                continue

            while pos_behind_box != start:
                grid.set(pos_behind_box, grid.get(pos_behind_box - dir_))
                pos_behind_box -= dir_
        else:
            # recursive mode
            wall_in_the_way = False
            Q = deque([next_pos])
            boxes = set()
            while Q:
                pos = Q.popleft()
                if grid.get(pos) == "#":
                    wall_in_the_way = True
                    break
                elif grid.get(pos) == ".":
                    continue

                Q.append(pos + dir_)
                if grid.get(pos) == "[":
                    boxes.add((pos, pos + (1, 0)))
                    Q.append(pos + dir_ + (1, 0))
                else:
                    boxes.add((pos + (-1, 0), pos))
                    Q.append(pos + dir_ + (-1, 0))

            if wall_in_the_way:
                continue

            for box_l, box_r in boxes:
                grid.set(box_l, ".")
                grid.set(box_r, ".")

            for box_l, box_r in boxes:
                grid.set(box_l + dir_, "[")
                grid.set(box_r + dir_, "]")

        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=".")
        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)