use grid::Grid;
use std::collections::HashSet;

fn solve(grid: &Grid<u32>) -> usize {
    let trailheads = grid.iter().enumerate()
        .filter_map(|(index, c)| (*c == 0).then_some(index))
        .collect::<Vec<usize>>();

    let mut sum = 0;

    for trailhead in &trailheads {

        let x = (trailhead % grid.cols()) as i32;
        let y = (trailhead / grid.cols()) as i32;
        let cell = grid.get(y, x).unwrap();

        let mut set = HashSet::new();
        set.extend(check_cell(grid, &cell, x + 1, y));
        set.extend(check_cell(grid, &cell, x, y + 1));
        set.extend(check_cell(grid, &cell, x - 1, y));
        set.extend(check_cell(grid, &cell, x, y - 1));

        sum += set.len();
    }

    sum
}

fn check_cell(grid: &Grid<u32>, prev: &u32, x: i32, y: i32) -> HashSet<(i32, i32)> {
    let mut set = HashSet::new();

    if let Some(cell) = grid.get(y, x) {
        if cell > prev && cell - prev == 1 {
            if *cell == 9 {
                set.insert((x, y));
            } else {
                set.extend(check_cell(grid, &cell, x + 1, y));
                set.extend(check_cell(grid, &cell, x, y + 1));
                set.extend(check_cell(grid, &cell, x - 1, y));
                set.extend(check_cell(grid, &cell, x, y - 1));
            }
        }
    }

    set
}

fn main() {
    let input = include_str!("../input.txt")
        .lines()
        .map(|l| l.chars().map(|c| c.to_digit(10).unwrap()).collect())
        .collect::<Vec<Vec<u32>>>();
    let cols = input.len();
    let grid = Grid::from_vec(input.into_iter().flatten().collect(), cols);

    let time = std::time::Instant::now();

    let answer = solve(&grid);

    println!("Time elapsed: {:?}", time.elapsed());
    println!("Answer: {}", answer);
}