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

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_11.txt")?;
    let galaxy_map = GalaxyMap::parser(&input).unwrap().1;

    let part_1_expanded_galaxy_map = galaxy_map.expand(2);
    dbg!(&part_1_expanded_galaxy_map.distance_sum_between_galaxies());

    let part_2_expanded_galaxy_map = galaxy_map.expand(1_000_000);
    dbg!(&part_2_expanded_galaxy_map.distance_sum_between_galaxies());

    Ok(())
}

#[derive(Debug)]
struct GalaxyMap(Vec<Point>);

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

#[derive(Debug, Clone)]
enum Token {
    Space,
    Galaxy,
}

impl GalaxyMap {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_list1(
                line_ending,
                many1(alt((
                    value(Token::Space, char('.')),
                    value(Token::Galaxy, char('#')),
                ))),
            ),
            |rows| {
                GalaxyMap(
                    rows.into_iter()
                        .enumerate()
                        .flat_map(move |(y, row)| {
                            row.into_iter()
                                .enumerate()
                                .filter_map(move |(x, token)| match token {
                                    Token::Space => None,
                                    Token::Galaxy => Some(Point { x, y }),
                                })
                        })
                        .collect(),
                )
            },
        )(input)
    }

    fn expand(&self, expansion_factor: usize) -> GalaxyMap {
        let max_y = self.0.iter().map(|p| p.y).max().unwrap();
        let empty_rows: Vec<usize> = (0..max_y)
            .filter(|y| !self.0.iter().any(|p| p.y == *y))
            .collect();

        let max_x = self.0.iter().map(|p| p.x).max().unwrap();
        let empty_cols: Vec<usize> = (0..max_x)
            .filter(|x| !self.0.iter().any(|p| p.x == *x))
            .collect();

        let extra_spaces = expansion_factor - 1;

        GalaxyMap(
            self.0
                .iter()
                .map(|p| Point {
                    x: p.x
                        + empty_cols.iter().filter(|empty_x| **empty_x < p.x).count()
                            * extra_spaces,
                    y: p.y
                        + empty_rows.iter().filter(|empty_y| **empty_y < p.y).count()
                            * extra_spaces,
                })
                .collect(),
        )
    }

    fn distance_sum_between_galaxies(&self) -> usize {
        let mut sum = 0;

        for i in 0..self.0.len() {
            for j in i + 1..self.0.len() {
                sum += self.0[i].distance_to(&self.0[j]);
            }
        }

        sum
    }
}

impl Point {
    fn distance_to(&self, other: &Point) -> usize {
        other.x.abs_diff(self.x) + other.y.abs_diff(self.y)
    }
}