from typing import NamedTuple

test_input = """
....#.....
.........#
..........
..#.......
.......#..
..........
.#..^.....
........#.
#.........
......#...
""".strip()

test_solution_p1 = 41
test_solution_p2 = 6


def solve_p1(puzzle_input: str) -> int:
    map_info = _parse_map(puzzle_input)
    traversal_analysis = _traverse_map(map_info)
    unique_locations = set((a, b) for a, b, _ in traversal_analysis.visited)
    return len(unique_locations)


def solve_p2(puzzle_input: str) -> int:
    map_info = _parse_map(puzzle_input)
    added_obstacles_candidates = set()

    for i, line in enumerate(map_info.map):
        for j, cell in enumerate(line):
            if not cell:
                map_info.map[i][j] = True
                if _traverse_map(map_info).stuck:
                    added_obstacles_candidates.add((i, j))
                map_info.map[i][j] = False

    # Remove guard location from candidates
    added_obstacles_candidates.discard((map_info.guard[:2]))

    return len(added_obstacles_candidates)


class MapInfo(NamedTuple):
    map: list[list[bool]]
    guard: tuple[int, int, str]


def _parse_map(puzzle_input: str) -> MapInfo:
    puzzle_map = []
    guard = None

    for x, line in enumerate(puzzle_input.split("\n")):
        if not line:
            continue
        line_list = list()
        for y, cell in enumerate(line):
            if cell in "^v<>":
                guard = (x, y, cell)
            line_list.append(cell == "#")
        puzzle_map.append(line_list)

    assert guard is not None, "Guard not found in map"
    assert all(
        len(line) == len(puzzle_map[0]) for line in puzzle_map
    ), "Map is not rectangular"

    return MapInfo(puzzle_map, guard)


class TraversalAnalysis(NamedTuple):
    visited: set[tuple[int, int]]
    stuck: bool


_rotation_map = {
    "^": ">",
    "v": "<",
    "<": "^",
    ">": "v",
}


def _traverse_map(map_info: MapInfo) -> TraversalAnalysis:
    x, y, direction = map_info.guard
    visited = set()
    visited.add((x, y, direction))

    while True:
        # Determine next cell
        match direction:
            case "^":
                next_x, next_y = x - 1, y
            case "v":
                next_x, next_y = x + 1, y
            case "<":
                next_x, next_y = x, y - 1
            case ">":
                next_x, next_y = x, y + 1

        # Guard patrols out of bounds
        if (
            next_x < 0
            or next_x >= len(map_info.map)
            or next_y < 0
            or next_y >= len(map_info.map[0])
        ):
            break

        # Turn right if an obstacle is encountered
        if map_info.map[next_x][next_y]:
            direction = _rotation_map[direction]

        # Otherwise, move to next cell
        else:
            x, y = next_x, next_y
            if (x, y, direction) in visited:
                return TraversalAnalysis(visited, True)
            visited.add((x, y, direction))

    return TraversalAnalysis(visited, False)