use std::collections::{BTreeSet, VecDeque};

use aoc_runner_derive::{aoc, aoc_generator};

pub struct Data(Vec<Vec<u8>>);

impl Data {
    fn shape(&self) -> (usize, usize) {
        if !self.0.iter().all(|v| v.len() == self.0[0].len()) {
            panic!("Not rectangular!");
        }
        (self.0.len(), self.0[0].len())
    }
    fn get(&self, x: i64, y: i64) -> Option<u8> {
        if x < 0 || y < 0 {
            return None;
        }
        self.0
            .get(x as usize)
            .map(|v| v.get(y as usize).copied())
            .flatten()
    }
}

#[aoc_generator(day09)]
pub fn input_generator(input: &str) -> Data {
    Data(
        input
            .lines()
            .map(|l| {
                l.chars()
                    .map(|c| c.to_digit(10).unwrap().try_into().unwrap())
                    .collect()
            })
            .collect(),
    )
}

fn get_low_points(input: &Data) -> Vec<(i64, i64, u8)> {
    let mut ret = vec![];
    let (sx, sy) = input.shape();
    let sx = sx as i64;
    let sy = sy as i64;
    for x in 0..sx {
        for y in 0..sy {
            let c = input.get(x, y);
            if c.is_none() {
                continue;
            }
            let c = c.unwrap();
            let low_point = input
                .get(x + 1, y)
                .iter()
                .chain(input.get(x - 1, y).iter())
                .chain(input.get(x, y + 1).iter())
                .chain(input.get(x, y - 1).iter())
                .all(|v| v > &c);
            if low_point {
                ret.push((x, y, c))
            }
        }
    }
    ret
}

#[aoc(day09, part1)] // 468
pub fn solve_part1(input: &Data) -> usize {
    get_low_points(input)
        .iter()
        .map(|(_, _, v)| (*v as usize) + 1)
        .sum()
}

#[aoc(day09, part2)]
pub fn solve_part2(input: &Data) -> usize {
    let mut basins: BTreeSet<usize> = BTreeSet::new();
    let mut q = VecDeque::new();
    let mut visited = BTreeSet::<(i64, i64)>::new();
    for (x, y, c) in &get_low_points(input) {
        visited.clear();
        q.clear();
        q.push_back((*x, *y, *c));
        while let Some((x, y, c)) = q.pop_front() {
            if c == 9 || !visited.insert((x, y)) {
                continue;
            }
            for d in [-1, 1] {
                if let Some(v) = input.get(x + d, y) {
                    if v > c {
                        q.push_back((x + d, y, v));
                    }
                }
                if let Some(v) = input.get(x, y + d) {
                    if v > c {
                        q.push_back((x, y + d, v));
                    }
                }
            }
        }
        basins.insert(visited.len());
    }
    return basins.iter().rev().take(3).product();
}