const XMAS: &str = "XMAS";
const MAS: &str = "MAS";

type Position = (usize, usize);
type Row = Vec<Position>;
type Rows = Vec<Row>;
type Move = (Position, (usize, isize));

struct Matrix {
    chars: Vec<Vec<char>>,
    limit: Position,
}

impl Matrix {
    fn new(text: &str) -> Self {
        let chars: Vec<Vec<char>> = text.lines().map(|row| row.chars().collect()).collect();
        let limit = (chars.len(), chars[0].len());
        Self { chars, limit }
    }

    fn check_iter<'a, I>(&self, it: I, needle: &str) -> bool
    where
        I: Iterator<Item = &'a Position>,
    {
        for (c, (i, j)) in needle.chars().zip(it) {
            if self.chars[*i][*j] != c {
                return false;
            }
        }

        true
    }

    fn check_row(&self, row: &Row, needle: &str) -> bool {
        self.check_iter(row.iter(), needle) || self.check_iter(row.iter().rev(), needle)
    }
}

struct Square<'a> {
    pos: Position,
    size: usize,
    matrix: &'a Matrix,
}

impl Square<'_> {
    fn rows<M>(&self, moves: M) -> Rows
    where
        M: Fn(&Square) -> Vec<Move>,
    {
        let mut rs: Rows = vec![];

        // Horizontal / vertical / diagonal
        for (p0, (di, dj)) in moves(self) {
            // Bounds check
            if p0.0 + di * self.size > self.matrix.limit.0
                || (dj > 0 && offplus(p0.1, dj * self.size as isize) > self.matrix.limit.1)
                || p0.1 >= self.matrix.limit.1
            {
                continue;
            }

            let mut p_acc = p0;
            let mut row: Row = vec![];
            for _idx in 0..self.size {
                row.push(p_acc);
                p_acc = (p_acc.0 + di, offplus(p_acc.1, dj));
            }
            rs.push(row);
        }

        rs
    }
}

fn offplus(x: usize, y: isize) -> usize {
    (x as isize + y) as usize
}

fn moves_xmas(Square { pos, size, .. }: &Square) -> Vec<Move> {
    let v = vec![
        (*pos, (0, 1)),                       // diag from pos, top-down and left-to-right
        (*pos, (1, 0)),                       // diag from pos, top-down and left-to-right
        (*pos, (1, 1)),                       // diag from pos, top-down and left-to-right
        ((pos.0, pos.1 + size - 1), (1, -1)), // the other diag, top-down and right-to-left
    ];
    v
}

fn p1(input: &str) -> String {
    let matrix: Matrix = Matrix::new(input);
    let mut result = 0;

    for i in 0..matrix.limit.0 {
        for j in 0..matrix.limit.1 {
            let sq = Square {
                pos: (i, j),
                size: XMAS.len(),
                matrix: &matrix,
            };
            for row in sq.rows(moves_xmas) {
                if matrix.check_row(&row, XMAS) {
                    result += 1;
                }
            }
        }
    }

    result.to_string()
}

fn moves_mas(Square { pos, size, .. }: &Square) -> Vec<Move> {
    let v = vec![
        (*pos, (1, 1)),                       // diag from pos, top-down and left-to-right
        ((pos.0, pos.1 + size - 1), (1, -1)), // the other diag, top-down and right-to-left
    ];
    v
}

fn p2(input: &str) -> String {
    let matrix: Matrix = Matrix::new(input);
    let mut result = 0;

    for i in 0..=matrix.limit.0 - MAS.len() {
        for j in 0..=matrix.limit.1 - MAS.len() {
            let sq = Square {
                pos: (i, j),
                size: MAS.len(),
                matrix: &matrix,
            };
            if sq.rows(moves_mas).iter().all(|r| matrix.check_row(r, MAS)) {
                result += 1;
            }
        }
    }

    result.to_string()
}

fn main() {
    aoc2024::run_day("4", p1, p2);
}