use nom::{
    branch::alt,
    character::complete::{char as nom_char, line_ending},
    combinator::{map, map_res, value},
    multi::{many1, separated_list1},
    sequence::tuple,
    IResult,
};
use std::{collections::BTreeSet, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_20.txt")?;
    let mut enhanceable_image = parse_enhanceable_image(&input).unwrap().1;
    for _ in 0..2 {
        enhanceable_image = enhanceable_image.enhance();
    }
    dbg!(enhanceable_image.image.count_light_spots().unwrap());
    for _ in 2..50 {
        enhanceable_image = enhanceable_image.enhance();
    }
    dbg!(enhanceable_image.image.count_light_spots().unwrap());

    Ok(())
}

#[derive(Debug)]
struct EnhanceableImage {
    enhancement_lights: [bool; 512],
    image: Image,
}

impl EnhanceableImage {
    fn enhance(&self) -> EnhanceableImage {
        let current_background_dark = matches!(self.image, Image::LightSpots(_));
        let next_background_dark =
            !self.enhancement_lights[if current_background_dark { 0 } else { 511 }];
        let mut spots = BTreeSet::new();

        let top_left = self.image.top_left(1);
        let bottom_right = self.image.bottom_right(1);
        for y in top_left.y..=bottom_right.y {
            for x in top_left.x..=bottom_right.x {
                let center = Point { x, y };
                let surrounds = center.surrounds();
                let number = self.image.to_number(surrounds);
                let center_is_light = self.enhancement_lights[number];
                if center_is_light == next_background_dark {
                    spots.insert(center);
                }
            }
        }

        EnhanceableImage {
            enhancement_lights: self.enhancement_lights.clone(),
            image: if next_background_dark {
                Image::LightSpots(spots)
            } else {
                Image::DarkSpots(spots)
            },
        }
    }
}

#[derive(Debug)]
enum Image {
    LightSpots(BTreeSet<Point>),
    DarkSpots(BTreeSet<Point>),
}

#[derive(Debug)]
enum ImageError {
    InfiniteLight,
}

impl Image {
    fn count_light_spots(&self) -> Result<usize, ImageError> {
        match self {
            Self::LightSpots(spots) => Ok(spots.len()),
            Self::DarkSpots(_) => Err(ImageError::InfiniteLight),
        }
    }

    fn top_left(&self, margin: i32) -> Point {
        let (Self::LightSpots(spots) | Self::DarkSpots(spots)) = self;
        let min_x = spots.iter().map(|p| p.x).min().unwrap_or(0);
        let min_y = spots.iter().map(|p| p.y).min().unwrap_or(0);
        Point {
            x: min_x - margin,
            y: min_y - margin,
        }
    }

    fn bottom_right(&self, margin: i32) -> Point {
        let (Self::LightSpots(spots) | Self::DarkSpots(spots)) = self;
        let max_x = spots.iter().map(|p| p.x).max().unwrap_or(0);
        let max_y = spots.iter().map(|p| p.y).max().unwrap_or(0);
        Point {
            x: max_x + margin,
            y: max_y + margin,
        }
    }

    fn to_number(&self, bit_locations: [Point; 9]) -> usize {
        let mut result = 0;
        for bit_location in bit_locations {
            result <<= 1;
            let next_is_1 = match self {
                Self::LightSpots(spots) => spots.contains(&bit_location),
                Self::DarkSpots(spots) => !spots.contains(&bit_location),
            };
            if next_is_1 {
                result += 1;
            }
        }
        result
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
struct Point {
    x: i32,
    y: i32,
}

impl Point {
    fn surrounds(&self) -> [Point; 9] {
        [
            Point {
                x: self.x - 1,
                y: self.y - 1,
            },
            Point {
                x: self.x,
                y: self.y - 1,
            },
            Point {
                x: self.x + 1,
                y: self.y - 1,
            },
            Point {
                x: self.x - 1,
                y: self.y,
            },
            Point {
                x: self.x,
                y: self.y,
            },
            Point {
                x: self.x + 1,
                y: self.y,
            },
            Point {
                x: self.x - 1,
                y: self.y + 1,
            },
            Point {
                x: self.x,
                y: self.y + 1,
            },
            Point {
                x: self.x + 1,
                y: self.y + 1,
            },
        ]
    }
}

fn parse_enhanceable_image(input: &str) -> IResult<&str, EnhanceableImage> {
    map(
        tuple((parse_enhancement_lights, many1(line_ending), parse_image)),
        |(enhancement_lights, _, image)| EnhanceableImage {
            enhancement_lights,
            image,
        },
    )(input)
}

fn parse_enhancement_lights(input: &str) -> IResult<&str, [bool; 512]> {
    map_res(many1(parse_pixel), |pixels| pixels.try_into())(input)
}

fn parse_image(input: &str) -> IResult<&str, Image> {
    map(separated_list1(line_ending, many1(parse_pixel)), |pixels| {
        let mut result = BTreeSet::new();
        for (y, row) in pixels.into_iter().enumerate() {
            for (x, light) in row.into_iter().enumerate() {
                if light {
                    result.insert(Point {
                        x: x as i32,
                        y: y as i32,
                    });
                }
            }
        }
        Image::LightSpots(result)
    })(input)
}

fn parse_pixel(input: &str) -> IResult<&str, bool> {
    alt((value(true, nom_char('#')), value(false, nom_char('.'))))(input)
}