use itertools::Itertools;
use std::cmp::max;

#[aoc_generator(day2)]
fn parse_input(input: &str) -> Vec<Vec<Vec<(i32, Color)>>> {
    input.split('\n')
        .map(|line| {
            let rounds = &line[line.find(':').unwrap() + 1..];
            rounds.split(';')
                .map(|game| {
                    game.split(',')
                        .map(|round| {
                            let (number, color) = round.split_whitespace().collect_tuple().unwrap();
                            let number = number.parse::<i32>().unwrap();
                            (number, color.into_color())
                        }).collect::<Vec<_>>()
                }).collect::<Vec<_>>()
        }).collect::<Vec<_>>()
}

#[aoc(day2, part1)]
fn solve_part1(input: &Vec<Vec<Vec<(i32, Color)>>>) -> usize {
    let (r_max, g_max, b_max) = (12, 13, 14);
    let mut sum: usize = 0;

    for (game_id, rounds) in input.iter().enumerate() {
        let mut sad = false;

        for round in rounds {
            let (mut r_cur, mut g_cur, mut b_cur) = (0, 0, 0);
            for (num, color) in round {
                match color {
                    Color::Red => r_cur += num,
                    Color::Blue => b_cur += num,
                    Color::Green => g_cur += num
                }
            }

            if r_cur > r_max || b_cur > b_max || g_cur > g_max {
                sad = true;
            }
        }

        if !sad {
            sum += game_id + 1;
        }
    }

    return sum;
}

#[aoc(day2, part2)]
fn solve_part2(input: &Vec<Vec<Vec<(i32, Color)>>>) -> usize {
    let mut sum: usize = 0;

    for rounds in input {
        let (mut r_max, mut g_max, mut b_max) = (0, 0, 0);
        for round in rounds {
            for (num, color) in round {
                match color {
                    Color::Red => r_max = max(*num, r_max),
                    Color::Blue => b_max = max(*num, b_max),
                    Color::Green => g_max = max(*num, g_max),
                }
            }
        }
        sum += (r_max * g_max * b_max) as usize;
    }

    sum
}

#[derive(Debug)]
enum Color {
    Red,
    Blue,
    Green
}

trait ColorConvertable {
    fn into_color(&self) -> Color;
}

impl ColorConvertable for str {
    fn into_color(&self) -> Color {
        match self {
            "red" => Color::Red,
            "blue" => Color::Blue,
            "green" => Color::Green,
            _ => panic!("Invalid Color")
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    static TEST_INPUT: &str = r"Game 1: 3 blue, 4 red; 1 red, 2 green, 6 blue; 2 green
Game 2: 1 blue, 2 green; 3 green, 4 blue, 1 red; 1 green, 1 blue
Game 3: 8 green, 6 blue, 20 red; 5 blue, 4 red, 13 green; 5 green, 1 red
Game 4: 1 green, 3 red, 6 blue; 3 green, 6 red; 3 green, 15 blue, 14 red
Game 5: 6 red, 1 blue, 3 green; 2 blue, 1 red, 2 green";

    #[test]
    fn part1_example() {
        assert_eq!(solve_part1(&parse_input(TEST_INPUT)), 8);
    }

    #[test]
    fn part2_example() {
        assert_eq!(solve_part2(&parse_input(TEST_INPUT)), 2286);
    }
}