use {
    anyhow::{
        anyhow,
        Context,
        Result,
    },
    std::{
        env::args,
        fs::File,
        io::{
            BufRead,
            BufReader,
        },
        path::Path,
    },
};

#[derive(Clone,Copy,Debug)]
enum Direction {
    North,
    South,
    East,
    West,
}

#[derive(Clone,Copy,Debug)]
struct Point {
    x: usize,
    y: usize,
}

impl Point {
    fn neighbour(&self, dir: Direction) -> Point {
        match dir {
            Direction::North => Point {x: self.x, y: self.y.overflowing_sub(1).0},
            Direction::South => Point {x: self.x, y: self.y + 1},
            Direction::East  => Point {x: self.x + 1, y: self.y},
            Direction::West  => Point {x: self.x.overflowing_sub(1).0, y: self.y},
        }
    }
}

#[derive(Debug)]
struct HeightMap {
    size: Point,
    levels: Vec<u8>,
}

#[derive(Debug)]
struct VisitMap {
    size: Point,
    visited: Vec<bool>,
}

impl VisitMap {
    fn new(size: Point) -> Self{
        let visited = vec![false; size.x*size.y];
        Self {
            size,
            visited,
        }
    }

    fn is_visited(&self, pos: Point) -> bool {
        if pos.x >= self.size.x || pos.y >= self.size.y {
            return true;
        }

        self.visited.get(pos.y * self.size.x + pos.x).copied().expect("Invalid index")
    }

    fn visit(&mut self, pos: Point) {
        self.visited[pos.y * self.size.x + pos.x] = true;
    }
}

impl HeightMap {
    fn level(&self, pos: &Point) -> Result<u8> {
        if pos.x >= self.size.x || pos.y >= self.size.y {
            return Err(anyhow!("Invalid index {}, {}", pos.x, pos.y));
        }
        self.levels.get(pos.y * self.size.x + pos.x).copied()
            .ok_or_else(|| anyhow!("Invalid index {}, {}", pos.x, pos.y))
    }

    fn is_lowpoint(&self, pos: &Point) -> bool {
        let us = self.level(pos).unwrap_or(0);

        let mut ok_count = 0;
        let directions = [Direction::North, Direction::South, Direction::East, Direction::West];

        directions.into_iter().for_each(|dir| {
            let neighbour = pos.neighbour(dir);
            match self.level(&neighbour) {
                Ok(them) => if us < them { ok_count += 1; },
                Err(_) => { ok_count += 1; },
            }
        });
        ok_count == directions.into_iter().count()
    }

    fn is_highpoint(&self, pos: &Point) -> bool {
        self.level(pos).unwrap_or(9) == 9
    }

    fn basin_size(&self, start: &Point) -> Result<usize> {
        let mut unexplored = vec![*start];
        let mut pos = None;
        let mut size = 0;

        let mut seen = VisitMap::new(self.size);

        while !unexplored.is_empty() {
            if pos.is_none() {
                if let Some(new_pos) = unexplored.pop() {
                    pos = Some(new_pos);
                    seen.visit(new_pos);
                    size += 1;
                } else {
                    unreachable!();
                }
            }

            if let Some(cur_pos) = pos {
                for dir in [Direction::North, Direction::South, Direction::East, Direction::West] {
                    if seen.is_visited(cur_pos.neighbour(dir)) {
                        continue;
                    }

                    let neighbour = cur_pos.neighbour(dir);
                    if self.is_highpoint(&neighbour) {
                        seen.visit(neighbour);
                        continue;
                    } else {
                        seen.visit(neighbour);
                        unexplored.push(neighbour);
                        continue;
                    }
                }
                if seen.is_visited(cur_pos.neighbour(Direction::North)) &&
                    seen.is_visited(cur_pos.neighbour(Direction::South)) &&
                    seen.is_visited(cur_pos.neighbour(Direction::East)) &&
                    seen.is_visited(cur_pos.neighbour(Direction::West))
                {
                    pos = None;
                }
            } else {
                unreachable!();
            }
        }

        Ok(size)
    }

    fn risk_level(&self, pos: &Point) -> Result<usize> {
        let level = self.level(pos)?;

        Ok(level as usize + 1)
    }
}

fn read_input<T: AsRef<Path>>(filename: T) -> Result<HeightMap> {
    let reader = BufReader::new(File::open(filename)?);
    let mut x_opt = None;
    let mut y = 0;

    let levels = reader.lines().map(|input| {
        let i = input?;
        let o: Vec<u8> = i.chars().map(|column| {
            let level: u8 = column.to_digit(10).ok_or(anyhow!("Invalid digit"))? as u8;
            Ok(level)
        }).collect::<Result<Vec<u8>>>()?;
        match (x_opt, o.len()) {
            (None, len) => x_opt = Some(len),
            (Some(this), other) => if this != other {
                return Err(anyhow!("Inconsistent map width encountered: {} != {}", this, other));
            },
        }
        y += 1;
        Ok(o)
    }).collect::<Result<Vec<Vec<u8>>>>()?.into_iter().flatten().collect();

    let x = x_opt.ok_or(anyhow!("Could not determine size of heightmap."))?;
    let size = Point {
        x,
        y,
    };
    Ok(HeightMap {
        size,
        levels,
    })
}

fn part1(map: &HeightMap) -> Result<usize> {
    let mut risk_level = 0;
    for y in 0..map.size.y {
        for x in 0..map.size.x {
            if map.is_lowpoint(&Point {x, y}) {
                risk_level += map.risk_level(&Point {x, y})?;
            }
        }
    }
    Ok(risk_level)
}

fn part2(map: &HeightMap) -> Result<usize> {
    let mut basin_sizes = vec![];
    for y in 0..map.size.y {
        for x in 0..map.size.x {
            if map.is_lowpoint(&Point {x, y}) {
                basin_sizes.push(map.basin_size(&Point {x, y})?);
            }
        }
    }
    basin_sizes.sort_unstable();

    Ok(basin_sizes.iter().rev().take(3).product())
}

fn main() -> Result<()> {
    let ( do_part_1, do_part_2 ) = aoc::do_parts();

    let filename = args().nth(1).ok_or(anyhow!("Missing input filename"))?;
    let input = read_input(filename).context("Could not read input")?;
    if do_part_1 {
        let solution = part1(&input).context("No solution for part 1")?;
        println!("Part1, solution found to be: {}", solution);
    }
    if do_part_2 {
        let solution = part2(&input).context("No solution for part 2")?;
        println!("Part2, solution found to be: {}", solution);
    }
    Ok(())
}