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

from puzzles._solver import Solver


class TrailSet(NamedTuple):
    trailhead: tuple[int, int]
    trailends: set[tuple[int, int]]
    count: int


@dataclass
class TopographicalMap:
    grid: list[list[str]]
    trailheads: set[tuple[int, int]]

    @classmethod
    def parse(cls, raw_map: str) -> "TopographicalMap":
        grid = []
        trailheads = set()

        for x, row in enumerate(raw_map.strip().splitlines()):
            grid_row = []

            for y, cell in enumerate(row):
                height = int(cell)
                grid_row.append(height)
                if height == 0:
                    trailheads.add((x, y))

            grid.append(grid_row)

        return cls(grid, trailheads)

    def trails(self) -> list[TrailSet]:
        trails = []

        for trailhead in self.trailheads:
            trailends: set[tuple, tuple] = set()
            count = 0
            x, y = trailhead
            height = self.grid[x][y]

            def follow_trail(x, y, height):
                nonlocal count

                if height == 9:
                    trailends.add((x, y))
                    count += 1
                    return

                if x > 0 and self.grid[x - 1][y] == height + 1:
                    follow_trail(x - 1, y, height + 1)
                if x < len(self.grid) - 1 and self.grid[x + 1][y] == height + 1:
                    follow_trail(x + 1, y, height + 1)
                if y > 0 and self.grid[x][y - 1] == height + 1:
                    follow_trail(x, y - 1, height + 1)
                if y < len(self.grid[0]) - 1 and self.grid[x][y + 1] == height + 1:
                    follow_trail(x, y + 1, height + 1)

            follow_trail(x, y, height)

            trails.append(TrailSet(trailhead, trailends, count))

        return trails


class DayTenSolver(Solver):
    topographical_map: TopographicalMap

    @override
    def __init__(self, puzzle_input: str):
        self.topographical_map = TopographicalMap.parse(puzzle_input)

    @override
    def solve_p1(self) -> int:
        trail_sets = self.topographical_map.trails()
        return sum(len(trail_set.trailends) for trail_set in trail_sets)

    @override
    def solve_p2(self) -> int:
        trail_sets = self.topographical_map.trails()
        return sum(trail_set.count for trail_set in trail_sets)


class TestDayTenSolver(TestCase):
    test_input = dedent(
        """
        89010123
        78121874
        87430965
        96549874
        45678903
        32019012
        01329801
        10456732
        """
    )

    def test_parse(self):
        solver = DayTenSolver(self.test_input)
        self.assertEqual(
            solver.topographical_map,
            TopographicalMap(
                grid=[
                    [8, 9, 0, 1, 0, 1, 2, 3],
                    [7, 8, 1, 2, 1, 8, 7, 4],
                    [8, 7, 4, 3, 0, 9, 6, 5],
                    [9, 6, 5, 4, 9, 8, 7, 4],
                    [4, 5, 6, 7, 8, 9, 0, 3],
                    [3, 2, 0, 1, 9, 0, 1, 2],
                    [0, 1, 3, 2, 9, 8, 0, 1],
                    [1, 0, 4, 5, 6, 7, 3, 2],
                ],
                trailheads={
                    (0, 2),
                    (0, 4),
                    (2, 4),
                    (4, 6),
                    (5, 2),
                    (5, 5),
                    (6, 0),
                    (6, 6),
                    (7, 1),
                },
            ),
        )

    def test_solve_p1(self):
        solver = DayTenSolver(self.test_input)
        self.assertEqual(solver.solve_p1(), 36)

    def test_solve_p2(self):
        solver = DayTenSolver(self.test_input)
        self.assertEqual(solver.solve_p2(), 81)