use std::collections::{BTreeMap, HashSet, VecDeque};

fn main() {
    let input = include_str!("../../input.txt");
    let grid = get_grid(input);
    let count = count_part_numbers(grid);

    println!("{count}")
}

fn count_part_numbers(grid: BTreeMap<(i32, i32), Component>) -> u32 {
    let neighbors = [
        (-1, -1),
        (-1, 0),
        (-1, 1),
        (0, 1),
        (1, 1),
        (1, 0),
        (1, -1),
        (0, -1),
    ];

    let mut visited: HashSet<(i32, i32)> = HashSet::new();
    let mut part_numbers: Vec<u32> = Vec::new();
    for ((r, c), comp) in grid.iter() {
        if let Component::Symbol = comp {
            for (dx, dy) in neighbors {
                let (x, y) = (r + dx, c + dy);
                if visited.contains(&(x, y)) {
                    continue;
                }

                let neighbor = grid.get(&(x, y));
                if let Some(&Component::Value(_)) = neighbor {
                    let num = capture_number((x, y), &grid, &mut visited);
                    part_numbers.push(num)
                }
            }
        }
    }

    part_numbers.into_iter().sum()
}

pub fn capture_number(
    (r, c): (i32, i32),
    grid: &BTreeMap<(i32, i32), Component>,
    visited: &mut HashSet<(i32, i32)>,
) -> u32 {
    let mut number: VecDeque<char> = VecDeque::new();

    // insert current
    if let Some(&Component::Value(val)) = grid.get(&(r, c)) {
        number.push_back(val);
        visited.insert((r, c));
    }
    // collect leftward
    for i in (0..c).rev() {
        if let Some(&Component::Value(val)) = grid.get(&(r, i)) {
            number.push_front(val);
            visited.insert((r, i));
        } else {
            break;
        }
    }
    // collect rightward
    for i in c + 1.. {
        if let Some(&Component::Value(val)) = grid.get(&(r, i)) {
            number.push_back(val);
            visited.insert((r, i));
        } else {
            break;
        }
    }

    number
        .into_iter()
        .collect::<String>()
        .parse::<u32>()
        .unwrap()
}

#[derive(Debug)]
pub enum Component {
    Value(char),
    Symbol,
    None,
}

impl From<char> for Component {
    fn from(value: char) -> Self {
        if value.is_ascii_digit() {
            return Self::Value(value);
        }

        match value {
            '.' => Self::None,
            _ => Self::Symbol,
        }
    }
}

pub fn get_grid(input: &str) -> BTreeMap<(i32, i32), Component> {
    input
        .lines()
        .enumerate()
        .flat_map(|(r, line)| {
            line.chars()
                .enumerate()
                .map(move |(c, char)| ((r as i32, c as i32), char.into()))
        })
        .collect()
}

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

    #[test]
    fn test_get() {
        let schematic = "467..114..
...*......
..35..633.
......#...
617*......
.....+.58.
..592.....
......755.
...$.*....
.664.598..";
        let grid = get_grid(schematic);
        let res = count_part_numbers(grid);
        assert_eq!(res, 4361);
    }
}