use anyhow::anyhow;
use std::{collections::HashMap, str::FromStr};

use crate::util;

#[derive(Debug, Clone)]
pub struct Board {
    pub numbers: HashMap<(usize, usize), u64>,
    pub marks: HashMap<(usize, usize), bool>,
}

impl FromStr for Board {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> anyhow::Result<Self> {
        let mut numbers = HashMap::new();
        let mut marks = HashMap::new();

        for (y, line) in s.trim().lines().enumerate() {
            for (x, n) in line.split_whitespace().enumerate() {
                numbers.insert((x, y), n.parse()?);
                marks.insert((x, y), false);
            }
        }

        Ok(Self { numbers, marks })
    }
}

impl Board {
    fn mark(&mut self, number: u64) {
        for (k, v) in &self.numbers {
            if v == &number {
                // SAFETY: coords are identical in both hashmaps
                *self.marks.get_mut(k).unwrap() = true
            }
        }
    }

    fn is_winner(&self) -> bool {
        // check rows
        for i in 0..5 {
            let winner = self
                .marks
                .iter()
                // grab the x axis
                .filter(|((x, _), _)| x == &i)
                // if all booleans are true, this entire row has been marked
                .all(|((_, _), b)| *b);

            if winner {
                return true;
            }
        }

        // check columns
        for j in 0..5 {
            let winner = self
                .marks
                .iter()
                // grab the y axis
                .filter(|((_, y), _)| y == &j)
                // if all booleans are true, this entire column has been marked
                .all(|((_, _), b)| *b);

            if winner {
                return true;
            }
        }

        false
    }

    fn get_unmarked(&self) -> Vec<u64> {
        self.marks
            .iter()
            .filter_map(|((x, y), b)| {
                match b {
                    // SAFETY: coords are identical in both hashmaps
                    false => Some(*self.numbers.get(&(*x, *y)).unwrap()),
                    true => None,
                }
            })
            .collect()
    }
}

pub fn get_input() -> anyhow::Result<(Vec<u64>, Vec<Board>)> {
    let input = util::get_input(2021, 4)?;
    let mut lines = input.lines();

    // get the numbers
    let numbers = lines
        .next()
        .ok_or_else(|| anyhow!("couldn't get numbers"))?
        .split(',')
        .map(|n| Ok(n.parse()?))
        .collect::<anyhow::Result<Vec<_>>>()?;

    // get the boards
    let mut boards = Vec::new();
    let mut board_iter = lines.peekable();

    while board_iter.peek().is_some() {
        let board = board_iter
            .by_ref()
            .skip(1)
            .take(5)
            // reintroduce the \n so that we can parse it in it's original 5x5 format
            .map(|l| l.to_owned() + "\n")
            .collect::<String>()
            .parse()?;

        boards.push(board);
    }

    Ok((numbers, boards))
}

pub fn part_one(numbers: &[u64], boards: &[Board]) -> Option<u64> {
    // clone the boards
    let mut boards = boards.to_vec();

    for number in numbers {
        // apply number to each board, checking each board has/hasn't one
        for board in boards.iter_mut() {
            board.mark(*number);
            if board.is_winner() {
                // calculate answer
                return Some(board.get_unmarked().iter().sum::<u64>() * number);
            }
        }
    }

    None
}

pub fn part_two(numbers: &[u64], boards: &[Board]) -> Option<u64> {
    // clone the boards
    let mut boards = boards.to_vec();
    let mut boards_remaining = boards.len();

    for number in numbers {
        for board in boards.iter_mut().peekable() {
            board.mark(*number);

            if board.is_winner() && boards_remaining == 1 {
                // SAFETY: boards_remaining moves in lockstep with boards, we know there's one board left
                return Some(boards.first().unwrap().get_unmarked().iter().sum::<u64>() * number);
            }
        }

        // now filter out the winners
        boards = boards
            .iter()
            .filter_map(|board| match board.is_winner() {
                false => Some(board.clone()),
                true => None,
            })
            .collect();

        boards_remaining = boards.len();
    }

    None
}

pub fn main() -> anyhow::Result<()> {
    let (numbers, boards) = get_input()?;

    log::info!("d04p01: {}", part_one(&numbers, &boards).unwrap());
    log::info!("d04p02: {}", part_two(&numbers, &boards).unwrap());

    Ok(())
}

#[cfg(test)]
mod tests {
    #[test]
    fn part_one() {
        let (numbers, boards) = super::get_input().unwrap();
        let result = super::part_one(&numbers, &boards).unwrap();
        assert_eq!(result, 16_674);
    }

    #[test]
    fn part_two() {
        let (numbers, boards) = super::get_input().unwrap();
        let result = super::part_two(&numbers, &boards).unwrap();
        assert_eq!(result, 7_075);
    }
}