# -*- coding: utf-8 -*-
from abc import ABC
from functools import lru_cache
from typing import Tuple, Iterator, Set, Any, Optional

from aoc import BaseAssignment
from aoc.datastructures import Coordinate


class Assignment(BaseAssignment, ABC):
    def __init__(self, path):
        super().__init__(path)

        self.cave_depth = None

    def coordinate_ranges(
        self, input: Iterator[str]
    ) -> Iterator[Tuple[Coordinate, Coordinate]]:
        for item in input:
            last = None

            for i in item.split(" -> "):
                x, y = i.split(",")

                coordinate = int(x), int(y)

                if last is not None:
                    yield last, coordinate

                last = coordinate

    def coordinates(self, a: Coordinate, b: Coordinate) -> Iterator[Coordinate]:
        for x in range(min(a[0], b[0]), max(a[0], b[0]) + 1):
            for y in range(min(a[1], b[1]), max(a[1], b[1]) + 1):
                yield x, y

    def generate_rocks(self, input: Iterator[str]) -> Set[Coordinate]:
        rocks = set(
            coordinate
            for coordinates in self.coordinate_ranges(input)
            for coordinate in self.coordinates(*coordinates)
        )

        self.cave_depth = self.get_cave_depth(rocks)

        return rocks

    def get_cave_depth(self, cave: Set[Coordinate]):
        return max([item[1] for item in cave])

    def invalid_position(self, sand: Coordinate, cave: set[Coordinate]):
        raise NotImplementedError()

    def blocking(self, sand: Coordinate, cave: set[Coordinate]):
        raise NotImplementedError()

    def drop_sand(
        self, sand: Coordinate, cave: Set[Coordinate]
    ) -> Optional[Coordinate]:
        if self.invalid_position(sand, cave):
            return None

        next_positions = [
            (sand[0], sand[1] + 1),
            (sand[0] - 1, sand[1] + 1),
            (sand[0] + 1, sand[1] + 1),
        ]

        non_blocking = [
            next_position
            for next_position in next_positions
            if not self.blocking(next_position, cave)
        ]

        if len(non_blocking) == 0:
            return sand

        return self.drop_sand(non_blocking[0], cave)

    def generate_sand(self, cave: set[Coordinate]) -> Iterator[Coordinate]:
        cave_copy = set(cave)
        sand = (500, 0)

        while True:
            coordinate = self.drop_sand(sand, cave_copy)

            if coordinate is None:
                return

            cave_copy.add(coordinate)
            yield coordinate

    def run(self, input: Iterator) -> Any:
        cave = self.generate_rocks(input)
        coordinates = set()

        for sand in self.generate_sand(cave):
            coordinates.add(sand)

        return len(coordinates)


class AssignmentOne(Assignment):
    example_result = 24

    def invalid_position(self, sand: Coordinate, cave: set[Coordinate]):
        return sand[1] > self.cave_depth

    def blocking(self, sand: Coordinate, cave: set[Coordinate]):
        return sand in cave


class AssignmentTwo(Assignment):
    example_result = 93

    def invalid_position(self, sand: Coordinate, cave: set[Coordinate]):
        return False

    def blocking(self, sand: Coordinate, cave: set[Coordinate]):
        return sand in cave or sand[1] >= self.cave_depth + 2

    def run(self, input: Iterator) -> Any:
        cave = self.generate_rocks(input)
        coordinates = set()

        for sand in self.generate_sand(cave):
            coordinates.add(sand)

            if sand == (500, 0):
                break

        return len(coordinates)