use std::{
    cmp::max,
    collections::{BTreeSet, VecDeque},
    str::FromStr,
};

use aoc_2022::*;

use color_eyre::Report;
use itertools::Itertools;
// use tracing::debug;

type Input = Vec<Line>;
type Output = usize;

type Coord = Vector2D<i32>;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
struct Line {
    from: Coord,
    to: Coord,
}

impl Line {
    fn direction(&self) -> Coord {
        self.to - self.from
    }

    fn extend_to_x(&self, x: i32) -> Option<Coord> {
        let d = self.direction();

        match d.x() {
            0 => None,
            _ => {
                let k = (self.from.x() - x) / d.x();
                let y = self.from.y() + k * d.y();

                Some(Coord::new(x, y))
            }
        }
    }

    fn contains(&self, coord: &Coord) -> bool {
        let mut xs = [self.from.x(), self.to.x()];
        xs.sort();

        let mut ys = [self.from.y(), self.to.y()];
        ys.sort();

        (xs[0]..=xs[1]).contains(&coord.x()) && (ys[0]..=ys[1]).contains(&coord.y())
    }
}

#[derive(Debug)]
struct Multiline {
    lines: BTreeSet<Line>,
}

impl FromStr for Multiline {
    type Err = Report;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        fn parse_coord(s: &str) -> Result<Coord, Report> {
            let (x, y) = s.split(',').collect_tuple().unwrap();

            Ok(Coord::new(x.parse()?, y.parse()?))
        }

        let segs: Vec<Coord> = s
            .split(" -> ")
            .map(|c| parse_coord(c).unwrap())
            .collect_vec();
        let mut lines: BTreeSet<Line> = BTreeSet::new();

        for i in 0..segs.len() - 1 {
            lines.insert(Line {
                from: segs[i],
                to: segs[i + 1],
            });
        }

        Ok(Multiline { lines })
    }
}

fn free_falling(grain: &Coord, lines: &[Line], set: &BTreeSet<Coord>) -> bool {
    set.iter().all(|g| g.x() != grain.x() || g.y() < grain.y())
        && lines
            .iter()
            .filter_map(|l| l.extend_to_x(*grain.x()))
            .all(|g| g.y() < grain.y())
}

fn can_move(grain: &Coord, lines: &[Line], set: &BTreeSet<Coord>) -> Option<Coord> {
    for direction in [Coord::new(0, 1), Coord::new(-1, 1), Coord::new(1, 1)] {
        let new_position = *grain + direction;

        if !set.contains(&new_position) && lines.iter().all(|l| !l.contains(&new_position)) {
            return Some(new_position);
        }
    }

    None
}

fn simulate_falling<F>(lines: &Input, terminate: F) -> usize
where  
    F: Fn(&Coord, &Input, &BTreeSet<Coord>, bool) -> bool
{
    let mut grains: VecDeque<Coord> = VecDeque::new();
    let mut set_grains: BTreeSet<Coord> = BTreeSet::new();

    loop {
        for i in (0..grains.len()).rev() {
            let grain = grains[i];

            let mut is_set = false;
            if let Some(new_grain) = can_move(&grain, lines, &set_grains) {
                grains[i] = new_grain;
            } else {
                let grain = grains.pop_back().unwrap();
                set_grains.insert(grain);
                is_set = true;
            }

            if terminate(&grain, lines, &set_grains, is_set) {
                return set_grains.len();
            }
        }

        grains.push_front(Coord::new(500, 0));
    }
}

struct Day14;
impl Solution<Input, Output> for Day14 {
    fn parse_input<P: AsRef<Path>>(pathname: P) -> Input {
        file_to_structs::<P, Multiline>(pathname)
            .into_iter()
            .flat_map(|m| m.lines)
            .collect_vec()
    }

    fn part_1(input: &Input) -> Output {
        fn terminate(grain: &Coord, lines: &Input, set_grains: &BTreeSet<Coord>, _is_set: bool) -> bool {
            free_falling(grain, lines, set_grains)
        }

        simulate_falling(input, terminate)
    }

    fn part_2(input: &Input) -> Output {
        fn terminate(grain: &Coord, _lines: &Input, _set_grains: &BTreeSet<Coord>, is_set: bool) -> bool {
            is_set && grain == &Coord::new(500, 0)
        }

        let y = input
            .iter()
            .map(|l| max(l.from.y(), l.to.y()))
            .max()
            .unwrap();

        let mut lines = input.clone();
        lines.push(Line {
            from: Coord::new(500 - (*y + 2), *y + 2),
            to: Coord::new(500 + (*y + 2), *y + 2),
        });

        simulate_falling(&lines, terminate)
    }
}

fn main() -> Result<()> {
    // Day14::run("sample")
    Day14::main()
}

test_sample!(day_14, Day14, 24, 93);