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

type Input = &'static str;
type Output = usize;

fn parse_input(input: &'static str) -> Input {
    input
}

pub fn main() {
    let input = parse_input(INPUT);
    println!("Part 1: {}", part_1(&input));
    println!("Part 2: {}", part_2(&input));
}

fn part_1(input: &Input) -> Output {
    let horizontal = input.match_indices("XMAS").count();
    let horizontal_reverse = input.match_indices("SAMX").count();

    let grid = input
        .lines()
        .map(str::chars)
        .map(Iterator::collect::<Vec<_>>)
        .collect::<Vec<_>>();
    let grid_size = grid.len();
    let mut vertical = 0;
    let mut vertical_reverse = 0;
    for i in 0..grid.len() {
        let column = grid
            .iter()
            .map(|line| line.iter().skip(i).next().unwrap())
            .collect::<String>();
        vertical += column.match_indices("XMAS").count();
        vertical_reverse += column.match_indices("SAMX").count();
    }
    let mut diagonal_nw_se = 0;
    let mut diagonal_ne_sw = 0;
    for y in 0..=(grid_size - "XMAS".len()) {
        for x in 0..=(grid_size - "XMAS".len()) {
            let chars = [
                grid[y + 0][x + 0],
                grid[y + 1][x + 1],
                grid[y + 2][x + 2],
                grid[y + 3][x + 3],
            ];
            if chars == ['X', 'M', 'A', 'S'] || chars == ['S', 'A', 'M', 'X'] {
                diagonal_nw_se += 1;
            }
            let reverse_chars = [
                grid[y + 0][x + 3],
                grid[y + 1][x + 2],
                grid[y + 2][x + 1],
                grid[y + 3][x + 0],
            ];
            if reverse_chars == ['X', 'M', 'A', 'S'] || reverse_chars == ['S', 'A', 'M', 'X'] {
                diagonal_ne_sw += 1;
            }
        }
    }

    horizontal + horizontal_reverse + vertical + vertical_reverse + diagonal_nw_se + diagonal_ne_sw
}

fn part_2(input: &Input) -> Output {
    let grid = input
        .lines()
        .map(str::chars)
        .map(Iterator::collect::<Vec<_>>)
        .collect::<Vec<_>>();
    let grid_size = grid.len();

    let mut count = 0;
    for y in 1..grid_size - 1 {
        for x in 1..grid_size - 1 {
            if grid[y][x] != 'A' {
                continue;
            }
            let chars = [
                grid[y - 1][x - 1],
                grid[y + 1][x + 1],
                grid[y - 1][x + 1],
                grid[y + 1][x - 1],
            ];
            if matches!((chars[0], chars[1]), ('S', 'M') | ('M', 'S'))
                && matches!((chars[2], chars[3]), ('S', 'M') | ('M', 'S'))
            {
                count += 1;
            }
        }
    }

    count
}

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

    const SAMPLE_INPUT: &str = r#"MMMSXXMASM
MSAMXMSMSA
AMXSXMAAMM
MSAMASMSMX
XMASAMXAMM
XXAMMXXAMA
SMSMSASXSS
SAXAMASAAA
MAMMMXMMMM
MXMXAXMASX"#;

    #[test]
    fn test_part_1_with_sample_solution() {
        let input = parse_input(SAMPLE_INPUT);
        assert_eq!(super::part_1(&input), 18);
    }

    #[test]
    fn test_part_1_with_solution() {
        let input = parse_input(super::INPUT);
        assert_eq!(super::part_1(&input), 2562);
    }

    #[test]
    fn test_part_2_with_sample_solution() {
        let input = parse_input(SAMPLE_INPUT);
        assert_eq!(super::part_2(&input), 9);
    }

    #[test]
    fn test_part_2_with_solution() {
        let input = parse_input(super::INPUT);
        assert_eq!(super::part_2(&input), 1902);
    }
}