from tools.aoc import AOCDay
from typing import Any


def get_winning_times(max_ms: int, dist: int) -> int:
    lower, higher = 0, 0
    for ms in range(max_ms):
        if ms * (max_ms - ms) > dist:
            lower = ms
            break

    for ms in range(max_ms, 0, -1):
        if ms * (max_ms - ms) > dist:
            higher = ms
            break

    return higher - lower + 1


class Day(AOCDay):
    inputs = [
        [
            (288, "input6_test"),
            (219849, "input6"),
        ],
        [
            (71503, "input6_test"),
            (29432455, "input6"),
        ],
    ]

    def part1(self) -> Any:
        race_times, race_dists = self.getIntsFromInput()
        ans = 1

        for i, duration in enumerate(race_times):
            ans *= get_winning_times(duration, race_dists[i])

        return ans

    def part2(self) -> Any:
        race_times, race_dists = self.getIntsFromInput()
        time = int("".join(map(str, race_times)))
        dist = int("".join(map(str, race_dists)))

        return get_winning_times(time, dist)


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