use std::collections::HashSet;
use std::io::stdin;
use std::io::BufRead;

fn main() {
    let arg = std::env::args().nth(1).unwrap_or("part2".to_string());
    match arg.as_str() {
        "part1" => part1(),
        _ => part2(),
    }
}

fn part1() {
    println!("Part 1");
    let mut numbers = Vec::new();
    let mut width = 0;
    let mut height = 0;
    for line in stdin().lock().lines() {
        let line = line.unwrap().replace(" ", "").replace("_", "");
        width = line.len();
        println!("{}", line);

        let mut parsed: Vec<u8> = line
            .chars()
            .map(|c| c.to_string().parse().unwrap())
            .collect();
        numbers.append(&mut parsed);
        height += 1;
    }

    dbg!(width, height);

    let field = Field::new(numbers, width, height);

    let mut risk_level: i32 = 0;

    for x in 0..height {
        for y in 0..width {
            let value = field.get((x, y));
            let neighbors = field.neighbors((x, y));
            if neighbors.iter().all(|&n| n > value) {
                risk_level += 1 + *value as i32;
                print!("\x1b[31m{}\x1b[0m", value);
            } else {
                print!("{}", value);
            }
        }
        println!();
    }

    dbg!(risk_level);
}

fn part2() {
    println!("Part 2");
    let mut numbers = Vec::new();
    let mut width = 0;
    let mut height = 0;
    for line in stdin().lock().lines() {
        let line = line.unwrap().replace(" ", "").replace("_", "");
        width = line.len();
        println!("{}", line);

        let mut parsed: Vec<u8> = line
            .chars()
            .map(|c| c.to_string().parse().unwrap())
            .collect();
        numbers.append(&mut parsed);
        height += 1;
    }

    dbg!(width, height);

    let field = Field::new(numbers, width, height);

    let mut basin_sizes = Vec::new();

    for x in 0..height {
        for y in 0..width {
            let value = field.get((x, y));
            let neighbors = field.neighbors((x, y));
            if neighbors.iter().all(|&n| n > value) {
                basin_sizes.push(field.basin_size((x, y)));
                print!("\x1b[31m{}\x1b[0m", value);
            } else if *value == 9 {
                print!("\x1b[33m{}\x1b[0m", value);
            } else {
                print!("\x1b[37m{}\x1b[0m", value);
            }
        }
        println!();
    }

    basin_sizes.sort();
    dbg!(&basin_sizes);

    let solution = basin_sizes.iter().rev().take(3).fold(1, |acc, el| acc * el);
    dbg!(solution);
}

struct Field {
    width: usize,
    height: usize,
    numbers: Vec<u8>,
}

impl<'a> Field {
    fn new(numbers: Vec<u8>, width: usize, height: usize) -> Field {
        Field {
            width,
            height,
            numbers,
        }
    }

    fn get(&'a self, pos: (usize, usize)) -> &'a u8 {
        let (x, y) = pos;
        let idx = x * self.width + y;
        self.numbers.get(idx).unwrap()
    }

    fn neighbor_coords(&'a self, pos: (usize, usize)) -> Vec<(usize, usize)> {
        let (x, y) = pos;
        let mut coords = Vec::new();
        let is_top = x == 0;
        let is_bot = x == self.height - 1;
        let is_left = y % self.width == 0;
        let is_right = (y + 1) % self.width == 0;

        if !is_top {
            coords.push((x - 1, y));
        }
        if !is_bot {
            coords.push((x + 1, y));
        }
        if !is_left {
            coords.push((x, y - 1))
        }
        if !is_right {
            coords.push((x, y + 1))
        }
        coords
    }

    fn neighbors(&'a self, pos: (usize, usize)) -> Vec<&'a u8> {
        self.neighbor_coords(pos)
            .iter()
            .map(|&pos| self.get(pos))
            .collect()
    }

    fn basin_size(&self, pos: (usize, usize)) -> i32 {
        let mut size = 0;
        let mut coords_to_check = HashSet::new();
        let mut coords_checked = HashSet::new();

        coords_to_check.insert(pos);

        while !coords_to_check.is_empty() {
            let pos = *coords_to_check.iter().next().unwrap();
            coords_to_check.remove(&pos);
            coords_checked.insert(pos);
            if *self.get(pos) == 9 {
                continue;
            }
            size += 1;
            let neighbor_set = HashSet::from_iter(self.neighbor_coords(pos));
            coords_to_check.extend(&neighbor_set - &coords_checked);
            // for neighbor in self.neighbor_coords(pos) {
            //     if !coords_checked.contains(neighbor){
            //         coords_to_check.add(neighbor)
            //     }
            // }
        }
        size
    }
}