#!/usr/bin/env python3

# https://adventofcode.com/2023/day/2


def get_lines(filename: str) -> list:
    with open(filename, "r") as file:
        return [line.strip() for line in file.readlines()]


def get_games(lines: list) -> dict:
    games = {}
    for line in lines:
        game_number, grabs = line.split(": ")
        game_number = game_number[5:]  # shave off "Game " from string
        games[int(game_number)] = grabs
    return games


def get_max_needed_cubes(grabs: str) -> dict:
    max_cubes = {}
    for grab in grabs.split(";"):
        for cube_string in grab.split(", "):
            num_cubes_str, colour = cube_string.split()
            num_cubes = int(num_cubes_str)
            if colour not in max_cubes or num_cubes > max_cubes[colour]:
                max_cubes[colour] = num_cubes
    return max_cubes


def game_is_valid(game: dict, cube_pools: dict) -> bool:
    for colour in game.keys():
        if colour not in cube_pools or game[colour] > cube_pools[colour]:
            return False
    return True


def get_sum_of_valid_games(games: dict, cube_pools: dict) -> int:
    valid_games = 0
    for game_num in games.keys():
        grabs = games[game_num]
        max_cubes = get_max_needed_cubes(grabs)
        if game_is_valid(max_cubes, cube_pools):
            valid_games += game_num
    return valid_games


def get_power_of_games(games: dict) -> int:
    power_of_games = 0
    for game_num in games.keys():
        power = 1
        grabs = games[game_num]
        max_cubes = get_max_needed_cubes(grabs)
        for colour in max_cubes.keys():
            power *= max_cubes[colour]
        power_of_games += power
    return power_of_games


def main():
    lines = get_lines("input.txt")
    cube_pools = {"red": 12, "green": 13, "blue": 14}
    games = get_games(lines)
    sum_of_valid_games = get_sum_of_valid_games(games, cube_pools)
    print(f"Part 1: The sum of valid games is: {sum_of_valid_games}")
    power_of_valid_games = get_power_of_games(games)
    print(f"Part 2: The power of valid games is: {power_of_valid_games}")


if __name__ == '__main__':
    main()