use aoc2024::direction::Direction::{self, *};
use aoc2024::matrix;
use core::cmp::max;
use std::collections::{BTreeSet, HashMap, HashSet};
use std::fmt;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum Score {
    Sc(u32, Direction),
    Inf,
}
use Score::*;

impl fmt::Display for Score {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let repr = match self {
            Self::Inf => "∞".to_string(),
            Self::Sc(x, _dir) => x.to_string(),
        };
        write!(f, "{}", repr)
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum Dot {
    Start,
    End,
    Wall,
    Empty,
}
use Dot::*;

impl fmt::Display for Dot {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let ch = match self {
            Start => "S",
            End => "E",
            Wall => "#",
            Empty => ".",
        };
        write!(f, "{}", ch)
    }
}

impl Dot {
    fn parse(c: char) -> Dot {
        match c {
            'S' => Start,
            'E' => End,
            '#' => Wall,
            '.' => Empty,
            _ => panic!("Invalid input"),
        }
    }
}

#[derive(Clone)]
struct M {
    matrix: matrix::Matrix<Dot>,
    unvisited: BTreeSet<(Score, matrix::Pos)>,
}

impl fmt::Display for M {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.matrix)
    }
}

impl M {
    fn new(text: &str) -> Self {
        let matrix = matrix::Matrix::new(text, Dot::parse);
        let mut unvisited = BTreeSet::new();
        for i in 0..matrix.limit.0 {
            for j in 0..matrix.limit.1 {
                let dot = matrix.get((i, j));
                let score = match dot {
                    Start => Some(Sc(0, Right)),
                    End => Some(Inf),
                    Wall => None,
                    Empty => Some(Inf),
                };
                if let Some(score) = score {
                    unvisited.insert((score, (i, j)));
                };
            }
        }
        M { matrix, unvisited }
    }

    fn adjacents(&self, pos: matrix::Pos, dir: Direction) -> Vec<(Direction, matrix::Pos)> {
        dir.around()
            .filter_map(|d| self.matrix.pos_move(pos, d.to_offset()).map(|p| (d, p)))
            .collect()
    }

    fn print_with_dots(&self, positions: &HashSet<matrix::Pos>, dot: char) {
        for i in 0..self.matrix.limit.0 {
            for j in 0..self.matrix.limit.1 {
                let p = (i, j);
                if positions.contains(&p) {
                    print!("{}", dot);
                } else {
                    print!("{}", self.matrix.get(p));
                }
            }
            println!();
        }
    }

    fn backtrack(
        &self,
        endpos: matrix::Pos,
        mut prevs: HashMap<matrix::Pos, HashSet<matrix::Pos>>,
    ) -> u32 {
        let mut visited: HashSet<matrix::Pos> = HashSet::new();
        visited.insert(endpos);
        let mut unvisited = vec![endpos];
        while let Some(node) = unvisited.pop() {
            if let Some(ps) = prevs.remove(&node) {
                for p in ps {
                    if visited.insert(p) {
                        unvisited.push(p);
                    }
                }
                // println!("Adding {} because of {node:?}'s prevs:", ps.len());
                // println!("  {:?}", ps);
                // unvisited.append(&mut ps.into_iter().collect());
            }
        }
        self.print_with_dots(&visited, 'O');
        // println!("Visited: {visited:?}");
        visited.len() as u32
    }

    fn dijkstra(&mut self) -> Option<(u32, u32)> {
        let mut visited: HashSet<matrix::Pos> = HashSet::new();
        let mut prevs: HashMap<matrix::Pos, HashSet<matrix::Pos>> = HashMap::new();
        let mut result = (0, (0, 0));
        while let Some((score, pos)) = self.unvisited.pop_first() {
            self.print_with_dots(&visited, 'O');
            self.print_with_dots(
                &self
                    .unvisited
                    .iter()
                    .filter_map(|(sc, p)| if let Sc(_, _) = sc { Some(*p) } else { None })
                    .collect(),
                '?',
            );
            let Sc(scoreint, dir) = score else {
                break;
            };
            let dot = self.matrix.get(pos);
            if dot == &End {
                // self.print_with_dots(&visited, 'O');
                result = (scoreint, pos);
            };
            // println!("Visiting ({pos:?}) with score = {score:?}");
            let adjs: Vec<(Score, matrix::Pos)> = self
                .adjacents(pos, dir)
                .into_iter()
                .map(|(d, p)| (d, p, self.matrix.get(p)))
                .filter(|&(_d, _p, dot)| dot != &Wall)
                .map(|(d, p, _dot)| (set_dot_score(scoreint, &score, dir, d), p))
                .collect();
            // println!("Adjacents when facing {:?}: {adjs:?}", dir);
            for adj in adjs {
                if visited.insert(adj.1) {
                    self.unvisited.insert((adj.0, adj.1));
                }
                prevs.entry(adj.1).or_default().insert(pos);
            }
            // println!("{}", self);
        }

        // }
        // unvisited.push((0, self.start));
        Some((result.0, self.backtrack(result.1, prevs)))
    }
}

fn set_dot_score(currscore: u32, dotscore: &Score, dir: Direction, newdir: Direction) -> Score {
    let new_score = Sc(if dir == newdir { 1 } else { 1001 } + currscore, newdir);
    max(*dotscore, new_score)
}

fn p1(input: &str) -> String {
    let mut m = M::new(input);
    println!("{m}");

    let result = m.dijkstra().unwrap();
    result.0.to_string()
}

fn p2(input: &str) -> String {
    let mut m = M::new(input);
    println!("{m}");

    let result = m.dijkstra().unwrap();
    result.1.to_string()
}

fn main() {
    aoc2024::run_day("16", p1, p2);
}