use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::line_ending,
    combinator::{map, value},
    multi::{many1, separated_list1},
    IResult,
};
use std::{collections::BTreeSet, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_16.txt")?;
    let device = LightDevice::parser(&input).unwrap().1;
    {
        let mut part_1_device = device.clone();
        part_1_device.energize(Point { x: 0, y: 0 }, Direction::East);
        dbg!(&part_1_device.count_energized_blocks());
    }

    dbg!(device.find_max_energization());

    Ok(())
}

#[derive(Debug, Clone)]
struct LightDevice {
    mirrors: Vec<Vec<LightBlock>>,
    light: Vec<Vec<BTreeSet<Direction>>>,
    bounds: Point,
}

#[derive(Debug, Clone, Copy)]
enum LightBlock {
    MirrorForwardLeaning,   // /
    MirrorBackwardsLeaning, // \
    Empty,                  // .
    HorizontalSplitter,     // -
    VerticalSplitter,       // |
}

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

#[derive(Debug, Clone)]
struct Point {
    x: usize,
    y: usize,
}

impl LightDevice {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_list1(line_ending, many1(LightBlock::parser)),
            |mirrors| LightDevice {
                bounds: Point {
                    x: mirrors[0].len(),
                    y: mirrors.len(),
                },
                light: mirrors
                    .iter()
                    .map(|mirror_row| vec![BTreeSet::new(); mirror_row.len()])
                    .collect(),
                mirrors,
            },
        )(input)
    }

    fn energize(&mut self, start_point: Point, start_direction: Direction) {
        let mut frontier = vec![(start_point.clone(), start_direction.clone())];
        self.light[start_point.y][start_point.x].insert(start_direction);

        while let Some((front_light_p, front_light_dir)) = frontier.pop() {
            let mirror = self.mirrors[front_light_p.y][front_light_p.x];
            let new_dirs: Vec<Direction> = match (mirror, front_light_dir) {
                (LightBlock::MirrorForwardLeaning, Direction::North) => vec![Direction::East],
                (LightBlock::MirrorForwardLeaning, Direction::South) => vec![Direction::West],
                (LightBlock::MirrorForwardLeaning, Direction::East) => vec![Direction::North],
                (LightBlock::MirrorForwardLeaning, Direction::West) => vec![Direction::South],
                (LightBlock::MirrorBackwardsLeaning, Direction::North) => vec![Direction::West],
                (LightBlock::MirrorBackwardsLeaning, Direction::South) => vec![Direction::East],
                (LightBlock::MirrorBackwardsLeaning, Direction::East) => vec![Direction::South],
                (LightBlock::MirrorBackwardsLeaning, Direction::West) => vec![Direction::North],
                (LightBlock::HorizontalSplitter, Direction::North)
                | (LightBlock::HorizontalSplitter, Direction::South) => {
                    vec![Direction::East, Direction::West]
                }
                (LightBlock::VerticalSplitter, Direction::East)
                | (LightBlock::VerticalSplitter, Direction::West) => {
                    vec![Direction::North, Direction::South]
                }
                (LightBlock::Empty, dir)
                | (LightBlock::HorizontalSplitter, dir)
                | (LightBlock::VerticalSplitter, dir) => vec![dir],
            };

            let new_lights: Vec<(Point, Direction)> = new_dirs
                .into_iter()
                .filter_map(|dir| front_light_p.go(dir, &self.bounds).map(|p| (p, dir)))
                .collect();

            for (new_light_p, new_light_dir) in new_lights {
                if !self.light[new_light_p.y][new_light_p.x].contains(&new_light_dir) {
                    self.light[new_light_p.y][new_light_p.x].insert(new_light_dir);
                    frontier.push((new_light_p, new_light_dir));
                }
            }
        }
    }

    fn count_energized_blocks(&self) -> usize {
        self.light
            .iter()
            .flat_map(|row| row.iter())
            .filter(|set| set.len() > 0)
            .count()
    }

    fn find_max_energization(&self) -> usize {
        (0..self.bounds.x)
            .flat_map(|x| {
                vec![
                    (Point { x, y: 0 }, Direction::South),
                    (
                        Point {
                            x,
                            y: self.bounds.y - 1,
                        },
                        Direction::North,
                    ),
                ]
            })
            .chain((0..self.bounds.y).flat_map(|y| {
                vec![
                    (Point { x: 0, y }, Direction::East),
                    (
                        Point {
                            x: self.bounds.x - 1,
                            y,
                        },
                        Direction::West,
                    ),
                ]
            }))
            .map(|(start_p, start_d)| {
                let mut energizer = self.clone();
                energizer.energize(start_p, start_d);
                energizer.count_energized_blocks()
            })
            .max()
            .unwrap()
    }
}

impl LightBlock {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            value(LightBlock::MirrorForwardLeaning, tag("/")),
            value(LightBlock::MirrorBackwardsLeaning, tag("\\")),
            value(LightBlock::Empty, tag(".")),
            value(LightBlock::HorizontalSplitter, tag("-")),
            value(LightBlock::VerticalSplitter, tag("|")),
        ))(input)
    }
}

impl Point {
    fn up(&self) -> Point {
        Point {
            x: self.x,
            y: self.y - 1,
        }
    }

    fn down(&self) -> Point {
        Point {
            x: self.x,
            y: self.y + 1,
        }
    }

    fn left(&self) -> Point {
        Point {
            x: self.x - 1,
            y: self.y,
        }
    }

    fn right(&self) -> Point {
        Point {
            x: self.x + 1,
            y: self.y,
        }
    }

    fn go(&self, dir: Direction, bounds: &Point) -> Option<Point> {
        match dir {
            Direction::North if self.y > 0 => Some(self.up()),
            Direction::South if self.y < bounds.y - 1 => Some(self.down()),
            Direction::West if self.x > 0 => Some(self.left()),
            Direction::East if self.x < bounds.x - 1 => Some(self.right()),
            _ => None,
        }
    }
}