from tools.aoc import AOCDay
from typing import Any
from tools.coordinate import Coordinate
from tools.grid import Grid

DIRECTIONS = [(0, -1), (1, 0), (0, 1), (-1, 0)]


class Day(AOCDay):
    inputs = [
        [
            (41, "input6_test"),
            (5208, "input6"),
        ],
        [
            (6, "input6_test"),
            (1972, "input6"),
        ],
    ]

    def parse_input(self) -> tuple[Grid, Coordinate]:
        grid = Grid.from_data(self.getInput(), translate={".": False, "#": True})
        start = list(grid.find("^"))[0]
        grid.set(start, False)

        return grid, start

    def walk_room(self, grid: Grid, start: Coordinate) -> bool:
        cur_pos = start
        dir_index = 0
        self.seen_reset()
        while grid.isWithinBoundaries(cur_pos):
            if self.seen((cur_pos, dir_index)):
                return True

            next_pos = cur_pos + DIRECTIONS[dir_index]
            while grid.get(next_pos):
                dir_index = (dir_index + 1) % len(DIRECTIONS)
                next_pos = cur_pos + DIRECTIONS[dir_index]

            cur_pos = next_pos

        return False

    def part1(self) -> Any:
        grid, cur_pos = self.parse_input()
        self.walk_room(grid, cur_pos)
        return len(set(x[0] for x in self.SEEN))

    def part2(self) -> Any:
        grid, start = self.parse_input()
        self.walk_room(grid, start)
        visits = set(x[0] for x in self.SEEN)

        loops = 0
        for obs_pos in visits:
            if obs_pos == start:
                continue

            grid.set(obs_pos, True)
            loops += self.walk_room(grid, start)
            grid.set(obs_pos, False)

        return loops


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