use aoc2024::matrix;
use std::collections::{HashMap, HashSet};
use std::fmt;

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

impl Direction {
    fn rotate(self) -> Self {
        match self {
            Up => Right,
            Right => Down,
            Down => Left,
            Left => Up,
        }
    }
}

#[derive(Clone)]
enum Dot {
    Empty,
    Obstacle,
    Visited,
    Guard(Direction),
}
use Dot::*;

impl fmt::Display for Dot {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let ch = match self {
            Empty => ".",
            Obstacle => "#",
            Visited => "X",
            Guard(Up) => "^",
            Guard(Right) => ">",
            Guard(Down) => "v",
            Guard(Left) => "<",
        };
        write!(f, "{}", ch)
    }
}

impl Dot {
    fn parse(c: char) -> Dot {
        match c {
            '.' => Empty,
            '#' => Obstacle,
            'X' => Visited,
            '^' => Guard(Up),
            '>' => Guard(Right),
            'v' => Guard(Down),
            '<' => Guard(Left),
            _ => panic!("Invalid input"),
        }
    }
}

#[derive(Clone)]
struct M {
    matrix: matrix::Matrix<Dot>,
    guard: matrix::Pos,
    visited: HashMap<matrix::Pos, HashSet<Direction>>,
}

impl M {
    fn new(text: &str) -> Self {
        let matrix = matrix::Matrix::new(text, Dot::parse);
        let mut guard = (0, 0);
        let mut visited = HashMap::new();
        for i in 0..matrix.limit.0 {
            for j in 0..matrix.limit.1 {
                if let Guard(dir) = matrix.get((i, j)) {
                    let mut dirhs = HashSet::new();
                    dirhs.insert(*dir);
                    visited.insert(guard, dirhs);
                    guard = (i, j);
                }
            }
        }
        M {
            matrix,
            guard,
            visited,
        }
    }
}

#[derive(Debug)]
enum SimResult {
    OutOfBounds,
    Loop,
}

impl M {
    fn step(&mut self) -> Option<SimResult> {
        let M {
            matrix,
            guard,
            visited,
        } = self;

        // Compute guard's next position
        let Guard(direction) = matrix.get(*guard) else {
            panic!("Impossible: {self}")
        };
        let newpos: (usize, usize) = match direction {
            Up if guard.0 > 0 => (guard.0 - 1, guard.1),
            Right if guard.1 < matrix.limit.1 - 1 => (guard.0, guard.1 + 1),
            Down if guard.0 < matrix.limit.0 - 1 => (guard.0 + 1, guard.1),
            Left if guard.1 > 0 => (guard.0, guard.1 - 1),
            _ => return Some(SimResult::OutOfBounds),
        };
        match (matrix.get(newpos), visited.get(guard)) {
            (Visited, Some(vs)) if vs.contains(direction) => return Some(SimResult::Loop),
            (Empty | Visited, _) => {
                let d = *direction;
                matrix.set(newpos, Guard(d));
                matrix.set(*guard, Visited);
                let hs = visited.entry(*guard).or_insert(HashSet::new());
                hs.insert(d);
                *guard = newpos;
            }
            (Obstacle, _) => {
                matrix.set(*guard, Guard(direction.rotate()));
            }
            _ => (),
        };
        // println!("{self}");
        None
    }

    fn step_forever(&mut self) -> SimResult {
        loop {
            if let Some(res) = self.step() { return res }
        }
    }

    fn count_visited(&self) -> u32 {
        let mut result: u32 = 0;
        let M { matrix, .. } = self;
        for i in 0..matrix.limit.0 {
            for j in 0..matrix.limit.1 {
                match matrix.get((i, j)) {
                    Visited | Guard(_) => result += 1,
                    _ => (),
                }
            }
        }
        result
    }
}

impl fmt::Display for M {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // write!(f, "{}[2J", 27 as char)?;
        write!(f, "{}", self.matrix)?;
        write!(f, "[Guard position: {:?}]", self.guard)
    }
}

fn p1(input: &str) -> String {
    let mut m: M = M::new(input);
    println!("{m}");
    println!("Result: {:?}", m.step_forever());
    println!("{m}");
    let result = m.count_visited();
    result.to_string()
}

fn count_loops(m: M) -> u32 {
    let mut result = 0;
    // let mut simulations = 0;
    for i in 0..m.matrix.limit.0 {
        for j in 0..m.matrix.limit.1 {
            if let Empty = m.matrix.get((i, j)) {
                // simulations += 1;
                let mut m1 = m.clone();
                m1.matrix.set((i, j), Obstacle);
                // println!("Simulation {simulations} @ ({i},{j}) ({result})...");
                // println!("{m1}");
                if let SimResult::Loop = m1.step_forever() {
                    result += 1;
                }
            }
        }
    }
    result
}

fn p2(input: &str) -> String {
    let m: M = M::new(input);
    let result = count_loops(m);
    result.to_string()
}

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