use std::{
    collections::{BinaryHeap, HashMap},
    fmt::Debug,
    fs::DirEntry,
};

use itertools::Itertools;

pub fn run() {
    let input = include_str!("../input/day17.txt");
    let map: Map = input.into();

    let path = map.dijkstra_ish(
        (0, 0),
        (map.width as isize - 1, map.height as isize - 1),
        Map::part1_edges,
    );
    dbg!(map.path_cost(&path.unwrap()));

    let path2 = map.dijkstra_ish(
        (0, 0),
        (map.width as isize - 1, map.height as isize - 1),
        Map::part2_edges,
    );
    dbg!(map.part2_path_cost(&path2.unwrap()));
}

#[derive(Clone)]
struct Map {
    costs: Vec<Vec<u8>>,
    width: usize,
    height: usize,
}

impl From<&str> for Map {
    fn from(value: &str) -> Self {
        let costs = value
            .trim()
            .lines()
            .map(|l| {
                l.trim()
                    .chars()
                    .map(|c| c.to_digit(10).unwrap() as u8)
                    .collect::<Vec<_>>()
            })
            .collect::<Vec<_>>();
        let width = costs[0].len();
        let height = costs.len();
        Self {
            costs,
            width,
            height,
        }
    }
}

impl Debug for Map {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Map(\n")?;
        for row in self.costs.iter() {
            for cost in row {
                write!(f, "{}", cost)?;
            }
        }
        write!(f, ")")
    }
}

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

impl Direction {
    fn all() -> [Direction; 4] {
        [
            Direction::Up,
            Direction::Down,
            Direction::Left,
            Direction::Right,
        ]
    }

    fn move_pos(&self, pos: (isize, isize)) -> (isize, isize) {
        self.move_pos_count(pos, 1)
    }

    fn move_pos_count(&self, (col, row): (isize, isize), count: isize) -> (isize, isize) {
        match self {
            Direction::Up => (col, row - count),
            Direction::Down => (col, row + count),
            Direction::Left => (col - count, row),
            Direction::Right => (col + count, row),
        }
    }

    fn opposite(&self) -> Self {
        match self {
            Direction::Up => Direction::Down,
            Direction::Down => Direction::Up,
            Direction::Left => Direction::Right,
            Direction::Right => Direction::Left,
        }
    }
}

#[derive(Debug, Clone, Copy)]
struct Edge {
    direction: Direction,
    dest: (isize, isize),
    cost: u8,
    steps: usize,
}

impl Map {
    fn in_bounds(&self, (col, row): (isize, isize)) -> bool {
        col >= 0 && col < self.width as isize && row >= 0 && row < self.height as isize
    }

    fn part1_edges(
        &self,
        pos: (isize, isize),
        incoming_dir: Option<(Direction, usize)>,
    ) -> Vec<Edge> {
        Direction::all()
            .iter()
            .filter_map(|&d| {
                if incoming_dir.is_some_and(|(x, c)| x == d.opposite() || (x == d && c >= 3)) {
                    return None;
                }
                let moved = d.move_pos(pos);
                if !self.in_bounds(moved) {
                    return None;
                }
                Some(Edge {
                    direction: d,
                    dest: moved,
                    cost: self.costs[moved.1 as usize][moved.0 as usize],
                    steps: 1,
                })
            })
            .collect()
    }

    fn part2_edges(
        &self,
        pos: (isize, isize),
        incoming_dir: Option<(Direction, usize)>,
    ) -> Vec<Edge> {
        let mut v = vec![];
        for d in Direction::all() {
            if incoming_dir.is_some_and(|(x, _)| x == d.opposite()) {
                continue;
            }
            for dist in 4..=10 {
                if incoming_dir.is_some_and(|(x, c)| x == d && (c + dist) > 10) {
                    continue;
                }

                let moved = d.move_pos_count(pos, dist as isize);
                if !self.in_bounds(moved) {
                    continue;
                }

                let costs_sum = (1..=dist)
                    .map(|c| {
                        let moved = d.move_pos_count(pos, c as isize);
                        self.costs[moved.1 as usize][moved.0 as usize]
                    })
                    .sum();
                v.push(Edge {
                    direction: d,
                    dest: moved,
                    cost: costs_sum,
                    steps: dist,
                });
            }
        }
        v
    }

    fn dijkstra_ish(
        &self,
        start: (isize, isize),
        end: (isize, isize),
        edges_fn: impl Fn(&Self, (isize, isize), Option<(Direction, usize)>) -> Vec<Edge>,
    ) -> Option<Vec<(isize, isize)>> {
        #[derive(Debug, Clone, Copy, PartialEq, Eq)]
        struct QueueEntry {
            neg_cost: isize,
            inbound: Option<(Direction, usize)>,
            pos: (isize, isize),
        }
        impl Ord for QueueEntry {
            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
                self.neg_cost.cmp(&other.neg_cost)
            }
        }
        impl PartialOrd for QueueEntry {
            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
                Some(self.cmp(other))
            }
        }

        let mut queue = BinaryHeap::<QueueEntry>::new();

        #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
        struct HistoryKey {
            pos: (isize, isize),
            inbound: Option<(Direction, usize)>,
        }

        let mut distances = HashMap::<HistoryKey, usize>::new();
        let mut prev = HashMap::<HistoryKey, HistoryKey>::new();

        distances.insert(
            HistoryKey {
                pos: start,
                inbound: None,
            },
            0,
        );
        queue.push(QueueEntry {
            neg_cost: 0,
            inbound: None,
            pos: start,
        });

        while let Some(entry) = queue.pop() {
            let entry_key = HistoryKey {
                pos: entry.pos,
                inbound: entry.inbound,
            };
            for edge in edges_fn(self, entry.pos, entry.inbound) {
                let alt = distances.get(&entry_key).unwrap_or(&usize::MAX) + edge.cost as usize;
                let new_inbound_dir_count = if let Some((d, c)) = entry.inbound {
                    if d == edge.direction {
                        c + edge.steps
                    } else {
                        edge.steps
                    }
                } else {
                    edge.steps
                };
                let edge_key = HistoryKey {
                    pos: edge.dest,
                    inbound: Some((edge.direction, new_inbound_dir_count)),
                };
                if alt < *distances.get(&edge_key).unwrap_or(&usize::MAX) {
                    distances.insert(edge_key, alt);
                    prev.insert(edge_key, entry_key);
                    let new_entry = QueueEntry {
                        neg_cost: -(alt as isize),
                        inbound: Some((edge.direction, new_inbound_dir_count)),
                        pos: edge.dest,
                    };
                    queue.push(new_entry);
                }
            }
        }

        // for row in 0..self.height {
        //     for col in 0..self.width {
        //         println!("({}, {}) ", col, row);
        //         for d in Direction::all() {
        //             println!(
        //                 "\t{:?} 0={}/{:?}, 1={}/{:?}, 2={}/{:?}, 3={}/{:?}",
        //                 d,
        //                 distances[row][col][d as usize][0],
        //                 prev[row][col][d as usize][0],
        //                 distances[row][col][d as usize][1],
        //                 prev[row][col][d as usize][1],
        //                 distances[row][col][d as usize][2],
        //                 prev[row][col][d as usize][2],
        //                 distances[row][col][d as usize][3],
        //                 prev[row][col][d as usize][3],
        //             );
        //         }
        //     }
        // }

        let mut path = vec![];
        let min = distances
            .iter()
            .filter(|(k, _)| k.pos == end)
            .min_by_key(|(_, d)| **d)
            .unwrap()
            .0;
        let mut u = Some(min);
        while let Some(key) = u {
            println!("({}, {})", key.pos.0, key.pos.1);
            path.insert(0, key.pos);
            u = prev.get(key);
        }
        Some(path)
    }

    fn path_cost(&self, path: &[(isize, isize)]) -> usize {
        path.iter()
            .skip(1)
            .map(|&(col, row)| self.costs[row as usize][col as usize] as usize)
            .sum()
    }

    fn part2_path_cost(&self, path: &[(isize, isize)]) -> usize {
        let mut cost = 0;
        for (a, b) in path.iter().zip(path.iter().skip(1)) {
            // a->b is always a cardinal direction, so one of these loops will only have one
            // iteration
            for row in a.1.min(b.1)..=a.1.max(b.1) {
                for col in a.0.min(b.0)..=a.0.max(b.0) {
                    cost += self.costs[row as usize][col as usize] as usize;
                }
            }

            // don't double count the beginning of each path segment
            // also don't count the beginning of the first segment
            cost -= self.costs[a.1 as usize][a.0 as usize] as usize;
        }
        cost
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_max_single_direction() {
        let map: Map = r#"
11111
22222
            "#
        .into();
        assert_eq!(
            map.dijkstra_ish((0, 0), (4, 1), Map::part1_edges),
            Some(vec![(0, 0), (1, 0), (2, 0), (3, 0), (3, 1), (4, 1)])
        );
    }

    #[test]
    fn test_example_small() {
        let map: Map = r#"
24134
32154
        "#
        .into();
        assert_eq!(
            map.dijkstra_ish(
                (0, 0),
                (map.width as isize - 1, map.height as isize - 1),
                Map::part1_edges
            ),
            Some(vec![(0, 0), (1, 0), (2, 0), (2, 1), (3, 1), (4, 1)])
        );
    }

    #[test]
    fn test_example() {
        let map: Map = r#"
2413432311323
3215453535623
3255245654254
3446585845452
4546657867536
1438598798454
4457876987766
3637877979653
4654967986887
4564679986453
1224686865563
2546548887735
4322674655533
            "#
        .into();
        let path = dbg!(map.dijkstra_ish(
            (0, 0),
            (map.width as isize - 1, map.height as isize - 1),
            Map::part1_edges
        ));
        assert_eq!(map.path_cost(&path.unwrap()), 102);
    }

    #[test]
    fn test_part2_example1() {
        let map: Map = r#"
2413432311323
3215453535623
3255245654254
3446585845452
4546657867536
1438598798454
4457876987766
3637877979653
4654967986887
4564679986453
1224686865563
2546548887735
4322674655533
            "#
        .into();
        let path = dbg!(map.dijkstra_ish(
            (0, 0),
            (map.width as isize - 1, map.height as isize - 1),
            Map::part2_edges
        ));
        assert_eq!(map.part2_path_cost(&path.unwrap()), 94);
    }

    #[test]
    fn test_part2_example2() {
        let map: Map = r#"
111111111111
999999999991
999999999991
999999999991
999999999991
            "#
        .into();
        let path = dbg!(map.dijkstra_ish(
            (0, 0),
            (map.width as isize - 1, map.height as isize - 1),
            Map::part2_edges
        ));
        assert_eq!(map.part2_path_cost(&path.unwrap()), 71);
    }

    #[test]
    fn test_part2_input_subgrid() {
        let map: Map = r#"
645655655636632423532352254234
353344463432453352632231125552
526342234456245465642135234421
346523323446362131122522334552
344544334245121341153154552255
545315342221553551233455324225
312521155444131243522241333531
153145243235542533153124432435
133522524342325151543155414215
            "#
        .into();
        assert_eq!(
            map.part2_path_cost(&[(0, 8), (9, 8), (9, 4), (19, 4), (19, 0), (29, 0),]),
            110
        );
    }
}