from dataclasses import dataclass
from functools import cached_property
from itertools import count
from textwrap import dedent
from typing import Iterator, NamedTuple, override
from unittest import TestCase

from more_itertools import first_true, ilen
from multidict import MultiDict

from puzzles._solver import Solver


@dataclass(frozen=True)
class Fence:
    inside: tuple[int, int]
    outside: tuple[int, int]

    def left(self) -> "Fence":
        (xi, yi) = self.inside
        (xo, yo) = self.outside
        dx, dy = (xo - xi, yo - yi)
        return Fence((xi + dy, yi + dx), (xo + dy, yo + dx))

    def right(self) -> "Fence":
        (xi, yi) = self.inside
        (xo, yo) = self.outside
        dx, dy = (xo - xi, yo - yi)
        return Fence((xi - dy, yi - dx), (xo - dy, yo - dx))


@dataclass
class Region:
    label: str
    plants: set[tuple[int, int]]

    @property
    def area(self) -> int:
        return len(self.plants)

    @cached_property
    def fences(self) -> set[Fence]:
        return {
            Fence((x, y), (x + dx, y + dy))
            for x, y in self.plants
            for dx, dy in ((0, 1), (1, 0), (0, -1), (-1, 0))
            if (x + dx, y + dy) not in self.plants
        }

    @property
    def perimeter(self) -> int:
        return len(self.fences)

    @cached_property
    def sides(self) -> set[tuple[int, int]]:
        fence_ids: dict[Fence, int] = {}
        id_gen = count(0)

        for fence in self.fences:
            if fence in fence_ids:
                continue

            id = next(id_gen)
            fence_ids[fence] = id

            left_fence = fence.left()
            while left_fence in self.fences:
                fence_ids[left_fence] = id
                left_fence = left_fence.left()

            right_fence = fence.right()
            while right_fence in self.fences:
                fence_ids[right_fence] = id
                right_fence = right_fence.right()

        return len(set(fence_ids.values()))


class DayTwelveSolver(Solver):
    regions: MultiDict[Region]

    @override
    def __init__(self, puzzle_input: str):
        self.regions = MultiDict()

        for x, row in enumerate(puzzle_input.strip().splitlines()):
            for y, plant in enumerate(row):
                if plant not in self.regions:
                    self.regions.add(plant, Region(plant, {(x, y)}))
                    continue

                similar_regions = self.regions.popall(plant)
                merge_top_region = first_true(
                    similar_regions,
                    None,
                    pred=lambda region: (x - 1, y) in region.plants,
                )
                merge_left_region = first_true(
                    similar_regions,
                    None,
                    pred=lambda region: (x, y - 1) in region.plants,
                )

                if merge_top_region:
                    similar_regions.remove(merge_top_region)
                    merge_top_region.plants.add((x, y))
                    if merge_left_region and merge_left_region is not merge_top_region:
                        similar_regions.remove(merge_left_region)
                        merge_top_region.plants.update(merge_left_region.plants)
                    self.regions.add(plant, merge_top_region)

                elif merge_left_region:
                    similar_regions.remove(merge_left_region)
                    merge_left_region.plants.add((x, y))
                    self.regions.add(plant, merge_left_region)

                else:
                    self.regions.add(plant, Region(plant, {(x, y)}))

                self.regions.extend(
                    [(plant, similar_region) for similar_region in similar_regions]
                )

    @override
    def solve_p1(self) -> int:
        return sum(region.area * region.perimeter for region in self.regions.values())

    @override
    def solve_p2(self) -> int:
        return sum(region.area * region.sides for region in self.regions.values())


class TestDayTwelveSolver(TestCase):
    def test_solve_p1(self):
        solver = DayTwelveSolver(
            dedent(
                """
                AAAA
                BBCD
                BBCC
                EEEC
                """
            )
        )
        self.assertEqual(solver.solve_p1(), 140)

        solver = DayTwelveSolver(
            dedent(
                """
                RRRRIICCFF
                RRRRIICCCF
                VVRRRCCFFF
                VVRCCCJFFF
                VVVVCJJCFE
                VVIVCCJJEE
                VVIIICJJEE
                MIIIIIJJEE
                MIIISIJEEE
                MMMISSJEEE
                """
            )
        )
        self.assertEqual(solver.solve_p1(), 1930)

    def test_solve_p2(self):
        solver = DayTwelveSolver(
            dedent(
                """
                AAAA
                BBCD
                BBCC
                EEEC
                """
            )
        )
        for region in solver.regions.values():
            print(region.label, region.sides)
        self.assertEqual(solver.solve_p2(), 80)

        solver = DayTwelveSolver(
            dedent(
                """
                RRRRIICCFF
                RRRRIICCCF
                VVRRRCCFFF
                VVRCCCJFFF
                VVVVCJJCFE
                VVIVCCJJEE
                VVIIICJJEE
                MIIIIIJJEE
                MIIISIJEEE
                MMMISSJEEE
                """
            )
        )
        self.assertEqual(solver.solve_p2(), 1206)