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

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_14.txt")?;
    let rock_field = RockField::parser(&input).unwrap().1;
    {
        let mut north_rock_field = rock_field.clone();
        north_rock_field.tilt_north();
        dbg!(&north_rock_field.north_load());
    }

    {
        let mut spin_rock_field = rock_field.clone();
        let mut last_east = BTreeMap::new();
        let mut i = 0;
        let target = 1000000000;
        while i < target {
            spin_rock_field.tilt_north();
            spin_rock_field.tilt_west();
            spin_rock_field.tilt_south();
            spin_rock_field.tilt_east();
            if let Some(last_i) = last_east.get(&spin_rock_field) {
                let interval = i - last_i;
                // relying on integer division to round down here, want to add
                // interval as many times as possible without going over target.
                i += ((target - i) / interval) * interval;
            } else {
                last_east.insert(spin_rock_field.clone(), i);
            }
            i += 1;
        }
        dbg!(&spin_rock_field.north_load());
    }

    Ok(())
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct RockField(Vec<Vec<Rock>>);

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum Rock {
    Rounded,
    Cubed,
    None,
}

impl RockField {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(separated_list1(line_ending, many1(Rock::parser)), RockField)(input)
    }

    fn tilt_north(&mut self) {
        for y in 0..self.0.len() {
            for x in 0..self.0[y].len() {
                if self.0[y][x] == Rock::Rounded {
                    let new_y = (0..y)
                        .rev()
                        .take_while(|new_y| self.0[*new_y][x] == Rock::None)
                        .last();
                    if let Some(new_y) = new_y {
                        self.0[new_y][x] = Rock::Rounded;
                        self.0[y][x] = Rock::None;
                    }
                }
            }
        }
    }

    fn tilt_west(&mut self) {
        for x in 0..self.0[0].len() {
            for y in 0..self.0.len() {
                if self.0[y][x] == Rock::Rounded {
                    let new_x = (0..x)
                        .rev()
                        .take_while(|new_x| self.0[y][*new_x] == Rock::None)
                        .last();
                    if let Some(new_x) = new_x {
                        self.0[y][new_x] = Rock::Rounded;
                        self.0[y][x] = Rock::None;
                    }
                }
            }
        }
    }

    fn tilt_south(&mut self) {
        for y in (0..self.0.len()).rev() {
            for x in 0..self.0[y].len() {
                if self.0[y][x] == Rock::Rounded {
                    let new_y = (y + 1..self.0.len())
                        .take_while(|new_y| self.0[*new_y][x] == Rock::None)
                        .last();
                    if let Some(new_y) = new_y {
                        self.0[new_y][x] = Rock::Rounded;
                        self.0[y][x] = Rock::None;
                    }
                }
            }
        }
    }

    fn tilt_east(&mut self) {
        for x in (0..self.0[0].len()).rev() {
            for y in 0..self.0.len() {
                if self.0[y][x] == Rock::Rounded {
                    let new_x = (x + 1..self.0[0].len())
                        .take_while(|new_x| self.0[y][*new_x] == Rock::None)
                        .last();
                    if let Some(new_x) = new_x {
                        self.0[y][new_x] = Rock::Rounded;
                        self.0[y][x] = Rock::None;
                    }
                }
            }
        }
    }

    fn north_load(&self) -> usize {
        self.0
            .iter()
            .enumerate()
            .map(|(y, row)| {
                row.iter().filter(|r| **r == Rock::Rounded).count() * (self.0.len() - y)
            })
            .sum()
    }
}

impl Rock {
    fn parser(input: &str) -> IResult<&str, Self> {
        alt((
            value(Rock::Rounded, char('O')),
            value(Rock::Cubed, char('#')),
            value(Rock::None, char('.')),
        ))(input)
    }
}