use super::TaskData;

/// Represents the four cardinal directions the guard can face.
#[derive(Debug, Clone, Copy)]
pub enum Direction {
    Up,
    Right,
    Down,
    Left,
}

/// Represents the type of each cell in the grid.
#[derive(Debug, Clone, Copy)]
pub enum CellType {
    Free,
    Obstacle,
}

/// Encapsulates the guard's position and facing direction.
#[derive(Debug, Clone, Copy)]
pub struct Guard {
    pub row: usize,
    pub col: usize,
    pub dir: Direction,
}

/// Solution module containing methods to solve both parts of the problem.
pub struct Solution;

impl Solution {
    /// Solves Part I: Calculates the number of distinct positions the guard visits
    /// before leaving the mapped area.
    pub fn solve_simple(data: &TaskData) -> i32 {
        let (moved_out, seen_dirs) =
            Self::simulate_guard(data.guard.clone(), &data.field, (data.rows, data.cols));
        if !moved_out {
            panic!("Guard is stuck in a loop.");
        }
        Self::count_seen_cells(&seen_dirs) as i32
    }

    /// Solves Part II: Determines the number of positions where placing a new obstacle
    /// would cause the guard to get stuck in a loop.
    pub fn solve_hard(data: &TaskData) -> i32 {
        let mut field = data.field.clone();
        let (_, seen_dirs) =
            Self::simulate_guard(data.guard.clone(), &field, (data.rows, data.cols));

        (0..data.rows)
            .flat_map(|r| (0..data.cols).map(move |c| (r, c)))
            .filter(|&(r, c)| (r, c) != (data.guard.row, data.guard.col) && seen_dirs[r][c] != 0)
            .filter_map(|(r, c)| {
                field[r][c] = CellType::Obstacle;
                let (moved_out, _) = 
                    Self::simulate_guard(data.guard.clone(), &field, (data.rows, data.cols));
                field[r][c] = CellType::Free;

                (!moved_out).then_some((r,c))
            }).count() as i32
    }

    /// Simulates the guard's movement until they either leave the grid or enter a loop.
    fn simulate_guard(
        mut guard: Guard,
        field: &[Vec<CellType>],
        limits: (usize, usize),
    ) -> (bool, Vec<Vec<i32>>) {
        let (rows, cols) = limits;
        let mut seen_dirs = vec![vec![0; cols]; rows];

        while seen_dirs[guard.row][guard.col] & guard.dir.bitmask() == 0 {
            // Mark the current cell as visited from the current direction.
            seen_dirs[guard.row][guard.col] |= guard.dir.bitmask();

            if !Self::step_guard(&mut guard, field, limits) {
                return (true, seen_dirs); // Guard has moved out of bounds.
            }
        }

        (false, seen_dirs) // Guard is stuck in a loop.
    }

    /// Attempts to move the guard forward. If an obstacle is ahead, the guard turns right.
    fn step_guard(guard: &mut Guard, field: &[Vec<CellType>], limits: (usize, usize)) -> bool {
        match guard.dir.next_cell(guard.row, guard.col, limits) {
            Some((nr, nc)) => {
                if matches!(field[nr][nc], CellType::Obstacle) {
                    guard.dir = guard.dir.turn_right();
                } else {
                    guard.row = nr;
                    guard.col = nc;
                }
                true
            }
            None => false, // Guard would move out of bounds
        }
    }

    /// Counts the number of distinct cells that the guard has visited.
    fn count_seen_cells(seen_dirs: &[Vec<i32>]) -> usize {
        seen_dirs.iter()
            .flat_map(|row| row.iter())
            .filter(|&&v| v != 0)
            .count()
    }
 }

impl Direction {
    /// Returns the direction obtained by turning right from the current direction.
    fn turn_right(self) -> Self {
        match self {
            Self::Up => Self::Right,
            Self::Right => Self::Down,
            Self::Down => Self::Left,
            Self::Left => Self::Up,
        }
    }

    /// Calculates the next cell position based on the current direction.
    fn next_cell(self, row: usize, col: usize, limits: (usize, usize)) -> Option<(usize, usize)> {
        match self {
            Self::Up if row > 0 => Some((row - 1, col)),
            Self::Right if col + 1 < limits.1 => Some((row, col + 1)),
            Self::Down if row + 1 < limits.0 => Some((row + 1, col)),
            Self::Left if col > 0 => Some((row, col - 1)),
            _ => None,
        }
    }

    /// Returns a bitmask corresponding to the current direction.
    fn bitmask(self) -> i32 {
        match self {
            Self::Up => 1,
            Self::Right => 2,
            Self::Down => 4,
            Self::Left => 8,
        }
    }
}