use std::collections::VecDeque;

use aoc_runner_derive::aoc;
use grid::Grid;

struct RaceTrack {
    map: Grid<u8>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
struct State {
    pos: (i64, i64),
    cost: u64,
}

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
struct CheatState {
    s: State,
    p: Vec<(i64, i64)>,
}

const DIRECTIONS: [(i64, i64); 4] = [(-1, 0), (1, 0), (0, -1), (0, 1)];

impl RaceTrack {
    fn valid_moves<'a>(&'a self, CheatState { s: state, p }: &'a CheatState) -> impl Iterator<Item = CheatState> + 'a {
        let mut new_path = p.clone();
        new_path.push(state.pos);
        DIRECTIONS
            .iter()
            .map(|dir| (state.pos.0 + dir.0, state.pos.1 + dir.1))
            .filter_map(move |pos| match &self.map.get(&pos) {
                Some(b'.') | Some(b'S') | Some(b'E') => Some(CheatState {
                    p: new_path.clone(),
                    s: State {
                        pos,
                        cost: state.cost + 1,
                    },
                }),
                _ => None,
            })
    }
    fn path_costs(&self, start: (i64, i64), goal: (i64, i64)) -> (Vec<(i64, i64)>, Grid<Option<u64>>) {
        let mut queue = VecDeque::new();
        let mut visited = self.map.same_shape(None);

        let start_state = CheatState {
            s: State { pos: start, cost: 0 },
            p: Vec::new(),
        };
        visited.set(&start, Some(0));
        queue.push_back(start_state);

        while let Some(state) = queue.pop_front() {
            if state.s.pos == goal {
                let mut final_path = state.p;
                final_path.push(goal);
                return (final_path, visited);
            }

            let moves = self.valid_moves(&state);
            for new_state in moves {
                if visited.get(&new_state.s.pos).unwrap().is_some() {
                    continue;
                }
                visited.set(&new_state.s.pos, Some(new_state.s.cost));
                queue.push_back(new_state);
            }
        }
        panic!("no path");
    }

    fn find_cheats(
        &self,
        path: &Vec<(i64, i64)>,
        costs: &Grid<Option<u64>>,
        min: u64,
    ) -> Vec<((i64, i64), (i64, i64), u64)> {
        let mut cheats = Vec::new();
        for pos in path {
            let local_cost = costs.get(pos).unwrap().unwrap();
            for ofs in DIRECTIONS {
                let cheat_start = (pos.0 + ofs.0, pos.1 + ofs.1);
                let cheat_exit = (pos.0 + ofs.0 * 2, pos.1 + ofs.1 * 2);
                if let Some(Some(cheat_cost)) = costs.get(&cheat_exit) {
                    if *cheat_cost > local_cost + 2 {
                        let cheat_savings = cheat_cost - local_cost - 2;
                        if cheat_savings >= min {
                            cheats.push((cheat_start, cheat_exit, cheat_savings));
                        }
                    }
                }
            }
        }
        cheats
    }

    fn taxi_dist(&self, from: &(i64, i64), to: &(i64, i64)) -> Option<u64> {
        if self.map.is_valid(to) {
            Some(from.0.abs_diff(to.0) + from.1.abs_diff(to.1))
        } else {
            None
        }
    }

    fn find_cheats_n(
        &self,
        path: &Vec<(i64, i64)>,
        costs: &Grid<Option<u64>>,
        max_length: u64,
        min: u64,
    ) -> Vec<((i64, i64), (i64, i64))> {
        let mut cheats = Vec::new();
        let mut solutions = self.map.clone();

        for pos in path {
            let from_cost = costs.get(pos).unwrap().unwrap();
            for x in 0..self.map.width() as i64 {
                for y in 0..self.map.height() as i64 {
                    if let Some(dist) = self.taxi_dist(pos, &(x, y)) {
                        if dist <= max_length && dist >= 2 {
                            if let Some(to_cost) = costs.get(&(x, y)).unwrap() {
                                solutions.set(pos, b'O');
                                solutions.set(&(x, y), b'O');
                                if *to_cost > (from_cost + dist) && (to_cost - (from_cost + dist) >= min) {
                                    cheats.push((*pos, (x, y)));
                                }
                            }
                        }
                    }
                    solutions.set(&(x, y), *self.map.get(&(x, y)).unwrap());
                }
            }
            solutions.set(pos, *self.map.get(pos).unwrap());
        }
        cheats
    }
}

fn parse(input: &str) -> RaceTrack {
    let map = input.as_bytes().into();
    RaceTrack { map }
}

fn part1_impl(input: &str, cheat_min: u64) -> i64 {
    let track = parse(input);
    let start = track.map.find(&b'S').unwrap();
    let goal = track.map.find(&b'E').unwrap();
    let (best_path, costs) = track.path_costs(start.into(), goal.into());
    let cheats = track.find_cheats(&best_path, &costs, cheat_min);

    cheats.len() as i64
}

fn part2_impl(input: &str, max_length: u64, cheat_min: u64) -> i64 {
    let track = parse(input);
    let start = track.map.find(&b'S').unwrap();
    let goal = track.map.find(&b'E').unwrap();
    let (best_path, costs) = track.path_costs(start.into(), goal.into());
    let cheats = track.find_cheats_n(&best_path, &costs, max_length, cheat_min);

    cheats.len() as i64
}

#[aoc(day20, part1)]
pub fn part1(input: &str) -> i64 {
    part1_impl(input, 100)
}

#[aoc(day20, part2)]
pub fn part2(input: &str) -> i64 {
    part2_impl(input, 20, 100)
}

#[cfg(test)]
mod tests {
    use super::*;
    const EXAMPLE: &str = "###############
#...#...#.....#
#.#.#.#.#.###.#
#S#...#.#.#...#
#######.#.#.###
#######.#.#...#
#######.#.###.#
###..E#...#...#
###.#######.###
#...###...#...#
#.#####.#.###.#
#.#...#.#.#...#
#.#.#.#.#.#.###
#...#...#...###
###############";

    #[test]
    fn part1_example() {
        assert_eq!(part1_impl(EXAMPLE, 0), 44);
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2_impl(EXAMPLE, 2, 0), 44);
        assert_eq!(part2_impl(EXAMPLE, 20, 50), 285);
    }
}