const INPUT: &str = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/input.txt"));

fn main() {

    let mut data = Vec::new();
    let mut cols = None;

    for line in INPUT.lines() {
        let mut count = 0;
        for ch in line.bytes() {
            data.push(ch);
            count += 1;
        }
        if let Some(cols) = cols {
            assert_eq!(cols, count);
        } else {
            cols = Some(count);
        }
    }

    let data = Data { data, cols: cols.unwrap(), };

    let mut total = 0;
    for row in 0..data.rows() as isize {
        for col in 0..data.cols as isize {
            total += xmas(&data, row, col);
        }
    }

    println!("Part 1: {}", total);

    let mut total = 0;
    for row in 0..data.rows() as isize {
        for col in 0..data.cols as isize {
            total += u32::from(x_mas(&data, row, col));
        }
    }
    println!("Part 2: {}", total);
}

struct Data {
    data: Vec<u8>,
    cols: usize,
}

impl Data {
    fn get(&self, row: isize, col: isize) -> Option<u8> {
        if row >= 0 && col >= 0 {
            let row = row as usize;
            let col = col as usize;
            self.data
                // From row start until end of array
                .get(row * self.cols..)
                // Just the desired row
                .and_then(|x| x.get(..self.cols))
                // The desired column
                .and_then(|x| x.get(col))
                .copied()
        } else {
            None
        }
    }

    fn rows(&self) -> usize {
        assert_eq!(self.data.len() % self.cols, 0);
        self.data.len() / self.cols
    }
}

fn x_mas(
    data: &Data,
    row: isize,
    col: isize,
) -> bool {
    if data.get(row, col) == Some(b'A') {
        let br = data.get(row + 1, col + 1);
        let bl = data.get(row + 1, col - 1);
        let tl = data.get(row - 1, col - 1);
        let tr = data.get(row - 1, col + 1);

        if let Some((((bl, br), tl), tr)) =  bl.zip(br).zip(tl).zip(tr) {
               br == bl && tr == tl && (
                 br == b'S' && tl == b'M' || br == b'M' && tl == b'S'
               )
            || br == tr && bl == tl && (
                 br == b'S' && tl == b'M' || br == b'M' && tl == b'S'
               )
        } else {
            false
        }
    } else {
        false
    }
}

fn xmas(
    data: &Data,
    row: isize,
    col: isize,
) -> u32 {
    let mut total = 0;
    for direction in [
        ( 0,  1),
        ( 1,  0),
        ( 0, -1),
        (-1,  0),
        ( 1,  1),
        (-1, -1),
        (-1,  1),
        ( 1, -1),
    ] {
        total += u32::from(xmas_one_direction(data, row, col, direction));
    }
    total
}

fn xmas_one_direction(
    data: &Data,
    mut row: isize,
    mut col: isize,
    direction: (isize, isize)
) -> bool {
    for ch in [b'X', b'M', b'A', b'S'] {
        if data.get(row, col) != Some(ch) {
            return false;
        }

        row += direction.0;
        col += direction.1;
    }

    return true;
}