use std::collections::HashSet;
use std::error::Error;
use std::fs::File;
use std::io::{self, BufRead};
use std::vec::Vec;

fn get_height(x: i32, y: i32, map: &Vec<Vec<u32>>) -> Option<u32> {
    if y >= 0 && y < map.len() as i32 && x >= 0 && x < map[y as usize].len() as i32 {
        Some(map[y as usize][x as usize])
    } else {
        None
    }
}

fn get_adjacent(x: i32, y: i32, map: &Vec<Vec<u32>>) -> Vec<u32> {
    let mut res = Vec::new();

    for (ax, ay) in [
        (x + 0, y - 1),
        (x + 0, y + 1),
        (x - 1, y + 0),
        (x + 1, y + 0),
    ] {
        if let Some(h) = get_height(ax, ay, map) {
            res.push(h)
        }
    }

    res
}

fn find_basin_size(x: i32, y: i32, visted: &mut HashSet<(i32, i32)>, map: &Vec<Vec<u32>>) -> u32 {
    if visted.contains(&(x, y)) {
        return 0;
    }

    visted.insert((x, y));

    let height = get_height(x, y, map);
    if height == Some(9) || height == None {
        return 0;
    }

    return 1
        + find_basin_size(x, y - 1, visted, map)
        + find_basin_size(x, y + 1, visted, map)
        + find_basin_size(x - 1, y, visted, map)
        + find_basin_size(x + 1, y, visted, map);
}

fn main() -> Result<(), Box<dyn Error>> {
    let file = File::open("inputs/day9.txt")?;
    let map: Vec<Vec<u32>> = io::BufReader::new(file)
        .lines()
        .map(|l| {
            l.unwrap()
                .chars()
                .map(|c| c.to_string().parse().unwrap())
                .collect()
        })
        .collect();

    let mut low_spots: Vec<u32> = Vec::new();
    let mut basin_sizes: Vec<u32> = Vec::new();
    for y in 0..map.len() {
        for x in 0..map[y].len() {
            let current_height = map[y][x];
            let adjacent = get_adjacent(x as i32, y as i32, &map);

            if adjacent.iter().all(|h| &current_height < h) {
                low_spots.push(current_height);

                let size = find_basin_size(x as i32, y as i32, &mut HashSet::new(), &map);
                basin_sizes.push(size);
            }
        }
    }

    println!("Low spots: {:?}", low_spots);
    let answer1 = low_spots.iter().fold(0, |a, b| a + b + 1);

    println!("Answer1: {}", answer1);

    basin_sizes.sort();
    basin_sizes.reverse();
    println!("basin sizes: {:?}", basin_sizes);

    let answer2 = basin_sizes[0] * basin_sizes[1] * basin_sizes[2];

    println!("Answer2: {}", answer2);

    Ok(())
}