use aoc_runner_derive::{aoc, aoc_generator};
use grid::{Coord2d, Grid};
use std::str::FromStr;

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

fn emit_beams(grid: &mut Grid<u8>, pos: Coord2d) -> u64 {
    match grid.set(&pos, b'|') {
        None | Some(b'|') => 0,
        Some(b'.') | Some(b'S') => emit_beams(grid, pos + &(0, 1)),
        Some(b'^') => 1 + emit_beams(grid, pos + &(-1, 0)) + emit_beams(grid, pos + &(1, 0)),
        _ => panic!(),
    }
}

#[aoc(day7, part1)]
fn part1(input: &Grid<u8>) -> u64 {
    let mut grid = input.clone();
    emit_beams(&mut grid, input.find(&b'S').unwrap())
}

#[aoc(day7, part2)]
fn part2(input: &Grid<u8>) -> u64 {
    let mut col_counts = vec![0; input.width()];
    col_counts[input.find(&b'S').unwrap().x as usize] = 1;

    // Start row is already set up
    for y in 1..input.height() {
        // for each column with non-zero counts
        for x in 0..input.width() {
            if col_counts[x] == 0 {
                continue;
            }
            match input.get(&(x, y)) {
                // if our current position is a caret, add that count to our neighbours
                Some(b'^') => {
                    for col in [x as i64 - 1, x as i64 + 1]
                        .into_iter()
                        .filter(|x| *x >= 0 && *x < input.width() as i64)
                    {
                        col_counts[col as usize] += col_counts[x];
                    }
                    col_counts[x] = 0;
                }
                // if our current position is a . do nothing, it's already counted
                Some(b'.') => continue,
                Some(c) => panic!("unexpected character {c}"),
                None => panic!("How did we end up outside the grid?"),
            }
        }
    }

    col_counts.iter().sum::<u64>()
}

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

    const TRIVIAL_EXAMPLE: &str = "...S...
.......
...^...
..^.^.^
...^.^.
.......
.......";

    const EXAMPLE: &str = ".......S.......
...............
.......^.......
...............
......^.^......
...............
.....^.^.^.....
...............
....^.^...^....
...............
...^.^...^.^...
...............
..^...^.....^..
...............
.^.^.^.^.^...^.
...............";

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

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

    #[test]
    fn part2_trivial_example() {
        assert_eq!(part2(&parse(TRIVIAL_EXAMPLE)), 7);
    }
}