use aoc_runner_derive::{aoc, aoc_generator};

#[derive(Debug, Clone)]
pub struct Data {
    pub numbers: Vec<u8>,
    pub boards: Vec<[[(u8, bool); 5]; 5]>,
}

#[aoc_generator(day04)]
pub fn input_generator(input: &str) -> Data {
    let mut lines = input.lines();
    let mut boards = vec![];
    let numbers = lines
        .next()
        .unwrap()
        .split(',')
        .map(|s| s.parse())
        .collect::<Result<Vec<u8>, _>>()
        .unwrap();
    let mut row = 0;
    let mut board = [[(0u8, false); 5]; 5];
    for line in lines {
        if line.is_empty() {
            continue;
        }
        let line = line
            .split_ascii_whitespace()
            .map(|s| s.parse())
            .collect::<Result<Vec<u8>, _>>()
            .unwrap();
        line.iter().enumerate().for_each(|(i, v)| {
            board[row][i] = (*v, false);
        });
        row = (row + 1) % 5;
        if row == 0 {
            boards.push(board);
            board = [[(0u8, false); 5]; 5];
        }
    }
    Data { numbers, boards }
}

#[aoc(day04, part1)]
pub fn solve_part1(input: &Data) -> usize {
    let Data {
        numbers,
        mut boards,
    } = input.clone();
    for number in numbers.iter() {
        for board in &mut boards {
            for row in board.iter_mut() {
                for (n, b) in row {
                    if n == number {
                        *b = true;
                    }
                }
            }
            let win = (0..5)
                .map(|v| {
                    [
                        board[v][0],
                        board[v][1],
                        board[v][2],
                        board[v][3],
                        board[v][4],
                    ]
                    .iter()
                    .all(|v| v.1)
                        | [
                            board[0][v],
                            board[1][v],
                            board[2][v],
                            board[3][v],
                            board[4][v],
                        ]
                        .iter()
                        .all(|v| v.1)
                })
                .any(|v| v);
            if win {
                return board
                    .iter()
                    .flatten()
                    .filter_map(|(v, b)| (!*b).then(|| *v as usize))
                    .sum::<usize>()
                    * (*number as usize);
            }
        }
    }
    unreachable!()
}

#[aoc(day04, part2)]
pub fn solve_part2(input: &Data) -> usize {
    let mut ret = None;
    let Data {
        numbers,
        mut boards,
    } = input.clone();
    let mut board_wins: Vec<bool> = vec![false; boards.len()];
    for number in numbers.iter() {
        for (board_n, board) in boards.iter_mut().enumerate() {
            for row in board.iter_mut() {
                for (n, b) in row {
                    if n == number {
                        *b = true;
                    }
                }
            }
            let win = (0..5)
                .map(|v| {
                    [
                        board[v][0],
                        board[v][1],
                        board[v][2],
                        board[v][3],
                        board[v][4],
                    ]
                    .iter()
                    .all(|v| v.1)
                        | [
                            board[0][v],
                            board[1][v],
                            board[2][v],
                            board[3][v],
                            board[4][v],
                        ]
                        .iter()
                        .all(|v| v.1)
                })
                .any(|v| v);
            if win && !board_wins[board_n] {
                ret = Some(
                    board
                        .iter()
                        .flatten()
                        .filter(|(_, b)| !*b)
                        .map(|(v, _)| *v as usize)
                        .sum::<usize>()
                        * (*number as usize),
                );
                board_wins[board_n] = true;
            }
        }
    }
    return ret.unwrap();
}