from tools.aoc import AOCDay
from typing import Any


def get_min_len(groups: tuple[int, ...]) -> int:
    return sum(groups) + len(groups) - 1


class Day(AOCDay):
    inputs = [
        [
            (21, "input12_test"),
            (101, "input12_debug"),
            (7286, "input12_dennis"),
            (6981, "input12"),
        ],
        [
            (525152, "input12_test"),
            (4546215031609, "input12"),
        ],
    ]

    def parse_input(self, part2: bool = False) -> list[tuple[str, tuple[int, ...]]]:
        records = []
        for line in self.getInput():
            springs, groups = line.split()
            groups = tuple(map(int, groups.split(",")))
            if part2:
                springs = "?".join([springs] * 5)
                groups *= 5
            records.append((springs, groups))

        return records

    def get_arrangements(self, springs: str, groups: tuple[int, ...]) -> int:
        if (springs, groups) in self.DP:
            return self.DP[(springs, groups)]

        if len(groups) == 0:
            return 0

        count = 0
        for i in range(len(springs) - get_min_len(groups) + 1):
            sub_springs = springs[i:]
            hash_len = groups[0]
            if (
                (i > 0 and springs[i - 1] == "#")
                or "." in sub_springs[:hash_len]
                or (len(sub_springs) > hash_len and sub_springs[hash_len] == "#")
            ):
                if "#" in springs[:i]:
                    break
                else:
                    continue

            if len(groups) > 1:
                count += self.get_arrangements(springs[i + groups[0] + 1 :], groups[1:])
            else:
                if "#" in springs[i + groups[0] :]:
                    continue
                count += 1

        self.DP[(springs, groups)] = count
        return self.DP[(springs, groups)]

    def part1(self) -> Any:
        return sum(self.get_arrangements(springs, groups) for springs, groups in self.parse_input())

    def part2(self) -> Any:
        return sum(self.get_arrangements(springs, groups) for springs, groups in self.parse_input(part2=True))


if __name__ == "__main__":
    day = Day(2023, 12)
    day.run(verbose=True)