from dataclasses import dataclass
from textwrap import dedent
from typing import Literal, NamedTuple, get_args, override
from unittest import TestCase

from puzzles._solver import Solver

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


class GuardPosition(NamedTuple):
    x: int
    y: int
    direction: tuple[int, int]


def move_position(position: GuardPosition) -> GuardPosition:
    return GuardPosition(
        position.x + position.direction[0],
        position.y + position.direction[1],
        position.direction,
    )


def rotate_position(position: GuardPosition) -> GuardPosition:
    return GuardPosition(
        position.x,
        position.y,
        (position.direction[1], -position.direction[0]),
    )


class TestGuardPosition(TestCase):
    def test(self):
        guard = GuardPosition(0, 0, (0, 1))
        self.assertEqual(move_position(guard), GuardPosition(0, 1, (0, 1)))
        self.assertEqual(rotate_position(guard), GuardPosition(0, 0, (1, 0)))


class TraversalAnalysis(NamedTuple):
    visited: set[GuardPosition]
    stuck: bool


def unique_locations(traversal_analysis: TraversalAnalysis) -> set[tuple[int, int]]:
    return set((p.x, p.y) for p in traversal_analysis.visited)


class TestTraversalAnalysis(TestCase):
    def test(self):
        analysis = TraversalAnalysis(
            {GuardPosition(0, 0, (0, 1)), GuardPosition(2, 5, (1, 0))},
            False,
        )
        self.assertEqual(
            unique_locations(analysis),
            {(0, 0), (2, 5)},
        )


class DaySixSolver(Solver):
    puzzle_map: list[list[bool]]
    guard_origin: GuardPosition

    @override
    def __init__(self, puzzle_input: str):
        self.puzzle_map = []
        guard_origin = None

        for x, line in enumerate(puzzle_input.strip().split("\n")):
            line_list = list()
            for y, cell in enumerate(line):
                if cell in direction_map:
                    guard_origin = GuardPosition(x, y, direction_map[cell])
                line_list.append(cell == "#")
            self.puzzle_map.append(line_list)

        if guard_origin is None:
            raise ValueError("Guard not found in map")

        self.guard_origin = guard_origin

        if not all(len(line) == len(self.puzzle_map[0]) for line in self.puzzle_map):
            raise ValueError("Map is not rectangular")

    @override
    def solve_p1(self) -> int:
        traversal_analysis = self.traverse_map()
        return len(unique_locations(traversal_analysis))

    @override
    def solve_p2(self) -> int:
        candidate_locations = unique_locations(self.traverse_map())
        candidate_locations.discard((self.guard_origin.x, self.guard_origin.y))
        verified_locations = set()

        for x, y in candidate_locations:
            if not self.puzzle_map[x][y]:
                self.puzzle_map[x][y] = True
                if self.traverse_map().stuck:
                    verified_locations.add((x, y))
                self.puzzle_map[x][y] = False

        return len(verified_locations)

    def out_of_bounds(self, guard_position: GuardPosition) -> bool:
        return (
            guard_position.x < 0
            or guard_position.x >= len(self.puzzle_map)
            or guard_position.y < 0
            or guard_position.y >= len(self.puzzle_map[0])
        )

    def traverse_map(self) -> TraversalAnalysis:
        position = self.guard_origin
        visited = set({position})

        while True:
            next_position = move_position(position)
            if self.out_of_bounds(next_position):
                return TraversalAnalysis(visited, False)

            # Obstacle is encountered
            if self.puzzle_map[next_position.x][next_position.y]:
                next_position = rotate_position(position)

            # Guard is stuck
            if next_position in visited:
                return TraversalAnalysis(visited, True)

            position = next_position
            visited.add(position)


class TestDaySixSolver(TestCase):
    def test(self):
        solver = DaySixSolver(
            dedent(
                """
                ....#.....
                .........#
                ..........
                ..#.......
                .......#..
                ..........
                .#..^.....
                ........#.
                #.........
                ......#...
                """
            )
        )
        self.assertEqual(solver.guard_origin, GuardPosition(6, 4, (-1, 0)))
        self.assertEqual(solver.solve_p1(), 41)
        self.assertEqual(solver.solve_p2(), 6)