use std::fs;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("input")?;
    let layout = parse_layout(&input)?;

    // Part 1
    let result = count_eventually_occupied_seats(&layout);
    println!("{}", result);

    Ok(())
}

#[derive(Copy, Clone, Debug, PartialEq)]
enum Seating {
    Occupied,
    Empty,
    Floor,
}

impl Seating {
    fn new(symbol: char) -> Result<Self, String> {
        match symbol {
            '#' => Ok(Seating::Occupied),
            'L' => Ok(Seating::Empty),
            '.' => Ok(Seating::Floor),
            rubbish => Err(format!("Found some rubbish on the floor, area must be clean to conserve human behavior: {}", rubbish))
        }
    }
}

fn parse_layout(input: &str) -> Result<Vec<Vec<Seating>>, String> {
    input
        .lines()
        .map(|line| line.chars().map(Seating::new).collect())
        .collect()
}

fn print_layout(layout: &Vec<Vec<Seating>>) {
    let output: String = layout
        .iter()
        .map(|row| {
            let mut row: String = row
                .iter()
                .map(|c| match c {
                    Seating::Occupied => '#',
                    Seating::Empty => 'L',
                    Seating::Floor => '.',
                })
                .collect();
            row.push('\n');
            row
        })
        .collect();

    println!("{}", output);
}

fn get_adjacent(layout: &Vec<Vec<Seating>>, i: usize, j: usize) -> Vec<Option<&Seating>> {
    let adjacent: [[usize; 2]; 8] = [
        [0, 0],
        [0, 1],
        [0, 2],
        [1, 0],
        [1, 2],
        [2, 0],
        [2, 1],
        [2, 2],
    ];

    adjacent
        .iter()
        .map(|[nudge_row, nudge_col]| {
            if let Some(row) = (i + nudge_row).checked_sub(1) {
                if let Some(row) = layout.get(row) {
                    if let Some(col) = (j + nudge_col).checked_sub(1) {
                        row.get(col)
                    } else {
                        None
                    }
                } else {
                    None
                }
            } else {
                None
            }
        })
        .collect()
}

fn simulate_people(layout: &Vec<Vec<Seating>>) -> (Vec<Vec<Seating>>, usize) {
    let mut changes = 0;

    let layout = layout
        .iter()
        .enumerate()
        .map(|(i, row)| {
            row.iter()
                .enumerate()
                .map(|(j, seat)| match seat {
                    Seating::Occupied => {
                        if get_adjacent(layout, i, j)
                            .iter()
                            .filter(|seat| **seat == Some(&Seating::Occupied))
                            .count()
                            >= 4
                        {
                            changes += 1;
                            Seating::Empty
                        } else {
                            Seating::Occupied
                        }
                    }
                    Seating::Empty => {
                        if get_adjacent(layout, i, j)
                            .iter()
                            .any(|seat| *seat == Some(&Seating::Occupied))
                        {
                            Seating::Empty
                        } else {
                            changes += 1;
                            Seating::Occupied
                        }
                    }
                    Seating::Floor => Seating::Floor,
                })
                .collect()
        })
        .collect();

    (layout, changes)
}

fn count_eventually_occupied_seats(layout: &Vec<Vec<Seating>>) -> usize {
    let mut loop_layout = layout.clone();

    loop {
        let (res, changes) = simulate_people(&loop_layout);
        loop_layout = res;

        if changes == 0 {
            break;
        }
    }

    loop_layout
        .iter()
        .flatten()
        .filter(|seat| **seat == Seating::Occupied)
        .count()
}

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

    #[test]
    fn test_simple() -> Result<(), Box<dyn std::error::Error>> {
        let input = indoc!(
            "L.LL.LL.LL
             LLLLLLL.LL
             L.L.L..L..
             LLLL.LL.LL
             L.LL.LL.LL
             L.LLLLL.LL
             ..L.L.....
             LLLLLLLLLL
             L.LLLLLL.L
             L.LLLLL.LL
            "
        );

        let layout = parse_layout(&input)?;
        let result = count_eventually_occupied_seats(&layout);

        assert_eq!(result, 37);
        Ok(())
    }
}