use std::{collections::HashSet, ops::Add};

use aoc_runner_derive::{aoc, aoc_generator};
use itertools::Itertools;

use crate::util::parse_num;

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
struct Coord {
    x: i32,
    y: i32,
}

impl Coord {
    fn parse(s: &str) -> Option<(Self, &str)> {
        let (x, s) = parse_num(s)?;
        let s = s.strip_prefix(',')?;
        let (y, s) = parse_num(s)?;
        Some((Self { x, y }, s))
    }
}

impl Add<(i32, i32)> for Coord {
    type Output = Self;

    fn add(self, rhs: (i32, i32)) -> Self::Output {
        Self {
            x: self.x + rhs.0,
            y: self.y + rhs.1,
        }
    }
}

struct Path(Vec<Coord>);

impl Path {
    fn parse(s: &str) -> Option<(Self, &str)> {
        let (first, mut s) = Coord::parse(s)?;
        let mut vec = Vec::from([first]);

        loop {
            let Some(s2) = s.strip_prefix(" -> ") else {
                return Some((Self(vec), s))
            };
            let (coord, s2) = Coord::parse(s2)?;

            if coord.x != vec.last().unwrap().x && coord.y != vec.last().unwrap().y {
                return None;
            }

            vec.push(coord);
            s = s2;
        }
    }

    fn intersects(&self, coord: Coord) -> bool {
        for (a, b) in self.0.iter().tuple_windows() {
            if a.y == b.y {
                if coord.y == a.y
                    && ((a.x <= coord.x && coord.x <= b.x) || (b.x <= coord.x && coord.x <= a.x))
                {
                    return true;
                }
            } else if a.x == b.x {
                if coord.x == a.x
                    && ((a.y <= coord.y && coord.y <= b.y) || (b.y <= coord.y && coord.y <= a.y))
                {
                    return true;
                }
            } else {
                unreachable!()
            }
        }

        false
    }
}

#[aoc_generator(day14)]
fn generator(input: &str) -> Vec<Path> {
    input.lines().map(|l| Path::parse(l).unwrap().0).collect()
}

fn is_blocked(rocks: &[Path], sand: &HashSet<Coord>, coord: Coord) -> bool {
    sand.contains(&coord) || rocks.iter().any(|path| path.intersects(coord))
}

#[aoc(day14, part1)]
fn part1(input: &[Path]) -> usize {
    let abyss_depth = input
        .iter()
        .flat_map(|path| path.0.iter().map(|coord| coord.y))
        .max()
        .unwrap()
        + 1; // for good measure

    let mut sand = HashSet::new();

    loop {
        let mut particle = Coord { x: 500, y: 0 };

        loop {
            if !is_blocked(input, &sand, particle + (0, 1)) {
                particle = particle + (0, 1);
            } else if !is_blocked(input, &sand, particle + (-1, 1)) {
                particle = particle + (-1, 1);
            } else if !is_blocked(input, &sand, particle + (1, 1)) {
                particle = particle + (1, 1);
            } else {
                sand.insert(particle);
                break;
            }

            if particle.y >= abyss_depth {
                return sand.len();
            }
        }
    }
}

#[aoc(day14, part2)]
fn part2(input: &[Path]) -> usize {
    let floor_depth = input
        .iter()
        .flat_map(|path| path.0.iter().map(|coord| coord.y))
        .max()
        .unwrap()
        + 2;

    let mut sand = HashSet::new();

    loop {
        let mut particle = Coord { x: 500, y: 0 };

        loop {
            if particle.y + 1 == floor_depth {
                sand.insert(particle);
                break;
            }

            if !is_blocked(input, &sand, particle + (0, 1)) {
                particle = particle + (0, 1);
            } else if !is_blocked(input, &sand, particle + (-1, 1)) {
                particle = particle + (-1, 1);
            } else if !is_blocked(input, &sand, particle + (1, 1)) {
                particle = particle + (1, 1);
            } else {
                sand.insert(particle);

                if particle == (Coord { x: 500, y: 0 }) {
                    return sand.len();
                }

                break;
            }
        }
    }
}