#[derive(Debug, PartialEq, Clone)]
enum Direction {
    North,
    South,
    East,
    West,
}

#[derive(Debug, PartialEq, Clone)]
struct RayState {
    x: usize,
    y: usize,
    direction: Direction,
}

impl RayState {
    fn advance(&self, rows: &usize, columns: &usize) -> Option<Self> {
        use Direction::*;
        match self.direction {
            North => {
                if self.y > 0 {
                    Some(Self {
                        x: self.x,
                        y: self.y - 1,
                        direction: North,
                })
                } else {
                    None
                }
            },
            South => {
                if self.y < columns - 1 {
                    Some(Self {
                        x: self.x,
                        y: self.y + 1,
                        direction: South,
                })
                } else {
                    None
                }
            },
            East => {
                if self.x < rows - 1 {
                    Some(Self {
                        x: self.x + 1,
                        y: self.y,
                        direction: East,
                })
                } else {
                    None
                }
            },
            West => {
                if self.x > 0 {
                    Some(Self {
                        x: self.x - 1,
                        y: self.y,
                        direction: West,
                })
                } else {
                    None
                }
            },
        }
    }
    fn handle_mirror(&self, mirror_char: u8) -> Vec<Self> {
        match mirror_char {
            b'|' => {
                match self.direction {
                    Direction::East | Direction::West => {
                        vec![
                            RayState{x: self.x, y: self.y, direction: Direction::North},
                            RayState{x: self.x, y: self.y, direction: Direction::South},
                        ]
                    },
                    _ => {vec![self.clone()]},
                }
            },
            b'-' => {
                match self.direction {
                    Direction::North | Direction::South => {
                        vec![
                            RayState{x: self.x, y: self.y, direction: Direction::East},
                            RayState{x: self.x, y: self.y, direction: Direction::West},
                        ]
                    },
                    _ => {vec![self.clone()]},
                }
            },
            b'/' => {
                match self.direction {
                    Direction::North => {
                        vec![RayState{x: self.x, y: self.y, direction: Direction::East}]
                    },
                    Direction::South => {
                        vec![RayState{x: self.x, y: self.y, direction: Direction::West}]
                    },
                    Direction::East => {
                        vec![RayState{x: self.x, y: self.y, direction: Direction::North}]
                    },
                    Direction::West => {
                        vec![RayState{x: self.x, y: self.y, direction: Direction::South}]
                    },
                }
            },
            b'\\' => {
                match self.direction {
                    Direction::North => {
                        vec![RayState{x: self.x, y: self.y, direction: Direction::West}]
                    },
                    Direction::South => {
                        vec![RayState{x: self.x, y: self.y, direction: Direction::East}]
                    },
                    Direction::East => {
                        vec![RayState{x: self.x, y: self.y, direction: Direction::South}]
                    },
                    Direction::West => {
                        vec![RayState{x: self.x, y: self.y, direction: Direction::North}]
                    },
                }
            },
            _ => {vec![self.clone()]},
        }
    }
}

fn get_energized(input: &str, start_rays: Vec<RayState>, rows: usize, columns: usize) -> usize {
    let mut history = Vec::<RayState>::new();
    let mut energized = std::iter::repeat(
            std::iter::repeat(false).take(columns).collect()
        )
        .take(rows)
        .collect::<Vec<Vec<bool>>>();
    start_rays.iter().for_each(|RayState {x, y, ..}| { energized[*y][*x] = true; });
    let _ = (0..).try_fold(start_rays, |rays, _| {
        let new_rays = rays.into_iter().filter_map(|ray| {
            ray.advance(&rows, &columns).and_then(|next_ray| {
                energized[next_ray.y][next_ray.x] = true;
                if history.contains(&next_ray) {
                    None
                } else {
                    history.push(next_ray.clone());
                    let mirror_bytes = input.lines().nth(next_ray.y).unwrap().as_bytes()[next_ray.x];
                    Some(next_ray.handle_mirror(mirror_bytes))
                }
            })
        }).flatten().collect::<Vec<RayState>>();
        if new_rays.is_empty() {
            Err("Empty")
        } else {
            Ok(new_rays)
        }
    });
    energized.into_iter().map(|line| {
        line.into_iter().map(|character| {
            if character {
                1
            } else {
                0
            }
        })
    }).flatten().sum::<usize>()
}

pub fn part1(input: &str) -> String {
    let rows = input.lines().count();
    let columns = input.chars().position(|c| c=='\n').unwrap();
    let first_char = input.lines().nth(0).unwrap().as_bytes()[0];
    let first = RayState{x: 0, y: 0, direction: Direction::East};
    let start_rays = first.handle_mirror(first_char);
    get_energized(input, start_rays, rows, columns).to_string()
}

pub fn part2(input: &str) -> String {
    let rows = input.lines().count();
    let columns = input.chars().position(|c| c=='\n').unwrap();
    let edges = 
        std::iter::zip(0..rows, std::iter::repeat(0)).into_iter().map(|(y,x)| RayState{x, y, direction: Direction::East})
        .chain(std::iter::zip(std::iter::repeat(0), 0..columns).into_iter().map(|(y,x)| RayState{x, y, direction: Direction::South}))
        .chain(std::iter::zip(std::iter::repeat(rows - 1), 0..columns).into_iter().map(|(y,x)| RayState{x, y, direction: Direction::North}))
        .chain(std::iter::zip(0..rows, std::iter::repeat(columns - 1)).into_iter().map(|(y,x)| RayState{x, y, direction: Direction::West}));
    edges.map(|edge| {
        let first_char = input.lines().nth(edge.y).unwrap().as_bytes()[edge.x];
        let start_rays = edge.handle_mirror(first_char);
        let temp = get_energized(input, start_rays, rows, columns);
        temp
    }).max().unwrap().to_string()
}

pub mod prelude {
    pub use super::part1;
    pub use super::part2;
}
