use std::collections::HashMap;

fn time_bingo<RI: IntoIterator<Item = CI>, CI: IntoIterator<Item = u8>>(
    draws: &HashMap<u8, usize>,
    board: RI,
) -> usize {
    let mut col_notes: [usize; 5] = [0; 5];
    let mut soonest_row = usize::MAX;
    for row in board {
        let mut this_row = 0;
        for (x, c) in row.into_iter().enumerate() {
            let t = draws.get(&c).copied().unwrap_or(usize::MAX);
            this_row = this_row.max(t);
            col_notes[x] = col_notes[x].max(t);
        }
        soonest_row = soonest_row.min(this_row);
    }
    std::iter::once(soonest_row)
        .chain(col_notes.iter().copied())
        .min()
        .unwrap()
}

fn score_board<RI: IntoIterator<Item = CI>, CI: IntoIterator<Item = u8>>(
    draws: &HashMap<u8, usize>,
    board: RI,
    step: usize,
) -> usize {
    board
        .into_iter()
        .flat_map(IntoIterator::into_iter)
        .filter(|n| match draws.get(n) {
            None => true,
            Some(s) => s > &step,
        })
        .map(<u8 as Into<usize>>::into)
        .sum()
}

fn parse_draws(line: &str) -> HashMap<u8, usize> {
    line.split(',')
        .flat_map(|x| x.parse().ok())
        .enumerate()
        .map(|(i, x)| (x, i))
        .collect()
}

fn parse_boards<I: Iterator<Item = S>, S: AsRef<str>>(
    input: &mut I,
) -> impl Iterator<Item = Vec<Vec<u8>>> + '_ {
    input
        .filter(|l| !l.as_ref().is_empty())
        .groups_of(5)
        .map(|chunk| {
            chunk
                .iter()
                .map(|line| {
                    line.as_ref()
                        .split_whitespace()
                        .flat_map(|ns| ns.parse().ok())
                        .collect()
                })
                .collect()
        })
}

fn main() {
    use std::io::BufRead;
    let filename = std::env::args().nth(1).expect("Expected filename");
    let file = std::io::BufReader::new(
        std::fs::File::open(<String as AsRef<std::path::Path>>::as_ref(
            &filename,
        ))
        .unwrap(),
    );
    let mut lines = file.lines();
    let draws = lines
        .next()
        .map(|l| parse_draws(&l.unwrap()))
        .expect("File is empty");
    let mut lines = lines.flat_map(Result::ok);
    let boards = parse_boards(&mut lines);
    let (score1, step1, score2, step2) = boards
        .map(|board| {
            let step =
                time_bingo(&draws, board.iter().map(|r| r.iter().copied()));
            (
                score_board(
                    &draws,
                    board.iter().map(|r| r.iter().copied()),
                    step,
                ),
                step,
            )
        })
        .fold(None, |acc, (score, step)| match acc {
            None => Some((score, step, score, step)),
            Some((score1, step1, score2, step2)) => {
                let (score1, step1) = if step < step1 {
                    (score, step)
                } else {
                    (score1, step1)
                };
                let (score2, step2) = if step > step2 {
                    (score, step)
                } else {
                    (score2, step2)
                };
                Some((score1, step1, score2, step2))
            }
        })
        .unwrap();
    println!("- Part 1 -");
    println!("Time: {}", step1);
    println!("Score: {}", score1);
    let (last_called, _) =
        draws.iter().find(|(_, step)| **step == step1).unwrap();
    println!("Last called: {}", last_called);
    println!("Product: {}", score1 * *last_called as usize);
    println!();
    println!("- Part 2 -");
    println!("Time: {}", step2);
    println!("Score: {}", score2);
    let (last_called, _) =
        draws.iter().find(|(_, step)| **step == step2).unwrap();
    println!("Last called: {}", last_called);
    println!("Product: {}", score2 * *last_called as usize);
}

struct Groups<I> {
    inner: I,
    group_size: usize,
}

impl<I: Sized + Iterator> Iterator for Groups<I> {
    type Item = Vec<<I as Iterator>::Item>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.group_size == 0 {
            None
        } else {
            let mut result = Vec::with_capacity(self.group_size);
            for _ in 0..self.group_size {
                match self.inner.next() {
                    Some(a) => {
                        result.push(a);
                    }
                    None => {
                        break;
                    }
                }
            }
            if result.is_empty() {
                self.group_size = 0;
                None
            } else {
                Some(result)
            }
        }
    }
}

trait Groupable: Sized + Iterator {
    fn groups_of(self, size: usize) -> Groups<Self>;
}

impl<I: Sized + Iterator> Groupable for I {
    fn groups_of(self, size: usize) -> Groups<Self> {
        Groups {
            inner: self,
            group_size: size,
        }
    }
}