use aoc_runner_derive::{aoc, aoc_generator};
use grid::Grid;
use itertools::Itertools;

#[aoc_generator(day4)]
fn parse(input: &str) -> Grid<u8> {
    input.parse().unwrap()
}

#[aoc(day4, part1, Brute)]
fn part1(input: &Grid<u8>) -> u64 {
    (0..input.height() * input.width())
        .filter(|i| *input.get(&input.coord(*i as i64).unwrap()).unwrap() == b'@')
        .map(|i| input.adjacent_count(&input.coord(i as i64).unwrap(), |c| *c == b'@'))
        .filter(|n| *n < 4)
        .count() as u64
}

#[aoc(day4, part2, Brute)]
fn part2(input: &Grid<u8>) -> u64 {
    let mut grid = input.clone();
    let mut removed = 0;

    loop {
        let mut removed_iteration = 0;
        for i in 0..grid.width() * grid.height() {
            let pos = grid.coord(i as i64).unwrap();
            if grid.get(&pos).is_some_and(|c| *c == b'@')
                && grid.adjacent_count(&pos, |c| *c == b'@') < 4
            {
                // remove the roll
                grid.set(&pos, b'.');
                removed_iteration += 1;
            }
        }
        if removed_iteration == 0 {
            return removed;
        }
        removed += removed_iteration
    }
}

#[aoc(day4, part2, RollList)]
fn part2_list(input: &Grid<u8>) -> u64 {
    let mut grid = input.clone();
    let mut to_check = grid.find_all(&b'@').collect_vec();
    let mut removed = 0;

    while let Some(roll) = to_check.pop() {
        if grid.get(&roll).is_none_or(|c| *c == b'.') {
            continue;
        }
        let adjacent_rolls = grid
            .adjacent_iter(&roll)
            .filter_map(|i| if *i.value == b'@' { Some(i.pos) } else { None })
            .collect_vec();
        if adjacent_rolls.len() < 4 {
            grid.set(&roll, b'.');
            removed += 1;
            to_check.extend_from_slice(&adjacent_rolls);
        }
    }

    removed
}

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

    const EXAMPLE: &str = "..@@.@@@@.
@@@.@.@.@@
@@@@@.@.@@
@.@@@@..@.
@@.@@@@.@@
.@@@@@@@.@
.@.@.@.@@@
@.@@@.@@@@
.@@@@@@@@.
@.@.@@@.@.";

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

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

    #[test]
    fn part2_list_example() {
        assert_eq!(part2_list(&parse(EXAMPLE)), 43);
    }
}