use anyhow::{bail, Error, Result};
use itertools::Itertools;
use std::collections::{HashMap, HashSet};

/// Represents the direction in which a particular Blizzard is blowing.
/// We have impls to allow converting to/from single-bit u8 values so
/// that a set of four directions can be represented as a u8.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum Direction {
    Down,
    Left,
    Up,
    Right,
}

impl TryFrom<u8> for Direction {
    type Error = Error;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            1 => Ok(Direction::Down),
            2 => Ok(Direction::Left),
            4 => Ok(Direction::Up),
            8 => Ok(Direction::Right),
            _ => bail!("No direction corresponding to {value}!"),
        }
    }
}

impl Direction {
    /// Convert a Direction to a single-bit u8 value
    const fn value(&self) -> u8 {
        match self {
            Direction::Down => 1,
            Direction::Left => 2,
            Direction::Up => 4,
            Direction::Right => 8,
        }
    }

    /// Get an array of all four Directions
    const fn all() -> [Direction; 4] {
        [
            Direction::Down,
            Direction::Left,
            Direction::Up,
            Direction::Right,
        ]
    }
}

/// Represents one or more blizzards occupying a single space in the Valley.
/// Because there's only one blizzard per space in the input, there can ever
/// only be one blizzard per direction in a single space at any time after
/// the blizzards begin moving in the Valley. So, we never need more than
/// the presence/absence of each of the four directions per Blizzard space.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
struct Blizzard(u8);

impl Blizzard {
    fn from(direction: Direction) -> Self {
        Blizzard(direction.value())
    }

    fn direction(&self) -> Result<Direction> {
        Direction::try_from(self.0)
    }

    /// Add a direction to the current Blizzard
    fn add(&mut self, direction: Direction) {
        self.0 |= direction.value();
    }

    /// Check if the Blizzard space has a Blizzard blowing in a
    /// particular direction.
    fn has(&self, direction: Direction) -> bool {
        self.0 & direction.value() > 0
    }
}

/// Represents a single space in the Valley. Can either be a space with Blizzards
/// blowing in one or more directions, an impassable Wall, or an Empty space where
/// the elven expedition can walk.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum Space {
    Blizzard(Blizzard),
    Wall,
    Empty,
}

/// Represents the entire Valley and all the Blizzards blowing through it at a
/// given point in time.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct Valley {
    rows: usize,
    cols: usize,
    spaces: Vec<Vec<Space>>,
}

/// Produce a new Valley from a 2D grid of Spaces.
impl From<Vec<Vec<Space>>> for Valley {
    fn from(spaces: Vec<Vec<Space>>) -> Self {
        let rows = spaces.len();
        let cols = spaces.first().map(|line| line.len()).unwrap_or_default();
        Valley { rows, cols, spaces }
    }
}

impl From<&str> for Valley {
    fn from(input: &str) -> Self {
        let mut spaces = Vec::new();
        for line in input.lines() {
            let row = line
                .chars()
                .map(|glyph| match glyph {
                    '>' => Space::Blizzard(Blizzard::from(Direction::Right)),
                    '<' => Space::Blizzard(Blizzard::from(Direction::Left)),
                    '^' => Space::Blizzard(Blizzard::from(Direction::Up)),
                    'v' => Space::Blizzard(Blizzard::from(Direction::Down)),
                    '.' => Space::Empty,
                    '#' => Space::Wall,
                    _ => unreachable!(),
                })
                .collect::<Vec<_>>();
            spaces.push(row);
        }

        Valley::from(spaces)
    }
}

impl Valley {
    /// Produce a clone of this Valley with its state advanced by one minute. All
    /// Blizzards move forward by one space in the direction they are facing,
    /// wrapping around the Valley when they encounter a Wall.
    fn advance(&self) -> Self {
        // Creates a new mutable clone of this Valley with only Empty Spaces
        let Valley { rows, cols, spaces } = self;
        let mut new_spaces = vec![vec![Space::Empty; *cols]; *rows];
        let mut new_state = Valley::from(new_spaces);

        // For each Space in the current Valley, update the appropriate space in the
        // new state. Move Blizzards and set Walls. Skip the Empties, since all the
        // Spaces in the new state are already Empty.
        for (row, col) in (0..*rows).cartesian_product(0..*cols) {
            match spaces[row][col] {
                Space::Blizzard(blizzard) => {
                    for direction in Direction::all() {
                        if !blizzard.has(direction) {
                            continue;
                        }
                        let (new_row, new_col) = self.next_position(row, col, direction);
                        new_state.add_blizzard(new_row, new_col, direction);
                    }
                }
                Space::Wall => new_state.spaces[row][col] = Space::Wall,
                Space::Empty => continue,
            }
        }

        // Return the new Valley
        new_state
    }

    /// Given a starting row and column and a Direction to move, return the row
    /// and column where you'd end up if you moved in that Direction.
    fn next_position(&self, row: usize, col: usize, direction: Direction) -> (usize, usize) {
        match direction {
            // Wrapping for increasing rows/columns
            Direction::Down => ((row % (self.rows - 2)) + 1, col),
            Direction::Right => (row, (col % (self.cols - 2)) + 1),

            // Wrapping for decreasing rows/columns
            Direction::Left => match col - 1 {
                0 => (row, self.cols - 2),
                _ => (row, col - 1),
            },
            Direction::Up => match row - 1 {
                0 => (self.rows - 2, col),
                _ => (row - 1, col),
            },
        }
    }

    /// Given a row and column and the Direction a blizzard is blowing, add that
    /// blizzard Direction to that Space. If the Space is empty, convert it to a
    /// Blizzard. If there's already a Blizzard there, just add the new Direction
    /// to the existing Blizzard. Attempts to add a Blizzard to a Wall will
    /// definitely fail.
    fn add_blizzard(&mut self, row: usize, col: usize, direction: Direction) {
        match &mut self.spaces[row][col] {
            Space::Blizzard(v) => v.add(direction),
            Space::Wall => panic!("Tried to add a blizzard to a wall!"),
            Space::Empty => self.spaces[row][col] = Space::Blizzard(Blizzard::from(direction)),
        }
    }
}

const INPUT: &str = include_str!("../../input/24/input.txt");

/// Parse the initial Valley state from the input, then advance the state
/// minute-by-minute until we reach a state we've seen before. That's right,
/// the state of the Valley cycles! This way, we can have in memory a map
/// of which Spaces are Empty and which ones have Blizzards for each point
/// in time. This means we don't need to re-calculate each state on the fly,
/// potentially re-creating the same state multiple times. Store the Valley
/// states in a Vector where the index indicates the minute at which that
/// state is valid.
fn read() -> Vec<Valley> {
    let mut valley = Valley::from(INPUT);
    let mut valley_states = Vec::new();
    let mut seen_states = HashSet::new();
    while !seen_states.contains(&valley) {
        seen_states.insert(valley.clone());
        valley_states.push(valley.clone());
        valley = valley.advance()
    }
    valley_states
}

use core::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};

/// Solve Day 24, Part 1
fn solve(input: &[Valley]) -> u32 {
    let first_state = input.get(0).expect("Valley should have an initial state!");

    // All examples and input have the start Space at index (0, 1) and the end
    // Space on the bottom row, next to the last column.
    let start_at = Expedition(0, 1);
    let end_at = Expedition(first_state.rows - 1, first_state.cols - 2);

    // Calculate the length of the shortest path through the Valley.
    if let Some(minutes) = start_at.shortest_path(end_at, 0, input) {
        return minutes as u32;
    }

    // Unless we can't find a path. Then, freak out! This doesn't happen,
    // though. Not anymore...
    panic!("Could not find a way through the valley. Died of frostbite!");
}

/// Represents the location of the elven Expedition through the Valley.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
struct Expedition(usize, usize);

impl Expedition {
    /// Given the current location of the elves and a state of the Valley,
    /// determine which next steps are possible, in order to avoid landing
    /// on a Space with an active Blizzard. There are a maximum of five
    /// possible moves, including waiting.
    fn possible_next_steps(&self, valley: &Valley) -> [Option<Expedition>; 5] {
        let Expedition(row, col) = *self;
        let mut possible_steps = [None; 5];

        // Attempt to move to the left
        if let Space::Empty = valley.spaces[row][col - 1] {
            possible_steps[0] = Some(Expedition(row, col - 1));
        }

        // Attempt to move to the right
        if let Space::Empty = valley.spaces[row][col + 1] {
            possible_steps[1] = Some(Expedition(row, col + 1));
        }

        // The upward move needs to account for the fact that the
        // final space is in the top row, which means that moving into
        // the top row is possible.
        if row > 0 {
            if let Space::Empty = valley.spaces[row - 1][col] {
                possible_steps[2] = Some(Expedition(row - 1, col));
            }
        }

        // The downard move needs to account for the beginning space
        // being on the last row, which means that moving into the last row
        // is possible.
        if row < (valley.spaces.len() - 1) {
            if let Space::Empty = valley.spaces[row + 1][col] {
                possible_steps[3] = Some(Expedition(row + 1, col));
            }
        }

        // Waiting is a necessary option if there's nothing in our current space
        if let Space::Empty = valley.spaces[row][col] {
            possible_steps[4] = Some(Expedition(row, col));
        }

        possible_steps
    }

    /// Find the shortest path from this Expedition's location to the `target`,
    /// assuming we start the journey at minute `start_time`. Pass in a reference
    /// to the time states of the Valley so we can know which Spaces are Empty
    /// for each minute. It's a Dijkstra's implementation.
    fn shortest_path(
        &self,
        target: Expedition,
        start_time: usize,
        valley_states: &[Valley],
    ) -> Option<usize> {
        // Sort locations in the Heap by the least number of minutes spent. Uniquely
        // identify states along the path by the location of the Expedition for a
        // given state of the Valley.
        let mut open = BinaryHeap::from([(Reverse(start_time), *self)]);
        let mut minutes = HashMap::from([((*self, start_time % valley_states.len()), start_time)]);

        // So long as there are states to explore, take the one with the fewest
        // number of minutes passed.
        while let Some((Reverse(minutes_passed), expedition)) = open.pop() {
            // If we've found the end, return the number of minutes passed
            if expedition == target {
                return Some(minutes_passed);
            }

            // Get the state of the Valley in the next minute to identify
            // which Spaces are available to be moved to.
            let state_idx = (minutes_passed + 1) % valley_states.len();
            let state = &valley_states[state_idx];

            // Check each next step to see if this is the fastest we've gotten
            // to that state. If so, keep it and add it to the heap. Otherwise,
            // keep moving on.
            for step in expedition.possible_next_steps(state).into_iter().flatten() {
                let next_minutes = minutes_passed + 1;
                let curr_minutes = *minutes.get(&(step, state_idx)).unwrap_or(&usize::MAX);
                if next_minutes >= curr_minutes {
                    continue;
                }
                open.push((Reverse(next_minutes), step));
                minutes.insert((step, state_idx), next_minutes);
            }
        }

        None // Something has gone terribly wrong...
    }
}

/// Solve Day 24, Part 2
fn solve(input: &[Valley]) -> u32 {
    let first_state = input.get(0).expect("Valley should have an initial state!");

    // All examples and input have the start Space at index (0, 1) and the end
    // Space on the bottom row, next to the last column.
    let start_at = Expedition(0, 1);
    let end_at = Expedition(first_state.rows - 1, first_state.cols - 2);

    // Start at zero minutes and move from the start to the end.
    let Some(minutes) = start_at.shortest_path(end_at, 0, input) else {
        panic!("Could not find a way through the valley. Died of frostbite!"); 
    };

    // Turn around and head back to the start.
    let Some(minutes) = end_at.shortest_path(start_at, minutes, input) else {
        panic!("Could not find a way back! Died in a blizzard!"); 
    };

    // Then turn around and head back to the end again.
    let Some(minutes) = start_at.shortest_path(end_at, minutes, input) else {
        panic!("Could not find a way back! Died of embarassment!"); 
    };

    // Return the total number of minutes it took to travel all that way.
    minutes as u32
}