use std::{
    cmp::{Ordering, Reverse},
    collections::{BinaryHeap, HashMap},
};

use aoc_runner_derive::{aoc, aoc_generator};

struct PuzzleInput {
    heights: Vec<Vec<u8>>,
    start: (usize, usize),
    end: (usize, usize),
}

#[aoc_generator(day12)]
fn generator(input: &str) -> PuzzleInput {
    let mut start = None;
    let mut end = None;

    let nodes = input
        .lines()
        .enumerate()
        .map(|(i, l)| {
            l.chars()
                .enumerate()
                .map(|(j, c)| match c {
                    'S' => {
                        assert!(start.replace((i, j)).is_none());
                        1
                    }
                    'E' => {
                        assert!(end.replace((i, j)).is_none());
                        26
                    }
                    'a'..='z' => c as u8 - b'a' + 1,
                    _ => panic!(),
                })
                .collect()
        })
        .collect();

    PuzzleInput {
        heights: nodes,
        start: start.unwrap(),
        end: end.unwrap(),
    }
}

fn neighbours(coord: (usize, usize), heights: &[Vec<u8>]) -> impl Iterator<Item = (usize, usize)> {
    let current_height = heights[coord.0][coord.1];

    let a = (coord.0 + 1, coord.1);
    let a = heights
        .get(a.0)
        .and_then(|v| (v[a.1] >= current_height - 1).then_some(a));

    let b = coord.0.checked_sub(1).map(|v| (v, coord.1));
    let b = b.filter(|b| heights[b.0][b.1] >= current_height - 1);

    let c = (coord.0, coord.1 + 1);
    let c = heights[c.0]
        .get(c.1)
        .and_then(|v| (*v >= current_height - 1).then_some(c));

    let d = coord.1.checked_sub(1).map(|v| (coord.0, v));
    let d = d.filter(|d| heights[d.0][d.1] >= current_height - 1);

    a.into_iter().chain(b).chain(c).chain(d)
}

struct CmpByFirst<T, U>(T, U);

impl<T: PartialEq, U> PartialEq for CmpByFirst<T, U> {
    fn eq(&self, other: &Self) -> bool {
        self.0.eq(&other.0)
    }
}

impl<T: Eq, U> Eq for CmpByFirst<T, U> {}

impl<T: PartialOrd, U> PartialOrd for CmpByFirst<T, U> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

impl<T: Ord, U> Ord for CmpByFirst<T, U> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.cmp(&other.0)
    }
}

fn fill_dist_map(input: &PuzzleInput) -> HashMap<(usize, usize), u32> {
    let mut to_visit = BinaryHeap::new();
    to_visit.push(CmpByFirst(Reverse(0), input.end));

    let mut map = HashMap::new();
    map.insert(input.end, 0);

    while let Some(CmpByFirst(Reverse(dist), coord)) = to_visit.pop() {
        for neighbour in neighbours(coord, &input.heights) {
            map.entry(neighbour)
                .and_modify(|v| *v = (*v).min(dist + 1))
                .or_insert_with(|| {
                    to_visit.push(CmpByFirst(Reverse(dist + 1), neighbour));

                    dist + 1
                });
        }
    }

    map
}

#[aoc(day12, part1)]
fn part1(input: &PuzzleInput) -> u32 {
    fill_dist_map(input)[&input.start]
}

#[aoc(day12, part2)]
fn part2(input: &PuzzleInput) -> u32 {
    let map = fill_dist_map(input);

    input
        .heights
        .iter()
        .enumerate()
        .flat_map(|(i, v)| v.iter().enumerate().map(move |(j, v)| ((i, j), v)))
        .filter_map(|(coord, h)| (*h == 1).then_some(coord))
        .map(|coord| *map.get(&coord).unwrap_or(&u32::MAX))
        .min()
        .unwrap()
}