from enum import Enum, auto
from typing import Any, Callable, Iterator, cast

from ..base import BaseSolver


class Cell(Enum):
    AIR = auto()
    ROCK = auto()
    SAND = auto()

    def __str__(self) -> str:
        return {Cell.AIR: ".", Cell.ROCK: "#", Cell.SAND: "O"}[self]


def flow(
    blocks: dict[tuple[int, int], Cell],
    stop_fn: Callable[[int, int], bool],
    fill_fn: Callable[[int, int], Cell],
) -> dict[tuple[int, int], Cell]:
    """
    Flow sands onto the given set of blocks

    Args:
        blocks: Blocks containing ROCK position. Modified in-place.
        stop_fn: Function called with the last (assumed) position of a grain of
            sand BEFORE adding it to blocks. If the function returns True, the grain
            is added and a new one is flowed, otherwise, the whole procedure stops
            and the function returns (without adding the final grain).
        fill_fn: Function called when the target position of a grain (during the
            flowing process) is missing from blocks.

    Returns:
        The input blocks.
    """

    y_max = max(y for _, y in blocks)

    while True:
        x, y = 500, 0

        while y <= y_max:
            moved = False
            for cx, cy in ((x, y + 1), (x - 1, y + 1), (x + 1, y + 1)):
                if (cx, cy) not in blocks and fill_fn(cx, cy) == Cell.AIR:
                    x, y = cx, cy
                    moved = True
                elif blocks[cx, cy] == Cell.AIR:
                    x, y = cx, cy
                    moved = True

                if moved:
                    break

            if not moved:
                break

        if stop_fn(x, y):
            break

        blocks[x, y] = Cell.SAND

    return blocks


# === inputs ===


class Solver(BaseSolver):
    def print_blocks(self, blocks: dict[tuple[int, int], Cell]):
        """
        Print the given set of blocks on a grid.

        Args:
            blocks: Set of blocks to print.
        """
        x_min, y_min, x_max, y_max = (
            min(x for x, _ in blocks),
            0,
            max(x for x, _ in blocks),
            max(y for _, y in blocks),
        )

        for y in range(y_min, y_max + 1):
            self.logger.info(
                "".join(
                    str(blocks.get((x, y), Cell.AIR)) for x in range(x_min, x_max + 1)
                )
            )

    def solve(self, input: str) -> Iterator[Any]:
        lines = [line.strip() for line in input.splitlines()]

        paths: list[list[tuple[int, int]]] = []
        for line in lines:
            parts = line.split(" -> ")
            paths.append(
                [
                    cast(
                        tuple[int, int], tuple(int(c.strip()) for c in part.split(","))
                    )
                    for part in parts
                ]
            )

        blocks: dict[tuple[int, int], Cell] = {}
        for path in paths:
            for start, end in zip(path[:-1], path[1:]):
                x_start = min(start[0], end[0])
                x_end = max(start[0], end[0]) + 1
                y_start = min(start[1], end[1])
                y_end = max(start[1], end[1]) + 1

                for x in range(x_start, x_end):
                    for y in range(y_start, y_end):
                        blocks[x, y] = Cell.ROCK

        self.print_blocks(blocks)

        y_max = max(y for _, y in blocks)

        # === part 1 ===

        blocks_1 = flow(
            blocks.copy(), stop_fn=lambda x, y: y > y_max, fill_fn=lambda x, y: Cell.AIR
        )
        self.print_blocks(blocks_1)
        yield sum(v == Cell.SAND for v in blocks_1.values())

        # === part 2 ===

        blocks_2 = flow(
            blocks.copy(),
            stop_fn=lambda x, y: x == 500 and y == 0,
            fill_fn=lambda x, y: Cell.AIR if y < y_max + 2 else Cell.ROCK,
        )
        blocks_2[500, 0] = Cell.SAND
        self.print_blocks(blocks_2)
        yield sum(v == Cell.SAND for v in blocks_2.values())