from functools import cache
from math import log
from typing import override
from unittest import TestCase

from puzzles._solver import Solver


class DayElevenSolver(Solver):
    stones: tuple[int, ...]

    @override
    def __init__(self, puzzle_input: str):
        self.stones = tuple(map(int, puzzle_input.strip().split(" ")))

    @override
    def solve_p1(self) -> int:
        return sum(sum_descendants(stone, 25) for stone in self.stones)

    @override
    def solve_p2(self) -> int:
        return sum(sum_descendants(stone, 75) for stone in self.stones)


@cache
def sum_descendants(stone: int, blinks: int) -> int:
    if blinks == 0:
        return 1

    if stone == 0:
        return sum_descendants(1, blinks - 1)

    digit_len = int(log(stone, 10)) + 1
    if digit_len % 2 == 0:
        left_digits, right_digits = divmod(stone, 10 ** (digit_len // 2))
        return sum_descendants(left_digits, blinks - 1) + sum_descendants(
            right_digits, blinks - 1
        )

    return sum_descendants(stone * 2024, blinks - 1)


class TestDayElevenSolver(TestCase):
    def test_sum_descendants(self):
        self.assertEqual(sum_descendants(0, 1), 1)
        self.assertEqual(sum_descendants(17, 1), 2)
        self.assertEqual(sum_descendants(125, 4), 3)

    def test_parse(self):
        solver = DayElevenSolver("123 4 56\n")
        self.assertEqual(solver.stones, (123, 4, 56))

    def test_solve_p1(self):
        solver = DayElevenSolver("125 17\n")
        self.assertEqual(solver.solve_p1(), 55312)