use std::collections::{hash_map::Entry, BinaryHeap, HashMap, HashSet, VecDeque};

use anyhow::Context;

use crate::{Problem, Solution};

pub struct Day17;

impl Problem for Day17 {
    const DAY: u8 = 17;

    const INPUT: &'static str = include_str!("../input/day_17.txt");
}

impl Solution for Day17 {
    type Answer1 = usize;

    type Answer2 = usize;

    fn part_1(input: &str) -> anyhow::Result<Self::Answer1> {
        let grid = parse_grid(input).context("Failed to parse grid")?;

        let path = astar(&grid, (0, 0), (grid.len() - 1, grid[0].len() - 1))
            .context("Failed to find path")?;

        let mut visited = HashSet::new();
        for (p, cost) in path {
            visited.insert(p);
            println!("{p:?}: {cost}");
        }

        for (row, v) in grid.iter().enumerate() {
            let mut output = String::new();
            for (col, _) in v.iter().enumerate() {
                if visited.contains(&(row, col)) {
                    output.push('#');
                } else {
                    output.push('.');
                }
            }
            println!("{output}")
        }

        todo!()
    }

    fn part_2(input: &str) -> anyhow::Result<Self::Answer2> {
        todo!()
    }
}

fn parse_grid(input: &str) -> Option<Grid> {
    input
        .lines()
        .map(|v| {
            v.chars()
                .map(|c| c.to_digit(10).map(|n| n as usize))
                .try_collect::<Vec<_>>()
        })
        .try_collect::<Grid>()
}

fn astar(grid: &Grid, start: Position, goal: Position) -> Option<Vec<(Position, usize)>> {
    let mut queue = BinaryHeap::from([Node::start(start)]);
    let mut cache = HashMap::from([(start, usize::MAX)]);

    while let Some(node) = queue.pop() {
        for next in node.successors(grid, goal) {
            if next.position == goal {
                return Some(rebuild_path(cache, next));
            }

            match cache.entry(next.position) {
                Entry::Vacant(e) => {
                    e.insert(next.cost);
                }
                Entry::Occupied(mut e) => {
                    if next.cost < *e.get() {
                        e.insert(next.cost);
                    } else {
                        continue;
                    }
                }
            }

            queue.push(next)
        }
    }

    None
}

fn rebuild_path(cache: HashMap<Position, usize>, node: Node) -> Vec<(Position, usize)> {
    let mut path = node
        .parents
        .into_iter()
        .map(|p| (p, *cache.get(&p).unwrap()))
        .collect::<Vec<_>>();
    path.push((node.position, node.cost));
    path
}

type Grid = Vec<Vec<usize>>;
type Position = (usize, usize);

#[derive(Debug, Default, Clone, PartialEq, Eq)]
struct Node {
    estimated: usize,
    cost: usize,
    position: Position,
    parents: Vec<Position>,
    trail: VecDeque<Edge>,
}

impl Node {
    fn start(position: Position) -> Self {
        Node {
            position,
            estimated: usize::MAX,
            ..Default::default()
        }
    }

    fn successors(&self, grid: &Grid, goal: Position) -> Vec<Self> {
        Edge::ALL
            .into_iter()
            .filter(|e| self.trail.len() < 3 || !self.trail.iter().all(|t| e == t))
            .inspect(|e| print!(" {e:?}"))
            .filter_map(|e| {
                let (row, col) = self.position;

                let position = match e {
                    Edge::North if row < grid.len() - 1 => (row.checked_add(1)?, col),
                    Edge::East if col < grid[row].len() - 1 => (row, col.checked_add(1)?),
                    Edge::South => (row.checked_sub(1)?, col),
                    Edge::West => (row, col.checked_sub(1)?),
                    _ => return None,
                };

                let mut next = self.clone();
                next.position = position;
                next.cost += grid[position.0][position.1];
                next.estimated = next.cost + next.dist(goal);
                next.parents.push(self.position);

                if next.trail.len() >= 3 {
                    next.trail.pop_front();
                }
                next.trail.push_back(e);

                Some(next)
            })
            .collect()
    }

    fn dist(&self, p2: Position) -> usize {
        self.position.0.abs_diff(p2.0) + self.position.1.abs_diff(p2.1)
    }
}

impl PartialOrd for Node {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for Node {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        match self.estimated.cmp(&other.estimated) {
            std::cmp::Ordering::Equal => self.cost.cmp(&other.cost),
            c => c,
        }
    }
}

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum Edge {
    North,
    #[default]
    East,
    South,
    West,
}

impl Edge {
    const ALL: [Edge; 4] = [Edge::North, Edge::East, Edge::South, Edge::West];
}

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

    const INPUT: &str = indoc::indoc! {r#"
        2413432311323
        3215453535623
        3255245654254
        3446585845452
        4546657867536
        1438598798454
        4457876987766
        3637877979653
        4654967986887
        4564679986453
        1224686865563
        2546548887735
        4322674655533
    "#};

    #[test]
    fn test_part_1() -> anyhow::Result<()> {
        Ok(assert_eq!(102, Day17::part_1(INPUT)?))
    }
}