use itertools::Itertools;

use crate::common::{Answer, Solution};

pub struct Day04;

impl Solution for Day04 {
    fn name(&self) -> &'static str {
        "Ceres Search"
    }

    fn part_a(&self, input: &str) -> Answer {
        let horizontal = input.lines().map(xmas_count).sum::<usize>();
        let vertical = rotate_90(input).map(xmas_count).sum::<usize>();
        let diag_45 = rotate_45(input).map(xmas_count).sum::<usize>();
        let diag_135 = rotate_45(&rotate_90(input).join("\n"))
            .map(xmas_count)
            .sum::<usize>();

        (horizontal + vertical + diag_45 + diag_135).into()
    }

    fn part_b(&self, input: &str) -> Answer {
        let input_90 = rotate_90(input).join("\n");
        let input_180 = rotate_90(&input_90).join("\n");
        let input_270 = rotate_90(&input_180).join("\n");

        let sum = cross_mas_count(input)
            + cross_mas_count(&input_90)
            + cross_mas_count(&input_180)
            + cross_mas_count(&input_270);

        sum.into()
    }
}

fn rotate_90(input: &str) -> impl Iterator<Item = String> + use<'_> {
    let col_count = input.lines().next().unwrap().chars().count();

    (0..col_count)
        .rev() // this is important otherwise 135 rotation doesn't work
        .map(|i| input.lines().map(|l| l.chars().nth(i).unwrap()).join(""))
}

fn rotate_45(input: &str) -> impl Iterator<Item = String> + use<'_> {
    let line_count = input.lines().count();
    let col_count = input.lines().next().unwrap().chars().count();

    let cols = (0..col_count).map(|i| traverse_diagonal(input, (0, i)));
    let lines = (1..line_count).map(|i| traverse_diagonal(input, (i, 0)));

    cols.chain(lines)
}

fn traverse_diagonal(input: &str, start: (usize, usize)) -> String {
    (start.0..)
        .zip(start.1..)
        .map_while(|(x, y)| input.lines().nth(x).and_then(|l| l.chars().nth(y)))
        .collect()
}

fn xmas_count(line: impl AsRef<str>) -> usize {
    line.as_ref().matches("XMAS").count() + line.as_ref().matches("SAMX").count()
}

fn cross_mas_count(input: &str) -> usize {
    let mut sum = 0;

    for (x, line) in input.lines().enumerate() {
        for (y, _) in line.match_indices('A') {
            let prev_line = x.checked_sub(1).and_then(|n| input.lines().nth(n));
            let next_line = input.lines().nth(x + 1);

            if let (Some(prev_line), Some(next_line)) = (prev_line, next_line) {
                let prev_y = y.checked_sub(1);
                let next_y = y + 1;

                let prev_prev = prev_y.and_then(|n| prev_line.chars().nth(n));
                let prev_next = prev_line.chars().nth(next_y);

                let next_prev = prev_y.and_then(|n| next_line.chars().nth(n));
                let next_next = next_line.chars().nth(next_y);

                if let (Some('M'), Some('M'), Some('S'), Some('S')) =
                    (prev_prev, prev_next, next_prev, next_next)
                {
                    sum += 1;
                }
            }
        }
    }

    sum
}

#[cfg(test)]
mod test {
    use super::Day04;
    use crate::common::Solution;

    use indoc::indoc;

    const INPUT: &str = indoc! {"
        MMMSXXMASM
        MSAMXMSMSA
        AMXSXMAAMM
        MSAMASMSMX
        XMASAMXAMM
        XXAMMXXAMA
        SMSMSASXSS
        SAXAMASAAA
        MAMMMXMMMM
        MXMXAXMASX
    "};

    #[test]
    fn part_a() {
        assert_eq!(Day04.part_a(INPUT), 18.into());
    }

    #[test]
    fn part_b() {
        assert_eq!(Day04.part_b(INPUT), 9.into());
    }
}