use nom::{
    character::complete::{line_ending, one_of},
    combinator::{map, map_res},
    multi::{many1, separated_list1},
    IResult,
};
use std::fs;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_9.txt")?;
    let height_map = parse_height_map(&input).unwrap().1;
    let risk_level_sum: RiskLevel = height_map.risk_levels().into_iter().sum();
    dbg!(risk_level_sum);

    let mut basin_sizes: Vec<u32> = height_map.basins.iter().map(|basin| basin.size).collect();
    basin_sizes.sort_unstable_by(|a, b| b.cmp(a));
    dbg!(basin_sizes.iter().take(3).product::<u32>());

    Ok(())
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
struct Height(u8);
#[derive(
    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, derive_more::Add, derive_more::Sum,
)]
struct RiskLevel(u32);

impl From<Height> for RiskLevel {
    fn from(h: Height) -> Self {
        RiskLevel(h.0 as u32 + 1)
    }
}

#[derive(Debug, Default)]
struct Basin {
    size: u32,
}

#[derive(Debug)]
struct HeightMap {
    heights: Vec<Vec<Height>>,
    low_points: Vec<(usize, usize)>,
    basins: Vec<Basin>,
    basin_map: Vec<Vec<Option<usize>>>,
}

impl HeightMap {
    fn new(heights: Vec<Vec<Height>>) -> HeightMap {
        let mut height_map = HeightMap {
            heights,
            low_points: Vec::new(),
            basins: Vec::new(),
            basin_map: Vec::new(),
        };
        height_map.init_low_points();
        height_map.init_basins();
        height_map
    }
}

impl HeightMap {
    fn init_low_points(&mut self) {
        self.low_points = Vec::new();
        for y in 0..self.heights.len() {
            for x in 0..self.heights[y].len() {
                let current = self.heights[y][x];
                let mut is_low_point = true;
                let edges = [
                    (x as isize - 1, y as isize),
                    (x as isize + 1, y as isize),
                    (x as isize, y as isize - 1),
                    (x as isize, y as isize + 1),
                ];
                for edge in edges {
                    if self
                        .get_height(edge.0, edge.1)
                        .map_or(false, |other| other <= current)
                    {
                        is_low_point = false;
                    }
                }

                if is_low_point {
                    self.low_points.push((x, y));
                }
            }
        }
    }

    fn init_basins(&mut self) {
        for low_point in self.low_points.clone() {
            if self
                .get_basin(low_point.0 as isize, low_point.1 as isize)
                .is_some()
            {
                continue;
            }

            let mut basin = Basin::default();
            let basin_index = self.basins.len();
            self.set_basin(basin_index, low_point.0, low_point.1);
            basin.size += 1;
            let mut boundary = vec![low_point];
            while boundary.len() > 0 {
                let (x, y) = boundary.pop().unwrap();
                let edges = [
                    (x as isize - 1, y as isize),
                    (x as isize + 1, y as isize),
                    (x as isize, y as isize - 1),
                    (x as isize, y as isize + 1),
                ];
                for edge in edges {
                    if self
                        .get_height(edge.0, edge.1)
                        .map_or(false, |other| other != Height(9))
                        && self.get_basin(edge.0, edge.1).is_none()
                    {
                        let x = edge.0 as usize;
                        let y = edge.1 as usize;
                        self.set_basin(basin_index, x, y);
                        basin.size += 1;
                        boundary.push((x, y));
                    }
                }
            }
            self.basins.push(basin);
        }
    }

    fn get_height(&self, x: isize, y: isize) -> Option<Height> {
        if x < 0 || y < 0 {
            None
        } else {
            let x = x as usize;
            let y = y as usize;
            self.heights.get(y).and_then(|row| row.get(x)).cloned()
        }
    }

    fn get_basin(&self, x: isize, y: isize) -> Option<usize> {
        if x < 0 || y < 0 {
            None
        } else {
            let x = x as usize;
            let y = y as usize;
            self.basin_map
                .get(y)
                .and_then(|row| row.get(x))
                .cloned()
                .flatten()
        }
    }

    fn set_basin(&mut self, basin: usize, x: usize, y: usize) {
        while self.basin_map.len() <= y {
            self.basin_map.push(Vec::new());
        }
        while self.basin_map[y].len() <= x {
            self.basin_map[y].push(None);
        }
        self.basin_map[y][x] = Some(basin);
    }

    fn risk_levels(&self) -> Vec<RiskLevel> {
        self.low_points
            .iter()
            .copied()
            .map(|(x, y)| self.heights[y][x].clone().into())
            .collect()
    }
}

fn parse_height_map(input: &str) -> IResult<&str, HeightMap> {
    map(separated_list1(line_ending, parse_row), HeightMap::new)(input)
}

fn parse_row(input: &str) -> IResult<&str, Vec<Height>> {
    many1(parse_height)(input)
}

fn parse_height(input: &str) -> IResult<&str, Height> {
    map_res(one_of("0123456789"), |digit| {
        digit.to_string().parse().map(Height)
    })(input)
}