use std::collections::VecDeque;

#[aoc_generator(day4)]
fn parse(input: &str) -> Vec<Vec<Vec<i32>>> {
    let input: Vec<Vec<Vec<i32>>> = input.lines()
        .map(|card| &card[(card.find(':').unwrap() + 1)..])
        .map(|card| card.trim())
        .map(|card| {
            card.split('|')
                .map(|numbers| {
                    numbers.split_whitespace()
                    .map(|num| num.parse::<i32>().unwrap())
                    .collect::<Vec<_>>()
                }).collect::<Vec<_>>()
        }).collect::<Vec<_>>();
    input
}

#[aoc(day4, part1)]
fn part1(input: &Vec<Vec<Vec<i32>>>) -> i32 {
    let mut total_pts = 0;

    for card in input {
        let mut card_pts = 0;
        let winning = card.first().unwrap();
        let ours = card.last().unwrap();

        for num in ours {
            if winning.contains(num) {
                if card_pts == 0 {
                    card_pts = 1;
                } else {
                    card_pts *= 2;
                }
            }
        }

        total_pts += card_pts;
    }

    total_pts
}

#[aoc(day4, part2)]
fn part2(input: &Vec<Vec<Vec<i32>>>) -> i32 {
    let mut queue = VecDeque::from((0..input.len()).collect::<Vec<usize>>());

    let mut total_cards = 0;
    while !queue.is_empty() {
        let card_num = queue.pop_front().unwrap();
        let card = input.get(card_num).unwrap();
        total_cards += 1;
        let mut dup_cards = 0;
        let winning_nums = card.first().unwrap();
        let our_nums = card.last().unwrap();

        for num in our_nums {
            if winning_nums.contains(num) {
                dup_cards += 1;
            }
        }
        for card in (card_num + 1)..=(card_num + dup_cards) {
            if card < input.len() {
                queue.push_back(card);
            }
        }
    }

    total_cards
}


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

    const EX: &str = r"Card 1: 41 48 83 86 17 | 83 86  6 31 17  9 48 53
Card 2: 13 32 20 16 61 | 61 30 68 82 17 32 24 19
Card 3:  1 21 53 59 44 | 69 82 63 72 16 21 14  1
Card 4: 41 92 73 84 69 | 59 84 76 51 58  5 54 83
Card 5: 87 83 26 28 32 | 88 30 70 12 93 22 82 36
Card 6: 31 18 13 56 72 | 74 77 10 23 35 67 36 11";

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse(EX)), 13);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse(EX)), 30);
    }
}