#![feature(let_chains)]

static MAS: [char; 3] = ['M', 'A', 'S'];
static DIRS: [Dir; 8] = [
    Dir::N,
    Dir::S,
    Dir::E,
    Dir::W,
    Dir::NE,
    Dir::NW,
    Dir::SE,
    Dir::SW,
];

fn main() {
    let input = std::fs::read_to_string("input").unwrap();
    println!("{}", pt1(&input));
    println!("{}", pt2(&input));
}

fn pt1(input: &str) -> u32 {
    let mut total = 0;
    let board = Board::new(input);

    for (r, row) in board.rows.iter().enumerate() {
        for (c, letter) in row.iter().enumerate() {
            if letter == &'X' {
                for &dir in DIRS.iter() {
                    if find_xmas(&board, (r, c), dir) {
                        total += 1;
                    }
                }
            }
        }
    }

    total
}

fn find_xmas(board: &Board, loc: (usize, usize), dir: Dir) -> bool {
    let (mut r, mut c) = loc;
    for l in MAS {
        if let Some(&nl) = board.next((r, c), dir) {
            if nl != l {
                return false;
            }
            (r, c) = dir.next((r, c));
        } else {
            return false;
        }
    }
    true
}

fn pt2(input: &str) -> u32 {
    let mut total = 0;
    let board = Board::new(input);

    for (r, row) in board.rows.iter().enumerate() {
        for (c, letter) in row.iter().enumerate() {
            if letter == &'A' {
                if find_x_mas(&board, (r, c)) {
                    total += 1;
                }
            }
        }
    }

    total
}

fn find_x_mas(board: &Board, loc: (usize, usize)) -> bool {
    let one = if let Some(&nw) = (board.next(loc, Dir::NW))
        && let Some(&se) = board.next(loc, Dir::SE)
    {
        if (nw == 'M' && se == 'S') || (nw == 'S' && se == 'M') {
            true
        } else {
            false
        }
    } else {
        false
    };
    let two = if let Some(&ne) = (board.next(loc, Dir::NE))
        && let Some(&sw) = board.next(loc, Dir::SW)
    {
        if (ne == 'M' && sw == 'S') || (ne == 'S' && sw == 'M') {
            true
        } else {
            false
        }
    } else {
        false
    };

    one && two
}

struct Board {
    rows: Vec<Vec<char>>,
}

impl Board {
    fn new(input: &str) -> Self {
        let rows = input.split('\n').map(|l| l.chars().collect()).collect();
        Board { rows }
    }

    fn at(&self, loc: (usize, usize)) -> Option<&char> {
        self.rows.get(loc.0).and_then(|r| r.get(loc.1))
    }

    fn next(&self, loc: (usize, usize), dir: Dir) -> Option<&char> {
        let (nr, nc) = dir.next(loc);
        self.at((nr, nc))
    }
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
enum Dir {
    N,
    S,
    E,
    W,
    NE,
    NW,
    SE,
    SW,
}

impl Dir {
    fn next(&self, loc: (usize, usize)) -> (usize, usize) {
        let (r, c) = loc;
        let (nr, nc) = match self {
            Dir::N => (r.wrapping_sub(1), c),
            Dir::S => (r + 1, c),
            Dir::E => (r, c + 1),
            Dir::W => (r, c.wrapping_sub(1)),
            Dir::NE => (r.wrapping_sub(1), c + 1),
            Dir::NW => (r.wrapping_sub(1), c.wrapping_sub(1)),
            Dir::SE => (r + 1, c + 1),
            Dir::SW => (r + 1, c.wrapping_sub(1)),
        };

        (nr, nc)
    }
}

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

    static INPUT: &str = "MMMSXXMASM
MSAMXMSMSA
AMXSXMAAMM
MSAMASMSMX
XMASAMXAMM
XXAMMXXAMA
SMSMSASXSS
SAXAMASAAA
MAMMMXMMMM
MXMXAXMASX";

    #[test]
    fn p1() {
        assert_eq!(18, pt1(INPUT));
    }

    #[test]
    fn p2() {
        assert_eq!(9, pt2(INPUT));
    }
}