use nom::{
    bytes::complete::tag,
    character::complete::{line_ending, space0, space1, u32 as nom_u32},
    combinator::{map, map_res},
    multi::{many1, separated_list1},
    sequence::{preceded, tuple},
    IResult,
};
use std::fs;
use thiserror::Error;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_4.txt")?;
    let mut bingo_game = parse_bingo_game(&input).unwrap().1;
    let total_boards = bingo_game.boards.len();

    let mut winning_game_iter = std::iter::repeat_with(|| bingo_game.do_draw()).flatten();

    let (winning_block, winning_board, winning_mask) = winning_game_iter.next().unwrap();
    dbg!(winning_board.score(&winning_mask) * winning_block);

    let (losing_block, losing_board, losing_mask) =
        winning_game_iter.nth(total_boards - 2).unwrap();
    dbg!(losing_board.score(&losing_mask) * losing_block);

    Ok(())
}

const BINGO_BOARD_WIDTH: usize = 5;

#[derive(Debug)]
struct BingoGame {
    draws: Vec<BingoBlock>,
    boards: Vec<BingoBoard>,
    masks: Vec<BingoBoardMask>,
}

impl BingoGame {
    fn do_draw(&mut self) -> Vec<(BingoBlock, BingoBoard, BingoBoardMask)> {
        let mut wins = Vec::new();
        let mut results = Vec::new();

        if let Some(block) = self.draws.pop() {
            for (index, (mask, board)) in self
                .masks
                .iter_mut()
                .zip(self.boards.iter())
                .enumerate()
                .rev()
            {
                if let Some((row, col)) = board.find_block(&block) {
                    mask.mark(row, col);
                    if mask.has_bingo() {
                        wins.push(index);
                    }
                }
            }
            for win in wins {
                results.push((
                    block.clone(),
                    self.boards.remove(win),
                    self.masks.remove(win),
                ))
            }
        }

        results
    }
}

#[derive(Debug, Error)]
enum BingoBoardParseError {
    #[error("input board was the wrong width")]
    WrongWidth,
    #[error("input board was the wrong height")]
    WrongHeight,
}

#[derive(Debug, Default, Clone, PartialEq, Eq)]
struct BingoBoardMask {
    data: [[bool; BINGO_BOARD_WIDTH]; BINGO_BOARD_WIDTH],
}

impl BingoBoardMask {
    fn has_bingo(&self) -> bool {
        for i in 0..BINGO_BOARD_WIDTH {
            let row_bingo = self.data[i].iter().all(|marked| *marked);
            let col_bingo = self.data.iter().all(|row| row[i]);
            if row_bingo || col_bingo {
                return true;
            }
        }
        false
    }

    fn mark(&mut self, row: usize, col: usize) {
        self.data[row][col] = true;
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct BingoBoard {
    data: [[BingoBlock; BINGO_BOARD_WIDTH]; BINGO_BOARD_WIDTH],
}

impl BingoBoard {
    fn find_block(&self, block: &BingoBlock) -> Option<(usize, usize)> {
        for (row, row_data) in self.data.iter().enumerate() {
            for (col, col_data) in row_data.iter().enumerate() {
                if col_data == block {
                    return Some((row, col));
                }
            }
        }
        None
    }

    fn score(&self, mask: &BingoBoardMask) -> BingoGameScore {
        let mut score = BingoGameScore::default();
        for row in 0..BINGO_BOARD_WIDTH {
            for col in 0..BINGO_BOARD_WIDTH {
                if !mask.data[row][col] {
                    score = score + self.data[row][col];
                }
            }
        }
        score
    }
}

impl BingoBoard {
    fn new(data: Vec<Vec<BingoBlock>>) -> Result<BingoBoard, BingoBoardParseError> {
        let vec_array_data: Vec<[BingoBlock; BINGO_BOARD_WIDTH]> = data
            .into_iter()
            .map(|row| row.try_into())
            .collect::<Result<Vec<_>, _>>()
            .map_err(|_| BingoBoardParseError::WrongWidth)?;

        let array_array_data: [[BingoBlock; BINGO_BOARD_WIDTH]; BINGO_BOARD_WIDTH] = vec_array_data
            .try_into()
            .map_err(|_| BingoBoardParseError::WrongHeight)?;
        Ok(BingoBoard {
            data: array_array_data,
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct BingoBlock(u32);

#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
struct BingoGameScore(u32);

impl std::ops::Add<BingoBlock> for BingoGameScore {
    type Output = BingoGameScore;
    fn add(self, rhs: BingoBlock) -> BingoGameScore {
        BingoGameScore(self.0 + rhs.0)
    }
}

impl std::ops::Mul<BingoBlock> for BingoGameScore {
    type Output = BingoGameScore;
    fn mul(self, rhs: BingoBlock) -> BingoGameScore {
        BingoGameScore(self.0 * rhs.0)
    }
}

fn parse_bingo_game(input: &str) -> IResult<&str, BingoGame> {
    map(
        tuple((
            parse_draws,
            many1(line_ending),
            separated_list1(many1(line_ending), parse_board),
        )),
        |(draws, _, boards)| BingoGame {
            draws: draws.into_iter().rev().collect(),
            masks: vec![BingoBoardMask::default(); boards.len()],
            boards,
        },
    )(input)
}

fn parse_draws(input: &str) -> IResult<&str, Vec<BingoBlock>> {
    separated_list1(tag(","), parse_block)(input)
}

fn parse_block(input: &str) -> IResult<&str, BingoBlock> {
    map(nom_u32, BingoBlock)(input)
}

fn parse_board(input: &str) -> IResult<&str, BingoBoard> {
    map_res(
        separated_list1(
            line_ending,
            preceded(space0, separated_list1(space1, parse_block)),
        ),
        BingoBoard::new,
    )(input)
}

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

    #[test]
    fn parses_a_board() {
        assert_eq!(
            parse_board(
                r##"86 31 71 11 56
99 12 17 10 46
 5 33 85 61  2
30  1 28 88 66
15 38 21 54 64"##
            ),
            Ok((
                "",
                BingoBoard {
                    data: [
                        [
                            BingoBlock(86),
                            BingoBlock(31),
                            BingoBlock(71),
                            BingoBlock(11),
                            BingoBlock(56)
                        ],
                        [
                            BingoBlock(99),
                            BingoBlock(12),
                            BingoBlock(17),
                            BingoBlock(10),
                            BingoBlock(46)
                        ],
                        [
                            BingoBlock(5),
                            BingoBlock(33),
                            BingoBlock(85),
                            BingoBlock(61),
                            BingoBlock(2)
                        ],
                        [
                            BingoBlock(30),
                            BingoBlock(1),
                            BingoBlock(28),
                            BingoBlock(88),
                            BingoBlock(66)
                        ],
                        [
                            BingoBlock(15),
                            BingoBlock(38),
                            BingoBlock(21),
                            BingoBlock(54),
                            BingoBlock(64)
                        ]
                    ]
                }
            ))
        );
    }
}