use std::collections::HashMap;
use std::collections::HashSet;

#[path = "../lib.rs"]
mod lib;

#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)]
enum Dir {
    Up,
    Right,
    Down,
    Left,
}

impl Dir {
    fn apply(self: &Self, x: usize, y: usize) -> (usize, usize) {
        return match self {
            Self::Up => (x, y - 1),
            Self::Right => (x + 1, y),
            Self::Down => (x, y + 1),
            Self::Left => (x - 1, y),
        };
    }

    fn unapply(self: &Self, x: usize, y: usize) -> (usize, usize) {
        return match self {
            Self::Up => (x, y + 1),
            Self::Right => (x - 1, y),
            Self::Down => (x, y - 1),
            Self::Left => (x + 1, y),
        };
    }
}

type Pos = (usize, usize, Dir);

fn parse_input() -> (Pos, (usize, usize), Vec<Vec<bool>>) {
    let mut start_x = 0;
    let mut start_y = 0;
    let mut end_x = 0;
    let mut end_y = 0;
    let mut map = vec![];

    for (y, line) in lib::iter_input().enumerate() {
        map.push(
            line.bytes()
                .enumerate()
                .map(|(x, b)| match b {
                    b'#' => false,
                    b'.' => true,
                    b'S' => {
                        start_x = x;
                        start_y = y;
                        true
                    }
                    b'E' => {
                        end_x = x;
                        end_y = y;
                        true
                    }
                    _ => unreachable!(),
                })
                .collect(),
        );
    }

    return ((start_x, start_y, Dir::Right), (end_x, end_y), map);
}

fn push1(pos: Pos, score: usize, cache: &mut HashMap<Pos, usize>, queue: &mut Vec<Pos>) {
    if score < *cache.get(&pos).unwrap_or(&usize::MAX) {
        cache.insert(pos, score);
        queue.push(pos);
    }
}

fn part1(
    start: Pos,
    end: (usize, usize),
    map: &Vec<Vec<bool>>,
    cache: &mut HashMap<Pos, usize>,
) -> usize {
    let mut queue = vec![];
    let mut best_score = usize::MAX;

    push1(start, 0, cache, &mut queue);

    while let Some((x, y, dir)) = queue.pop() {
        let score = *cache.get(&(x, y, dir)).unwrap();

        if (x, y) == end {
            if score < best_score {
                best_score = score;
            }
            continue;
        }

        let (x2, y2) = dir.apply(x, y);
        if map[y2][x2] {
            push1((x2, y2, dir), score + 1, cache, &mut queue);
        }

        match dir {
            Dir::Up | Dir::Down => {
                push1((x, y, Dir::Right), score + 1000, cache, &mut queue);
                push1((x, y, Dir::Left), score + 1000, cache, &mut queue);
            }
            Dir::Right | Dir::Left => {
                push1((x, y, Dir::Up), score + 1000, cache, &mut queue);
                push1((x, y, Dir::Down), score + 1000, cache, &mut queue);
            }
        }

        // PERF: makes a major difference, probably because preliminary
        // results are not propagated
        queue.sort_by_key(|(xx, yy, _)| xx.abs_diff(end.0) + yy.abs_diff(end.1));
    }

    return best_score;
}

fn push2(pos: Pos, score: usize, cache: &HashMap<Pos, usize>, queue: &mut Vec<Pos>) {
    if score == *cache.get(&pos).unwrap_or(&usize::MAX) {
        queue.push(pos);
    }
}

fn part2(end: (usize, usize), best_score: usize, cache: &mut HashMap<Pos, usize>) -> usize {
    let mut queue = vec![];
    let mut paths = HashSet::new();

    for dir in [Dir::Up, Dir::Right, Dir::Down, Dir::Left] {
        push2((end.0, end.1, dir), best_score, cache, &mut queue);
    }

    while let Some((x, y, dir)) = queue.pop() {
        if let Some(score) = cache.remove(&(x, y, dir)) {
            paths.insert((x, y));

            let (x2, y2) = dir.unapply(x, y);
            push2((x2, y2, dir), score - 1, &cache, &mut queue);

            match dir {
                Dir::Up | Dir::Down => {
                    push2((x, y, Dir::Right), score - 1000, cache, &mut queue);
                    push2((x, y, Dir::Left), score - 1000, cache, &mut queue);
                }
                Dir::Right | Dir::Left => {
                    push2((x, y, Dir::Up), score - 1000, cache, &mut queue);
                    push2((x, y, Dir::Down), score - 1000, cache, &mut queue);
                }
            }
        }
    }

    return paths.len();
}

fn main() {
    let (start, end, map) = parse_input();
    let mut cache = HashMap::new();

    let best_score = part1(start, end, &map, &mut cache);
    let paths = part2(end, best_score, &mut cache);

    println!("part1: {}", best_score);
    println!("part2: {}", paths);
}